Add P2P trigger event and handler

Signed-off-by: peimingming <peimingming@corp.netease.com>

src/controller/artifact/controller.go

@@ -45,6 +45,7 @@ import (
 	"github.com/goharbor/harbor/src/pkg/immutabletag/match/rule"
 	"github.com/goharbor/harbor/src/pkg/label"
 	"github.com/goharbor/harbor/src/pkg/notification"
+	"github.com/goharbor/harbor/src/pkg/notifier/event"
 	"github.com/goharbor/harbor/src/pkg/registry"
 	"github.com/goharbor/harbor/src/pkg/repository"
 	"github.com/goharbor/harbor/src/pkg/signature"
@@ -483,8 +484,27 @@ func (c *controller) GetAddition(ctx context.Context, artifactID int64, addition
 	return processor.Get(artifact.MediaType).AbstractAddition(ctx, artifact, addition)
 }
 
-func (c *controller) AddLabel(ctx context.Context, artifactID int64, labelID int64) error {
-	return c.labelMgr.AddTo(ctx, labelID, artifactID)
+func (c *controller) AddLabel(ctx context.Context, artifactID int64, labelID int64) (err error) {
+	defer func() {
+		if err == nil {
+			// trigger label artifact event
+			e := &event.Event{}
+			metaData := &metadata.ArtifactLabeledMetadata{
+				ArtifactID: artifactID,
+				LabelID:    labelID,
+				Ctx:        ctx,
+			}
+			if err := e.Build(metaData); err == nil {
+				if err := e.Publish(); err != nil {
+					log.Error(errors.Wrap(err, "mark label to resource handler: event publish"))
+				}
+			} else {
+				log.Error(errors.Wrap(err, "mark label to resource handler: event build"))
+			}
+		}
+	}()
+	err = c.labelMgr.AddTo(ctx, labelID, artifactID)
+	return
 }
 
 func (c *controller) RemoveLabel(ctx context.Context, artifactID int64, labelID int64) error {

src/controller/artifact/controller_test.go

@@ -477,7 +477,7 @@ func (c *controllerTestSuite) TestGetAddition() {
 
 func (c *controllerTestSuite) TestAddTo() {
 	c.labelMgr.On("AddTo").Return(nil)
-	err := c.ctl.AddLabel(nil, 1, 1)
+	err := c.ctl.AddLabel(context.Background(), 1, 1)
 	c.Require().Nil(err)
 }
 

src/controller/event/handler/init.go

@@ -4,6 +4,7 @@ import (
 	"github.com/goharbor/harbor/src/controller/event"
 	"github.com/goharbor/harbor/src/controller/event/handler/auditlog"
 	"github.com/goharbor/harbor/src/controller/event/handler/internal"
+	"github.com/goharbor/harbor/src/controller/event/handler/p2p"
 	"github.com/goharbor/harbor/src/controller/event/handler/replication"
 	"github.com/goharbor/harbor/src/controller/event/handler/webhook/artifact"
 	"github.com/goharbor/harbor/src/controller/event/handler/webhook/chart"
@@ -34,6 +35,11 @@ func init() {
 	notifier.Subscribe(event.TopicCreateTag, &replication.Handler{})
 	notifier.Subscribe(event.TopicDeleteTag, &replication.Handler{})
 
+	// p2p preheat
+	notifier.Subscribe(event.TopicPushArtifact, &p2p.Handler{Context: orm.Context})
+	notifier.Subscribe(event.TopicScanningCompleted, &p2p.Handler{Context: orm.Context})
+	notifier.Subscribe(event.TopicArtifactLabeled, &p2p.Handler{Context: orm.Context})
+
 	// audit logs
 	notifier.Subscribe(event.TopicPushArtifact, &auditlog.Handler{})
 	notifier.Subscribe(event.TopicPullArtifact, &auditlog.Handler{})

src/controller/event/handler/p2p/preheat.go

@@ -0,0 +1,97 @@
+// 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 p2p
+
+import (
+	"context"
+	"github.com/goharbor/harbor/src/controller/artifact"
+	"github.com/goharbor/harbor/src/controller/artifact/processor/image"
+	"github.com/goharbor/harbor/src/controller/event"
+	"github.com/goharbor/harbor/src/controller/p2p/preheat"
+	"github.com/goharbor/harbor/src/lib/errors"
+	"github.com/goharbor/harbor/src/lib/log"
+)
+
+// Handler ...
+type Handler struct {
+	Context func() context.Context
+}
+
+// Handle ...
+func (p *Handler) Handle(value interface{}) error {
+	switch value.(type) {
+	case *event.PushArtifactEvent:
+		pushArtEvent, _ := value.(*event.PushArtifactEvent)
+		return p.handlePushArtifact(pushArtEvent)
+	case *event.ScanImageEvent:
+		scanImageEvent, _ := value.(*event.ScanImageEvent)
+		return p.handleImageScanned(scanImageEvent)
+	case *event.ArtifactLabeledEvent:
+		artifactLabeledEvent, _ := value.(*event.ArtifactLabeledEvent)
+		return p.handleArtifactLabeled(artifactLabeledEvent)
+	default:
+		return errors.New("unsupported type")
+	}
+}
+
+// IsStateful ...
+func (p *Handler) IsStateful() bool {
+	return false
+}
+
+func (p *Handler) handlePushArtifact(event *event.PushArtifactEvent) error {
+	if event.Artifact.Type != image.ArtifactTypeImage {
+		return nil
+	}
+	log.Debugf("preheat artifact event %s:%s", event.Artifact.RepositoryName, event.Artifact.Digest)
+
+	art, err := artifact.Ctl.Get(p.Context(), event.Artifact.ID, &artifact.Option{
+		WithTag:   true,
+		WithLabel: true,
+	})
+	if err != nil {
+		return err
+	}
+	_, err = preheat.Enf.PreheatArtifact(p.Context(), art)
+	return err
+}
+
+func (p *Handler) handleImageScanned(event *event.ScanImageEvent) error {
+	log.Debugf("preheat image scanned %s:%s", event.Artifact.Repository, event.Artifact.Tag)
+	art, err := artifact.Ctl.GetByReference(p.Context(), event.Artifact.Repository, event.Artifact.Digest,
+		&artifact.Option{
+			WithTag:   true,
+			WithLabel: true,
+		})
+	if err != nil {
+		return err
+	}
+	_, err = preheat.Enf.PreheatArtifact(p.Context(), art)
+	return err
+}
+
+func (p *Handler) handleArtifactLabeled(event *event.ArtifactLabeledEvent) error {
+	art, err := artifact.Ctl.Get(p.Context(), event.ArtifactID, &artifact.Option{
+		WithTag:   true,
+		WithLabel: true,
+	})
+	if art.Type != image.ArtifactTypeImage {
+		return nil
+	}
+	log.Debugf("preheat artifact labeled %s:%s", art.Artifact.RepositoryName, art.Artifact.Digest)
+
+	_, err = preheat.Enf.PreheatArtifact(p.Context(), art)
+	return err
+}

src/controller/event/handler/p2p/preheat_test.go

@@ -0,0 +1,144 @@
+package p2p
+
+import (
+	"context"
+	"github.com/goharbor/harbor/src/controller/artifact"
+	"github.com/goharbor/harbor/src/controller/event"
+	"github.com/goharbor/harbor/src/controller/p2p/preheat"
+	pkg_artifact "github.com/goharbor/harbor/src/pkg/artifact"
+	"github.com/goharbor/harbor/src/pkg/scan/rest/v1"
+	test_artifact "github.com/goharbor/harbor/src/testing/controller/artifact"
+	test_preheat "github.com/goharbor/harbor/src/testing/controller/p2p/preheat"
+	"github.com/goharbor/harbor/src/testing/mock"
+	"github.com/stretchr/testify/suite"
+	"testing"
+	"time"
+)
+
+// PreheatTestSuite is a test suite of testing preheat handler
+type PreheatTestSuite struct {
+	suite.Suite
+	artifactCtl artifact.Controller
+	preheatEnf  preheat.Enforcer
+
+	handler *Handler
+}
+
+// TestPreheat is an entry method of running PreheatTestSuite
+func TestPreheat(t *testing.T) {
+	suite.Run(t, &PreheatTestSuite{})
+}
+
+// SetupSuite prepares env for running PreheatTestSuite
+func (suite *PreheatTestSuite) SetupSuite() {
+	fakeArtifactCtl := &test_artifact.Controller{}
+	fakeArtifactCtl.On("GetByReference",
+		context.TODO(),
+		mock.AnythingOfType("string"),
+		mock.AnythingOfType("string"),
+		mock.AnythingOfType("*artifact.Option"),
+	).Return(&artifact.Artifact{}, nil)
+	fakeArtifactCtl.On("Get",
+		context.TODO(),
+		mock.AnythingOfType("int64"),
+		mock.AnythingOfType("*artifact.Option"),
+	).Return(&artifact.Artifact{
+		Artifact: pkg_artifact.Artifact{
+			Type: "IMAGE",
+		},
+	}, nil)
+
+	fakeEnforcer := &test_preheat.FakeEnforcer{}
+	fakeEnforcer.On("PreheatArtifact",
+		context.TODO(),
+		mock.AnythingOfType("*artifact.Artifact"),
+	).Return(nil, nil)
+
+	suite.artifactCtl = artifact.Ctl
+	artifact.Ctl = fakeArtifactCtl
+	suite.preheatEnf = preheat.Enf
+	preheat.Enf = fakeEnforcer
+	suite.handler = &Handler{
+		Context: context.TODO,
+	}
+}
+
+// TearDownSuite cleans the testing env
+func (suite *PreheatTestSuite) TearDownSuite() {
+	artifact.Ctl = suite.artifactCtl
+	preheat.Enf = suite.preheatEnf
+}
+
+// TestIsStateful ...
+func (suite *PreheatTestSuite) TestIsStateful() {
+	b := suite.handler.IsStateful()
+	suite.False(b, "handler is stateful")
+}
+
+// TestHandle ...
+func (suite *PreheatTestSuite) TestHandle() {
+	type args struct {
+		data interface{}
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "PreheatHandler String Error",
+			args: args{
+				data: "",
+			},
+			wantErr: true,
+		},
+		{
+			name: "PreheatHandler 1",
+			args: args{
+				data: &event.PushArtifactEvent{
+					ArtifactEvent: &event.ArtifactEvent{
+						Artifact: &pkg_artifact.Artifact{
+							Type:         "IMAGE",
+							ID:           11,
+							RepositoryID: 23,
+						},
+						Tags: []string{"v1.1", "v1.2"},
+					},
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "PreheatHandler 2",
+			args: args{
+				data: &event.ScanImageEvent{
+					OccurAt: time.Now(),
+					Artifact: &v1.Artifact{
+						Repository: "library",
+						Digest:     "sha256:1359608115b94599e5641638bac5aef1ddfaa79bb96057ebf41ebc8d33acf8a7",
+					},
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "PreheatHandler 3",
+			args: args{
+				data: &event.ArtifactLabeledEvent{
+					OccurAt:    time.Now(),
+					ArtifactID: 1,
+					LabelID:    2,
+				},
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		err := suite.handler.Handle(tt.args.data)
+		if tt.wantErr {
+			suite.Error(err, tt.name)
+		} else {
+			suite.NoError(err, tt.name)
+		}
+	}
+}

src/controller/event/metadata/label.go

@@ -0,0 +1,47 @@
+// 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 metadata
+
+import (
+	"context"
+	"github.com/goharbor/harbor/src/common/security"
+	event2 "github.com/goharbor/harbor/src/controller/event"
+	"github.com/goharbor/harbor/src/pkg/notifier/event"
+	"time"
+)
+
+// ArtifactLabeledMetadata is the metadata from which the artifact labeled event can be resolved
+type ArtifactLabeledMetadata struct {
+	Ctx        context.Context
+	ArtifactID int64
+	LabelID    int64
+	Operator   string
+}
+
+// Resolve to the event from the metadata
+func (al *ArtifactLabeledMetadata) Resolve(event *event.Event) error {
+	data := &event2.ArtifactLabeledEvent{
+		ArtifactID: al.ArtifactID,
+		LabelID:    al.LabelID,
+		OccurAt:    time.Now(),
+	}
+	ctx, exist := security.FromContext(al.Ctx)
+	if exist {
+		data.Operator = ctx.GetUsername()
+	}
+	event.Topic = event2.TopicArtifactLabeled
+	event.Data = data
+	return nil
+}

src/controller/event/topic.go

@@ -39,12 +39,13 @@ const (
 	TopicScanningFailed    = "SCANNING_FAILED"
 	TopicScanningCompleted = "SCANNING_COMPLETED"
 	// QuotaExceedTopic is topic for quota warning event, the usage reaches the warning bar of limitation, like 85%
-	TopicQuotaWarning  = "QUOTA_WARNING"
-	TopicQuotaExceed   = "QUOTA_EXCEED"
-	TopicUploadChart   = "UPLOAD_CHART"
-	TopicDownloadChart = "DOWNLOAD_CHART"
-	TopicDeleteChart   = "DELETE_CHART"
-	TopicReplication   = "REPLICATION"
+	TopicQuotaWarning    = "QUOTA_WARNING"
+	TopicQuotaExceed     = "QUOTA_EXCEED"
+	TopicUploadChart     = "UPLOAD_CHART"
+	TopicDownloadChart   = "DOWNLOAD_CHART"
+	TopicDeleteChart     = "DELETE_CHART"
+	TopicReplication     = "REPLICATION"
+	TopicArtifactLabeled = "ARTIFACT_LABELED"
 )
 
 // CreateProjectEvent is the creating project event
@@ -280,3 +281,11 @@ type ReplicationEvent struct {
 	OccurAt           time.Time
 	Status            string
 }
+
+// ArtifactLabeledEvent is event data of artifact labeled
+type ArtifactLabeledEvent struct {
+	ArtifactID int64
+	LabelID    int64
+	OccurAt    time.Time
+	Operator   string
+}

src/controller/p2p/preheat/enforcer.go

@@ -100,6 +100,11 @@ type extURLGetter func(c *selector.Candidate) (string, error)
 // The purpose of defining such a func template is decoupling code
 type accessCredMaker func(c *selector.Candidate) (string, error)
 
+var (
+	// Enf default enforcer
+	Enf = NewEnforcer()
+)
+
 // defaultEnforcer is default implementation of Enforcer.
 type defaultEnforcer struct {
 	// for policy management
@@ -121,7 +126,7 @@ type defaultEnforcer struct {
 	credMaker accessCredMaker
 }
 
-// NewEnforcer creat a new enforcer
+// NewEnforcer create a new enforcer
 func NewEnforcer() Enforcer {
 	return &defaultEnforcer{
 		policyMgr:    policy.New(),