handle blob status chanage in put blob middlware (#12315)

* handle blob status chanage in put blob middlware

After blob is uploaded success, the middleware will update the blob status accordingly.

Signed-off-by: wang yan <wangyan@vmware.com>

src/server/middleware/blob/put_blob_upload.go

@@ -15,17 +15,57 @@
 package blob
 
 import (
-	"net/http"
+	"fmt"
-	"strconv"
+	"github.com/goharbor/harbor/src/lib/errors"
-
 	"github.com/goharbor/harbor/src/lib/log"
+	blob_models "github.com/goharbor/harbor/src/pkg/blob/models"
 	"github.com/goharbor/harbor/src/pkg/distribution"
 	"github.com/goharbor/harbor/src/server/middleware"
+	"github.com/goharbor/harbor/src/server/middleware/requestid"
+	"net/http"
+	"strconv"
 )
 
-// PutBlobUploadMiddleware middleware to create Blob and ProjectBlob after PUT /v2/<name>/blobs/uploads/<session_id> success
+// PutBlobUploadMiddleware middleware is to update the blob status according to the different situation before the request passed into proxy(distribution).
+// And it creates Blob and ProjectBlob after PUT /v2/<name>/blobs/uploads/<session_id>?digest=<digest> success - http.StatusCreated
+// Why to use the middleware to handle blob status?
+// 1, As Put blob will always happen after head blob gets a 404, but the 404 could be caused by blob status is deleting, which is marked by GC.
+// 2, It has to deal with the concurrence blob push.
 func PutBlobUploadMiddleware() func(http.Handler) http.Handler {
-	return middleware.AfterResponse(func(w http.ResponseWriter, r *http.Request, statusCode int) error {
+
+	before := middleware.BeforeRequest(func(r *http.Request) error {
+		ctx := r.Context()
+		logger := log.G(ctx)
+
+		v := r.URL.Query()
+		digest := v.Get("digest")
+
+		// digest empty is handled by the blob controller GET method
+		bb, err := blobController.Get(r.Context(), digest)
+		if err != nil {
+			if errors.IsNotFoundErr(err) {
+				return nil
+			}
+			return err
+		}
+
+		switch bb.Status {
+		case blob_models.StatusNone, blob_models.StatusDelete, blob_models.StatusDeleteFailed:
+			err := blobController.Touch(r.Context(), bb)
+			if err != nil {
+				logger.Errorf("failed to update blob: %s status to StatusNone, error:%v", bb.Digest, err)
+				return errors.Wrapf(err, fmt.Sprintf("the request id is: %s", r.Header.Get(requestid.HeaderXRequestID)))
+			}
+		case blob_models.StatusDeleting:
+			logger.Warningf(fmt.Sprintf("the asking blob is in GC, mark it as non existing, request id: %s", r.Header.Get(requestid.HeaderXRequestID)))
+			return errors.New(nil).WithMessage(fmt.Sprintf("the asking blob is in GC, mark it as non existing, request id: %s", r.Header.Get(requestid.HeaderXRequestID))).WithCode(errors.NotFoundCode)
+		default:
+			return nil
+		}
+		return nil
+	})
+
+	after := middleware.AfterResponse(func(w http.ResponseWriter, r *http.Request, statusCode int) error {
 		if statusCode != http.StatusCreated {
 			return nil
 		}
@@ -63,4 +103,6 @@ func PutBlobUploadMiddleware() func(http.Handler) http.Handler {
 
 		return nil
 	})
+
+	return middleware.Chain(before, after)
 }

src/server/middleware/blob/put_blob_upload_test.go

@@ -16,6 +16,8 @@ package blob
 
 import (
 	"fmt"
+	pkg_blob "github.com/goharbor/harbor/src/pkg/blob"
+	blob_models "github.com/goharbor/harbor/src/pkg/blob/models"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -37,12 +39,11 @@ func (suite *PutBlobUploadMiddlewareTestSuite) SetupSuite() {
 
 func (suite *PutBlobUploadMiddlewareTestSuite) TestDataInBody() {
 	suite.WithProject(func(projectID int64, projectName string) {
-		req := suite.NewRequest(http.MethodPut, fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s", projectName, uuid.New().String()), nil)
+		digest := suite.DigestString()
+		req := suite.NewRequest(http.MethodPut, fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s?digest=%s", projectName, uuid.New().String(), digest), nil)
 		req.Header.Set("Content-Length", "512")
 		res := httptest.NewRecorder()
 
-		digest := suite.DigestString()
-
 		next := suite.NextHandler(http.StatusCreated, map[string]string{"Docker-Content-Digest": digest})
 		PutBlobUploadMiddleware()(next).ServeHTTP(res, req)
 
@@ -60,7 +61,8 @@ func (suite *PutBlobUploadMiddlewareTestSuite) TestDataInBody() {
 func (suite *PutBlobUploadMiddlewareTestSuite) TestWithoutBody() {
 	suite.WithProject(func(projectID int64, projectName string) {
 		sessionID := uuid.New().String()
-		path := fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s", projectName, sessionID)
+		digest := suite.DigestString()
+		path := fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s?digest=%s", projectName, sessionID, digest)
 
 		{
 			req := httptest.NewRequest(http.MethodPatch, path, nil)
@@ -74,8 +76,6 @@ func (suite *PutBlobUploadMiddlewareTestSuite) TestWithoutBody() {
 		req := suite.NewRequest(http.MethodPut, path, nil)
 		res := httptest.NewRecorder()
 
-		digest := suite.DigestString()
-
 		next := suite.NextHandler(http.StatusCreated, map[string]string{"Docker-Content-Digest": digest})
 		PutBlobUploadMiddleware()(next).ServeHTTP(res, req)
 		suite.Equal(http.StatusCreated, res.Code)
@@ -92,6 +92,71 @@ func (suite *PutBlobUploadMiddlewareTestSuite) TestWithoutBody() {
 	})
 }
 
+func (suite *PutBlobUploadMiddlewareTestSuite) TestBlobInDeleting() {
+	suite.WithProject(func(projectID int64, projectName string) {
+		digest := suite.DigestString()
+
+		id, err := blob.Ctl.Ensure(suite.Context(), digest, "application/octet-stream", 512)
+		suite.Nil(err)
+
+		// status-none -> status-delete -> status-deleting
+		_, err = pkg_blob.Mgr.UpdateBlobStatus(suite.Context(), &blob_models.Blob{ID: id, Status: blob_models.StatusDelete})
+		suite.Nil(err)
+		_, err = pkg_blob.Mgr.UpdateBlobStatus(suite.Context(), &blob_models.Blob{ID: id, Status: blob_models.StatusDeleting, Version: 1})
+		suite.Nil(err)
+
+		req := suite.NewRequest(http.MethodPut, fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s?digest=%s", projectName, uuid.New().String(), digest), nil)
+		req.Header.Set("Content-Length", "512")
+		res := httptest.NewRecorder()
+
+		next := suite.NextHandler(http.StatusCreated, map[string]string{"Docker-Content-Digest": digest})
+		PutBlobUploadMiddleware()(next).ServeHTTP(res, req)
+		suite.Equal(http.StatusNotFound, res.Code)
+	})
+}
+
+func (suite *PutBlobUploadMiddlewareTestSuite) TestBlobInDelete() {
+	suite.WithProject(func(projectID int64, projectName string) {
+		digest := suite.DigestString()
+
+		id, err := blob.Ctl.Ensure(suite.Context(), digest, "application/octet-stream", 512)
+		suite.Nil(err)
+
+		_, err = pkg_blob.Mgr.UpdateBlobStatus(suite.Context(), &blob_models.Blob{ID: id, Status: blob_models.StatusDelete})
+		suite.Nil(err)
+
+		req := suite.NewRequest(http.MethodPut, fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s?digest=%s", projectName, uuid.New().String(), digest), nil)
+		req.Header.Set("Content-Length", "512")
+		res := httptest.NewRecorder()
+
+		next := suite.NextHandler(http.StatusCreated, map[string]string{"Docker-Content-Digest": digest})
+		PutBlobUploadMiddleware()(next).ServeHTTP(res, req)
+		suite.Equal(http.StatusCreated, res.Code)
+
+		exist, err := blob.Ctl.Exist(suite.Context(), digest, blob.IsAssociatedWithProject(projectID))
+		suite.Nil(err)
+		suite.True(exist)
+
+		blob, err := blob.Ctl.Get(suite.Context(), digest)
+		if suite.Nil(err) {
+			suite.Equal(digest, blob.Digest)
+			suite.Equal(int64(512), blob.Size)
+			suite.Equal(blob_models.StatusNone, blob.Status)
+		}
+	})
+}
+
+func (suite *PutBlobUploadMiddlewareTestSuite) TestRequestWithoutDigest() {
+	suite.WithProject(func(projectID int64, projectName string) {
+		req := suite.NewRequest(http.MethodPut, fmt.Sprintf("/v2/%s/photon/blobs/uploads/%s", projectName, uuid.New().String()), nil)
+		req.Header.Set("Content-Length", "512")
+		next := suite.NextHandler(http.StatusCreated, map[string]string{})
+		res := httptest.NewRecorder()
+		PutBlobUploadMiddleware()(next).ServeHTTP(res, req)
+		suite.Equal(http.StatusBadRequest, res.Code)
+	})
+}
+
 func TestPutBlobUploadMiddlewareTestSuite(t *testing.T) {
 	suite.Run(t, &PutBlobUploadMiddlewareTestSuite{})
 }