From bc1f7b807984cdb5387dbca1dc543732a630a7de Mon Sep 17 00:00:00 2001 From: He Weiwei Date: Sun, 26 Apr 2020 15:56:48 +0000 Subject: [PATCH] feat(scan): support to scan artifact automatic after it pushed Closes #11692 Signed-off-by: He Weiwei --- src/controller/event/handler/init.go | 4 +- .../event/handler/internal/artifact.go | 23 +++- src/controller/event/handler/internal/util.go | 40 +++++++ .../event/handler/internal/util_test.go | 112 ++++++++++++++++++ src/lib/orm/orm.go | 6 + 5 files changed, 178 insertions(+), 7 deletions(-) create mode 100644 src/controller/event/handler/internal/util.go create mode 100644 src/controller/event/handler/internal/util_test.go diff --git a/src/controller/event/handler/init.go b/src/controller/event/handler/init.go index 2acb34932..7018d1fe7 100644 --- a/src/controller/event/handler/init.go +++ b/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}) } diff --git a/src/controller/event/handler/internal/artifact.go b/src/controller/event/handler/internal/artifact.go index b144f8b0a..b696e2b2d 100644 --- a/src/controller/event/handler/internal/artifact.go +++ b/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 { - ctx := orm.NewContext(context.Background(), beegorm.NewOrm()) +func (a *Handler) onPull(ctx context.Context, event *event.ArtifactEvent) error { 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 +} diff --git a/src/controller/event/handler/internal/util.go b/src/controller/event/handler/internal/util.go new file mode 100644 index 000000000..5f7a1cde3 --- /dev/null +++ b/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) +} diff --git a/src/controller/event/handler/internal/util_test.go b/src/controller/event/handler/internal/util_test.go new file mode 100644 index 000000000..dcb24e53e --- /dev/null +++ b/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{}) +} diff --git a/src/lib/orm/orm.go b/src/lib/orm/orm.go index 1c1277258..91a0eb19f 100644 --- a/src/lib/orm/orm.go +++ b/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 {