feat: upload size server-size validation (#1001)

Signed-off-by: Miguel Martinez Trivino <[email protected]>

Author: migmartri

app/artifact-cas/internal/server/grpc_test.go

@@ -96,7 +96,7 @@ func TestJWTAuthFunc(t *testing.T) {
 
 			b, err := robotaccount.NewBuilder(opts...)
 			require.NoError(t, err)
-			token, err := b.GenerateJWT("backend-type", "secret-id", tc.audience, robotaccount.Downloader)
+			token, err := b.GenerateJWT("backend-type", "secret-id", tc.audience, robotaccount.Downloader, 0)
 			require.NoError(t, err)
 
 			// add bearer token to context

app/artifact-cas/internal/service/bytestream.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2023 The Chainloop Authors.
+// Copyright 2024 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -92,7 +92,7 @@ func (s *ByteStreamService) Write(stream bytestream.ByteStream_WriteServer) erro
 
 	s.log.Infow("msg", "artifact does not exist, uploading", "digest", req.resource.Digest, "name", req.resource.FileName)
 	// Create a buffer that will be filled in the background before sending its content to the backend
-	buffer := newStreamReader()
+	buffer := newStreamReader(info.MaxBytes)
 	// Add data from first request
 	if err = buffer.Write(req.GetData()); err != nil {
 		return sl.LogAndMaskErr(err, s.log)
@@ -213,7 +213,7 @@ func bufferStream(ctx context.Context, stream bytestream.ByteStream_WriteServer,
 				return
 			}
 
-			log.Debugw("msg", "upload chunk received", "digest", req.resource.Digest, "bufferSize", buffer.size, "chunkSize", len(req.GetData()))
+			log.Debugw("msg", "upload chunk received", "digest", req.resource.Digest, "currentSize", buffer.size, "maxSize", buffer.maxSize, "chunkSize", len(req.GetData()))
 		}
 	}
 }
@@ -222,22 +222,32 @@ type streamReader struct {
 	*bytes.Buffer
 	// total size of the in-memory buffer in bytes
 	size int64
+	// Max size allowed to be uploaded
+	maxSize int64
 	// there was an error during stream data filling
 	errorChan chan error
 }
 
 // Wrapper around a buffer that adds
 // the ability to record the total size of the data that went through it
 // and a channel to be used by the clients to signal when the buffer has been filled
