Merge pull request #11757 from heww/fix-issue-11692

feat(scan): support to scan artifact automatic after it pushed

src/controller/event/handler/init.go

@@ -9,6 +9,7 @@ import (
 	"github.com/goharbor/harbor/src/controller/event/handler/webhook/chart"
 	"github.com/goharbor/harbor/src/controller/event/handler/webhook/quota"
 	"github.com/goharbor/harbor/src/controller/event/handler/webhook/scan"
+	"github.com/goharbor/harbor/src/lib/orm"
 	"github.com/goharbor/harbor/src/pkg/notifier"
 )
 
@@ -44,5 +45,6 @@ func init() {
 	notifier.Subscribe(event.TopicDeleteTag, &auditlog.Handler{})
 
 	// internal
-	notifier.Subscribe(event.TopicPullArtifact, &internal.Handler{})
+	notifier.Subscribe(event.TopicPullArtifact, &internal.Handler{Context: orm.Context})
+	notifier.Subscribe(event.TopicPushArtifact, &internal.Handler{Context: orm.Context})
 }

src/controller/event/handler/internal/artifact.go

@@ -16,26 +16,28 @@ package internal
 
 import (
 	"context"
-	beegorm "github.com/astaxie/beego/orm"
+	"time"
+
 	"github.com/goharbor/harbor/src/controller/artifact"
 	"github.com/goharbor/harbor/src/controller/event"
 	"github.com/goharbor/harbor/src/controller/repository"
 	"github.com/goharbor/harbor/src/controller/tag"
 	"github.com/goharbor/harbor/src/lib/log"
-	"github.com/goharbor/harbor/src/lib/orm"
 	"github.com/goharbor/harbor/src/lib/q"
-	"time"
 )
 
 // Handler preprocess artifact event data
 type Handler struct {
+	Context func() context.Context
 }
 
 // Handle ...
 func (a *Handler) Handle(value interface{}) error {
 	switch v := value.(type) {
 	case *event.PullArtifactEvent:
-		return a.handle(v.ArtifactEvent)
+		return a.onPull(a.Context(), v.ArtifactEvent)
+	case *event.PushArtifactEvent:
+		return a.onPush(a.Context(), v.ArtifactEvent)
 	default:
 		log.Errorf("Can not handler this event type! %#v", v)
 	}
@@ -47,8 +49,7 @@ func (a *Handler) IsStateful() bool {
 	return false
 }
 
-func (a *Handler) handle(event *event.ArtifactEvent) error {
+func (a *Handler) onPull(ctx context.Context, event *event.ArtifactEvent) error {
-	ctx := orm.NewContext(context.Background(), beegorm.NewOrm())
 	go func() { a.updatePullTime(ctx, event) }()
 	go func() { a.addPullCount(ctx, event) }()
 	return nil
@@ -81,3 +82,13 @@ func (a *Handler) addPullCount(ctx context.Context, event *event.ArtifactEvent)
 	}
 	return
 }
+
+func (a *Handler) onPush(ctx context.Context, event *event.ArtifactEvent) error {
+	go func() {
+		if err := autoScan(ctx, &artifact.Artifact{Artifact: *event.Artifact}); err != nil {
+			log.Errorf("scan artifact %s@%s failed, error: %v", event.Artifact.RepositoryName, event.Artifact.Digest, err)
+		}
+	}()
+
+	return nil
+}

src/controller/event/handler/internal/util.go

@@ -0,0 +1,40 @@
+// Copyright Project Harbor Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package internal
+
+import (
+	"context"
+
+	"github.com/goharbor/harbor/src/controller/artifact"
+	"github.com/goharbor/harbor/src/controller/project"
+	"github.com/goharbor/harbor/src/controller/scan"
+	"github.com/goharbor/harbor/src/lib/orm"
+)
+
+// autoScan scan artifact when the project of the artifact enable auto scan
+func autoScan(ctx context.Context, a *artifact.Artifact) error {
+	proj, err := project.Ctl.Get(ctx, a.ProjectID)
+	if err != nil {
+		return err
+	}
+	if !proj.AutoScan() {
+		return nil
+	}
+
+	// transaction here to work with the image index
+	return orm.WithTransaction(func(ctx context.Context) error {
+		return scan.DefaultController.Scan(ctx, a)
+	})(ctx)
+}

src/controller/event/handler/internal/util_test.go

@@ -0,0 +1,112 @@
+// Copyright Project Harbor Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package internal
+
+import (
+	"testing"
+
+	"github.com/goharbor/harbor/src/common/models"
+	"github.com/goharbor/harbor/src/controller/artifact"
+	"github.com/goharbor/harbor/src/controller/project"
+	"github.com/goharbor/harbor/src/controller/scan"
+	"github.com/goharbor/harbor/src/lib/errors"
+	"github.com/goharbor/harbor/src/lib/orm"
+	projecttesting "github.com/goharbor/harbor/src/testing/controller/project"
+	scantesting "github.com/goharbor/harbor/src/testing/controller/scan"
+	ormtesting "github.com/goharbor/harbor/src/testing/lib/orm"
+	"github.com/goharbor/harbor/src/testing/mock"
+	"github.com/stretchr/testify/suite"
+)
+
+type AutoScanTestSuite struct {
+	suite.Suite
+
+	originalProjectController project.Controller
+	projectController         *projecttesting.Controller
+
+	originalScanController scan.Controller
+	scanController         *scantesting.Controller
+}
+
+func (suite *AutoScanTestSuite) SetupTest() {
+	suite.originalProjectController = project.Ctl
+	suite.projectController = &projecttesting.Controller{}
+	project.Ctl = suite.projectController
+
+	suite.originalScanController = scan.DefaultController
+	suite.scanController = &scantesting.Controller{}
+	scan.DefaultController = suite.scanController
+}
+
+func (suite *AutoScanTestSuite) TearDownTest() {
+	project.Ctl = suite.originalProjectController
+	scan.DefaultController = suite.originalScanController
+}
+
+func (suite *AutoScanTestSuite) TestGetProjectFailed() {
+	mock.OnAnything(suite.projectController, "Get").Return(nil, errors.NotFoundError(nil))
+
+	ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
+	art := &artifact.Artifact{}
+
+	suite.Error(autoScan(ctx, art))
+}
+
+func (suite *AutoScanTestSuite) TestAutoScanDisabled() {
+	mock.OnAnything(suite.projectController, "Get").Return(&models.Project{
+		Metadata: map[string]string{
+			models.ProMetaAutoScan: "false",
+		},
+	}, nil)
+
+	ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
+	art := &artifact.Artifact{}
+
+	suite.Nil(autoScan(ctx, art))
+}
+
+func (suite *AutoScanTestSuite) TestAutoScan() {
+	mock.OnAnything(suite.projectController, "Get").Return(&models.Project{
+		Metadata: map[string]string{
+			models.ProMetaAutoScan: "true",
+		},
+	}, nil)
+
+	mock.OnAnything(suite.scanController, "Scan").Return(nil)
+
+	ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
+	art := &artifact.Artifact{}
+
+	suite.Nil(autoScan(ctx, art))
+}
+
+func (suite *AutoScanTestSuite) TestAutoScanFailed() {
+	mock.OnAnything(suite.projectController, "Get").Return(&models.Project{
+		Metadata: map[string]string{
+			models.ProMetaAutoScan: "true",
+		},
+	}, nil)
+
+	mock.OnAnything(suite.scanController, "Scan").Return(errors.ConflictError(nil))
+
+	ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
+	art := &artifact.Artifact{}
+
+	suite.Error(autoScan(ctx, art))
+}
+
+func TestAutoScanTestSuite(t *testing.T) {
+	suite.Run(t, &AutoScanTestSuite{})
+}

src/lib/orm/orm.go

@@ -17,6 +17,7 @@ package orm
 import (
 	"context"
 	"errors"
+
 	"github.com/astaxie/beego/orm"
 	"github.com/goharbor/harbor/src/lib/log"
 )
@@ -40,6 +41,11 @@ func NewContext(ctx context.Context, o orm.Ormer) context.Context {
 	return context.WithValue(ctx, ormKey{}, o)
 }
 
+// Context returns a context with an orm
+func Context() context.Context {
+	return NewContext(context.Background(), orm.NewOrm())
+}
+
 // WithTransaction a decorator which make f run in transaction
 func WithTransaction(f func(ctx context.Context) error) func(ctx context.Context) error {
 	return func(ctx context.Context) error {