fix: limit the file size of the cnai model processor (#21759)

Signed-off-by: chlins <chlins.zhang@gmail.com>

src/controller/artifact/processor/cnai/cnai_test.go

@@ -207,7 +207,7 @@ func (p *ProcessorTestSuite) TestAbstractAddition() {
 				p.Require().NoError(err)
 				r.On("PullManifest", mock.Anything, mock.Anything).Return(manifest, "", nil)
 			},
-			expectContent: `[{"name":"config.json","type":"file","size":50},{"name":"model","type":"directory","children":[{"name":"weights.bin","type":"file","size":100}]}]`,
+			expectContent: `[{"name":"model","type":"directory","children":[{"name":"weights.bin","type":"file","size":100}]},{"name":"config.json","type":"file","size":50}]`,
 			expectType:    "application/json; charset=utf-8",
 		},
 	}

src/controller/artifact/processor/cnai/parser/base.go

@@ -20,10 +20,16 @@ import (
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 
+	"github.com/goharbor/harbor/src/lib/errors"
 	"github.com/goharbor/harbor/src/pkg/artifact"
 	"github.com/goharbor/harbor/src/pkg/registry"
 )
 
+var (
+	// errFileTooLarge is returned when the file is too large to be processed.
+	errFileTooLarge = errors.New("The file is too large to be processed")
+)
+
 const (
 	// contentTypeTextPlain is the content type of text/plain.
 	contentTypeTextPlain = "text/plain; charset=utf-8"
@@ -31,6 +37,9 @@ const (
 	contentTypeMarkdown = "text/markdown; charset=utf-8"
 	// contentTypeJSON is the content type of application/json.
 	contentTypeJSON = "application/json; charset=utf-8"
+
+	// defaultFileSizeLimit is the default file size limit.
+	defaultFileSizeLimit = 1024 * 1024 * 4 // 4MB
 )
 
 // newBase creates a new base parser.
@@ -51,6 +60,10 @@ func (b *base) Parse(_ context.Context, artifact *artifact.Artifact, layer *ocis
 		return "", nil, fmt.Errorf("artifact or manifest cannot be nil")
 	}
 
+	if layer.Size > defaultFileSizeLimit {
+		return "", nil, errors.RequestEntityTooLargeError(errFileTooLarge)
+	}
+
 	_, stream, err := b.regCli.PullBlob(artifact.RepositoryName, layer.Digest.String())
 	if err != nil {
 		return "", nil, fmt.Errorf("failed to pull blob from registry: %w", err)

src/controller/artifact/processor/cnai/parser/files.go

@@ -100,8 +100,12 @@ func traverseFileNode(node *FileNode) []FileList {
 		})
 	}
 
-	// sort the children by name.
+	// sort the children by type (directories first) and then by name.
 	sort.Slice(children, func(i, j int) bool {
+		if children[i].Type != children[j].Type {
+			return children[i].Type == TypeDirectory
+		}
+
 		return children[i].Name < children[j].Name
 	})
 

src/controller/artifact/processor/cnai/parser/files_test.go

@@ -141,11 +141,6 @@ func TestFilesParser(t *testing.T) {
 			},
 			expectedType: contentTypeJSON,
 			expectedOutput: []FileList{
-				{
-					Name: "README.md",
-					Type: TypeFile,
-					Size: 100,
-				},
 				{
 					Name: "models",
 					Type: TypeDirectory,
@@ -174,6 +169,11 @@ func TestFilesParser(t *testing.T) {
 						},
 					},
 				},
+				{
+					Name: "README.md",
+					Type: TypeFile,
+					Size: 100,
+				},
 			},
 		},
 		{

src/go.mod

@@ -79,7 +79,7 @@ require (
 )
 
 require (
-	github.com/CloudNativeAI/model-spec v0.0.1
+	github.com/CloudNativeAI/model-spec v0.0.3
 	github.com/prometheus/client_model v0.6.1
 	go.pinniped.dev v0.37.0
 )

src/go.sum

@@ -40,8 +40,8 @@ github.com/Azure/go-ntlmssp v0.0.0-20221128193559-754e69321358/go.mod h1:chxPXzS
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/toml v1.2.1/go.mod h1:CxXYINrC8qIiEnFrOxCa7Jy5BFHlXnUU2pbicEuybxQ=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
-github.com/CloudNativeAI/model-spec v0.0.1 h1:BgVIStKTLuL1DrLC5A/gmHcR8TEhFCDz9+fYdCUa/CY=
+github.com/CloudNativeAI/model-spec v0.0.3 h1:5mvgFQ+3pyupzxYjtV5XAeg9zRe6+46pLPBFNHlOrqE=
-github.com/CloudNativeAI/model-spec v0.0.1/go.mod h1:3U/4zubBfbUkW59ATSg41HnkYyKrKUcKFH/cVdoPQnk=
+github.com/CloudNativeAI/model-spec v0.0.3/go.mod h1:3U/4zubBfbUkW59ATSg41HnkYyKrKUcKFH/cVdoPQnk=
 github.com/FZambia/sentinel v1.1.0 h1:qrCBfxc8SvJihYNjBWgwUI93ZCvFe/PJIPTHKmlp8a8=
 github.com/FZambia/sentinel v1.1.0/go.mod h1:ytL1Am/RLlAoAXG6Kj5LNuw/TRRQrv2rt2FT26vP5gI=
 github.com/HdrHistogram/hdrhistogram-go v1.1.2/go.mod h1:yDgFjdqOqDEKOvasDdhWNXYg9BVp4O+o5f6V/ehm6Oo=

src/lib/errors/const.go

@@ -47,6 +47,8 @@ const (
 	MANIFESTINVALID = "MANIFEST_INVALID"
 	// UNSUPPORTED is for digest UNSUPPORTED error
 	UNSUPPORTED = "UNSUPPORTED"
+	// RequestEntityTooLargeCode is the error code for request entity too large error.
+	RequestEntityTooLargeCode = "REQUEST_ENTITY_TOO_LARGE"
 )
 
 // NotFoundError is error for the case of object not found
@@ -94,6 +96,11 @@ func UnknownError(err error) *Error {
 	return New("unknown").WithCode(GeneralCode).WithCause(err)
 }
 
+// RequestEntityTooLargeError is error for the case of request entity too large.
+func RequestEntityTooLargeError(err error) *Error {
+	return New("request entity too large").WithCode(RequestEntityTooLargeCode).WithCause(err)
+}
+
 // IsNotFoundErr returns true when the error is NotFoundError
 func IsNotFoundErr(err error) bool {
 	return IsErr(err, NotFoundCode)

src/lib/http/error.go

@@ -43,6 +43,7 @@ var (
 		errors.ViolateForeignKeyConstraintCode: http.StatusPreconditionFailed,
 		errors.PROJECTPOLICYVIOLATION:          http.StatusPreconditionFailed,
 		errors.GeneralCode:                     http.StatusInternalServerError,
+		errors.RequestEntityTooLargeCode:       http.StatusRequestEntityTooLarge,
 	}
 )