-func newStreamReader() *streamReader {
+func newStreamReader(maxSize int64) *streamReader {
 	return &streamReader{
 		Buffer:    bytes.NewBuffer(nil),
 		errorChan: make(chan error),
+		maxSize:   maxSize,
 	}
 }
 
 func (r *streamReader) Write(data []byte) error {
 	r.size += int64(len(data))
+
+	// Check if the size of the buffer has exceeded the maximum allowed size
+	// if maxSize is 0, then there is no limit
+	if r.maxSize != 0 && r.size > r.maxSize {
+		return fmt.Errorf("max size of upload exceeded: want=%d, max=%d", r.size, r.maxSize)
+	}
+
 	_, err := r.Buffer.Write(data)
 	return err
 }

app/artifact-cas/internal/service/bytestream_test.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2023 The Chainloop Authors.
+// Copyright 2024 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -46,7 +46,7 @@ import (
 )
 
 func (s *bytestreamSuite) TestStreamReader() {
-	buffer := newStreamReader()
+	buffer := newStreamReader(0)
 	// Write twice and check the length
 	err := buffer.Write([]byte("hello"))
 	s.NoError(err)
@@ -68,6 +68,18 @@ func (s *bytestreamSuite) TestStreamReader() {
 	s.Equal(0, buffer.Len())
 }
 
+func (s *bytestreamSuite) TestStreamReaderOverflow() {
+	// a buffer with 8 bytes limit
+	buffer := newStreamReader(8)
+	// Write twice and check the length
+	err := buffer.Write([]byte("hello"))
+	s.NoError(err)
+	s.Equal(int64(5), buffer.size)
+	err = buffer.Write([]byte("chainloop"))
+	s.Error(err)
+	s.ErrorContains(err, "max size of upload exceeded")
+}
+
 func (s *bytestreamSuite) TestWrite() {
 	ctx := s.upCtx
 

app/controlplane/internal/service/attestation.go

@@ -336,7 +336,7 @@ func (s *AttestationService) GetUploadCreds(ctx context.Context, req *cpAPI.Atte
 	// Return the backend information and associated credentials (if applicable)
 	resp := &cpAPI.AttestationServiceGetUploadCredsResponse_Result{Backend: bizCASBackendToPb(backend)}
 	if backend.SecretName != "" {
-		ref := &biz.CASCredsOpts{BackendType: string(backend.Provider), SecretPath: backend.SecretName, Role: casJWT.Uploader}
+		ref := &biz.CASCredsOpts{BackendType: string(backend.Provider), SecretPath: backend.SecretName, Role: casJWT.Uploader, MaxBytes: backend.Limits.MaxBytes}
 		t, err := s.casCredsUseCase.GenerateTemporaryCredentials(ref)
 		if err != nil {
 			return nil, handleUseCaseErr(err, s.log)

app/controlplane/internal/service/cascredential.go

@@ -121,7 +121,7 @@ func (s *CASCredentialsService) Get(ctx context.Context, req *pb.CASCredentialsS
 		return nil, errors.BadRequest("invalid argument", "cannot upload or download artifacts from an inline CAS backend")
 	}
 
-	ref := &biz.CASCredsOpts{BackendType: string(backend.Provider), SecretPath: backend.SecretName, Role: role}
+	ref := &biz.CASCredsOpts{BackendType: string(backend.Provider), SecretPath: backend.SecretName, Role: role, MaxBytes: backend.Limits.MaxBytes}
 	t, err := s.casUC.GenerateTemporaryCredentials(ref)
 	if err != nil {
 		return nil, handleUseCaseErr(err, s.log)

app/controlplane/internal/service/casredirect.go

@@ -115,7 +115,7 @@ func (s *CASRedirectService) GetDownloadURL(ctx context.Context, req *pb.GetDown
 
 	// 2- add authentication token to the query params ?t=[token]
 	if backend.SecretName != "" {
-		ref := &biz.CASCredsOpts{BackendType: string(backend.Provider), SecretPath: backend.SecretName, Role: casJWT.Downloader}
+		ref := &biz.CASCredsOpts{BackendType: string(backend.Provider), SecretPath: backend.SecretName, Role: casJWT.Downloader, MaxBytes: backend.Limits.MaxBytes}
 		t, err := s.casCredsUseCase.GenerateTemporaryCredentials(ref)
 		if err != nil {
 			return nil, handleUseCaseErr(err, s.log)

app/controlplane/pkg/biz/cascredentials.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2023 The Chainloop Authors.
+// Copyright 2024 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -47,8 +47,9 @@ type CASCredsOpts struct {
 	BackendType string // i.e OCI, S3
 	SecretPath  string // path to for example the OCI secret in the vault
 	Role        robotaccount.Role
+	MaxBytes    int64
 }
 
 func (uc *CASCredentialsUseCase) GenerateTemporaryCredentials(backendRef *CASCredsOpts) (string, error) {
-	return uc.jwtBuilder.GenerateJWT(backendRef.BackendType, backendRef.SecretPath, jwt.CASAudience, backendRef.Role)
+	return uc.jwtBuilder.GenerateJWT(backendRef.BackendType, backendRef.SecretPath, jwt.CASAudience, backendRef.Role, backendRef.MaxBytes)
 }

internal/robotaccount/cas/robotaccount.go

@@ -37,6 +37,7 @@ type Claims struct {
 	Role           Role   `json:"role"`      // either downloader or uploader
 	StoredSecretID string `json:"secret-id"` // path to the OCI secret in the vault
 	BackendType    string `json:"backend"`   // backend to use, i.e OCI
+	MaxBytes       int64  `json:"maxbytes"`  // max bytes to upload
 }
 
 type Role string
@@ -102,7 +103,7 @@ func NewBuilder(opts ...NewOpt) (*Builder, error) {
 	return b, nil
 }
 
-func (ra *Builder) GenerateJWT(backendType, secretID, audience string, role Role) (string, error) {
+func (ra *Builder) GenerateJWT(backendType, secretID, audience string, role Role, maxBytes int64) (string, error) {
 	if backendType == "" {
 		return "", fmt.Errorf("backend type is required")
 	}
@@ -131,6 +132,11 @@ func (ra *Builder) GenerateJWT(backendType, secretID, audience string, role Role
 		},
 	}
 
+	// If there is limit on the size of the upload we store it as claim
+	if maxBytes != 0 {
+		claims.MaxBytes = maxBytes
+	}
+
 	if ra.expiration != nil {
 		claims.ExpiresAt = jwt.NewNumericDate(time.Now().Add(*ra.expiration))
 	}

internal/robotaccount/cas/robotaccount_test.go

@@ -150,7 +150,7 @@ func TestGenerateJWT(t *testing.T) {
 	)
 
 	require.NoError(t, err)
-	token, err := b.GenerateJWT("OCI", "secret-id", JWTAudience, Uploader)
+	token, err := b.GenerateJWT("OCI", "secret-id", JWTAudience, Uploader, 123)
 	assert.NoError(t, err)
 	assert.NotEmpty(t, token)
 
@@ -166,6 +166,7 @@ func TestGenerateJWT(t *testing.T) {
 	assert.Equal(t, Uploader, claims.Role)
 	assert.Equal(t, "my-issuer", claims.Issuer)
 	assert.Contains(t, claims.Audience, "artifact-cas.chainloop")
+	assert.Equal(t, claims.MaxBytes, int64(123))
 	assert.WithinDuration(t, time.Now(), claims.ExpiresAt.Time, 10*time.Second)
 }