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

Signed-off-by: chlins <chlins.zhang@gmail.com>
This commit is contained in:
Chlins Zhang 2025-03-21 15:17:31 +08:00 committed by GitHub
parent 8081d52c09
commit b37da544d2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 35 additions and 10 deletions

View File

@ -207,7 +207,7 @@ func (p *ProcessorTestSuite) TestAbstractAddition() {
p.Require().NoError(err) p.Require().NoError(err)
r.On("PullManifest", mock.Anything, mock.Anything).Return(manifest, "", nil) 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", expectType: "application/json; charset=utf-8",
}, },
} }

View File

@ -20,10 +20,16 @@ import (
ocispec "github.com/opencontainers/image-spec/specs-go/v1" 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/artifact"
"github.com/goharbor/harbor/src/pkg/registry" "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 ( const (
// contentTypeTextPlain is the content type of text/plain. // contentTypeTextPlain is the content type of text/plain.
contentTypeTextPlain = "text/plain; charset=utf-8" contentTypeTextPlain = "text/plain; charset=utf-8"
@ -31,6 +37,9 @@ const (
contentTypeMarkdown = "text/markdown; charset=utf-8" contentTypeMarkdown = "text/markdown; charset=utf-8"
// contentTypeJSON is the content type of application/json. // contentTypeJSON is the content type of application/json.
contentTypeJSON = "application/json; charset=utf-8" contentTypeJSON = "application/json; charset=utf-8"
// defaultFileSizeLimit is the default file size limit.
defaultFileSizeLimit = 1024 * 1024 * 4 // 4MB
) )
// newBase creates a new base parser. // 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") 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()) _, stream, err := b.regCli.PullBlob(artifact.RepositoryName, layer.Digest.String())
if err != nil { if err != nil {
return "", nil, fmt.Errorf("failed to pull blob from registry: %w", err) return "", nil, fmt.Errorf("failed to pull blob from registry: %w", err)

View File

@ -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 { 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 return children[i].Name < children[j].Name
}) })

View File

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

View File

@ -79,7 +79,7 @@ require (
) )
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 github.com/prometheus/client_model v0.6.1
go.pinniped.dev v0.37.0 go.pinniped.dev v0.37.0
) )

View File

@ -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 v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/BurntSushi/toml v1.2.1/go.mod h1:CxXYINrC8qIiEnFrOxCa7Jy5BFHlXnUU2pbicEuybxQ= 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/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 h1:qrCBfxc8SvJihYNjBWgwUI93ZCvFe/PJIPTHKmlp8a8=
github.com/FZambia/sentinel v1.1.0/go.mod h1:ytL1Am/RLlAoAXG6Kj5LNuw/TRRQrv2rt2FT26vP5gI= 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= github.com/HdrHistogram/hdrhistogram-go v1.1.2/go.mod h1:yDgFjdqOqDEKOvasDdhWNXYg9BVp4O+o5f6V/ehm6Oo=

View File

@ -47,6 +47,8 @@ const (
MANIFESTINVALID = "MANIFEST_INVALID" MANIFESTINVALID = "MANIFEST_INVALID"
// UNSUPPORTED is for digest UNSUPPORTED error // UNSUPPORTED is for digest UNSUPPORTED error
UNSUPPORTED = "UNSUPPORTED" 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 // 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) 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 // IsNotFoundErr returns true when the error is NotFoundError
func IsNotFoundErr(err error) bool { func IsNotFoundErr(err error) bool {
return IsErr(err, NotFoundCode) return IsErr(err, NotFoundCode)

View File

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