Add extra header to signed url (flyteorg#4971)

Signed-off-by: Kevin Su <[email protected]>
yubofredwang · Mar 26, 2024 · 8078467 · 8078467
1 parent 8316b44
commit 8078467
Show file tree

Hide file tree

Showing 35 changed files with 484 additions and 270 deletions.
diff --git a/datacatalog/go.mod b/datacatalog/go.mod
@@ -43,7 +43,7 @@ require (
 	github.com/emicklei/go-restful/v3 v3.9.0 // indirect
 	github.com/evanphx/json-patch/v5 v5.6.0 // indirect
 	github.com/fatih/color v1.13.0 // indirect
-	github.com/flyteorg/stow v0.3.8 // indirect
+	github.com/flyteorg/stow v0.3.9 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/ghodss/yaml v1.0.0 // indirect
 	github.com/go-gormigrate/gormigrate/v2 v2.1.1 // indirect

diff --git a/datacatalog/go.sum b/datacatalog/go.sum
@@ -113,8 +113,8 @@ github.com/evanphx/json-patch/v5 v5.6.0 h1:b91NhWfaz02IuVxO9faSllyAtNXHMPkC5J8sJ
 github.com/evanphx/json-patch/v5 v5.6.0/go.mod h1:G79N1coSVB93tBe7j6PhzjmR3/2VvlbKOFpnXhI9Bw4=
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
-github.com/flyteorg/stow v0.3.8 h1:4a6BtfgDR86fUwa48DkkZTcp6WK4oQXSfewPd/kN0Z4=
-github.com/flyteorg/stow v0.3.8/go.mod h1:fArjMpsYJNWkp/hyDKKdbcv07gxbuLmKFcb7YT1aSOM=
+github.com/flyteorg/stow v0.3.9 h1:oA9tKMSvdCBnTnVpx8vncwrjRf8h/w8Ks9c4VGm2Img=
+github.com/flyteorg/stow v0.3.9/go.mod h1:fArjMpsYJNWkp/hyDKKdbcv07gxbuLmKFcb7YT1aSOM=
 github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4HY=
 github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
 github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=

diff --git a/flyteadmin/dataproxy/service.go b/flyteadmin/dataproxy/service.go
@@ -49,11 +49,13 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 	// If it doesn't exist, then proceed as normal.
 
 	if len(req.Project) == 0 || len(req.Domain) == 0 {
+		logger.Infof(ctx, "project and domain are required parameters. Project [%v]. Domain [%v]", req.Project, req.Domain)
 		return nil, errors.NewFlyteAdminErrorf(codes.InvalidArgument, "project and domain are required parameters")
 	}
 
 	// At least one of the hash or manually given prefix must be provided.
 	if len(req.FilenameRoot) == 0 && len(req.ContentMd5) == 0 {
+		logger.Infof(ctx, "content_md5 or filename_root is a required parameter. FilenameRoot [%v], ContentMD5 [%v]", req.FilenameRoot, req.ContentMd5)
 		return nil, errors.NewFlyteAdminErrorf(codes.InvalidArgument,
 			"content_md5 or filename_root is a required parameter")
 	}
@@ -63,10 +65,12 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 		knownLocation, err := createStorageLocation(ctx, s.dataStore, s.cfg.Upload,
 			req.Project, req.Domain, req.FilenameRoot, req.Filename)
 		if err != nil {
+			logger.Errorf(ctx, "failed to create storage location. Error %v", err)
 			return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to create storage location, Error: %v", err)
 		}
 		metadata, err := s.dataStore.Head(ctx, knownLocation)
 		if err != nil {
+			logger.Errorf(ctx, "failed to check if file exists at location [%s], Error: %v", knownLocation.String(), err)
 			return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to check if file exists at location [%s], Error: %v", knownLocation.String(), err)
 		}
 		if metadata.Exists() {
@@ -76,12 +80,18 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 			if len(req.ContentMd5) == 0 {
 				return nil, errors.NewFlyteAdminErrorf(codes.AlreadyExists, "file already exists at location [%v], specify a matching hash if you wish to rewrite", knownLocation)
 			}
-			// Re-encode the hash 3-ways to support matching, hex, base32 and base64
-			hexDigest := hex.EncodeToString(req.ContentMd5)
-			base32Digest := base32.StdEncoding.EncodeToString(req.ContentMd5)
 			base64Digest := base64.StdEncoding.EncodeToString(req.ContentMd5)
-			if hexDigest != metadata.Etag() && base32Digest != metadata.Etag() && base64Digest != metadata.Etag() {
-				logger.Debugf(ctx, "File already exists at location [%v] but hashes do not match", knownLocation)
+			if len(metadata.ContentMD5()) == 0 {
+				// For backward compatibility, dataproxy assumes that the Etag exists if ContentMD5 is not in the metadata.
+				// Data proxy won't allow people to overwrite the file if both the Etag and the ContentMD5 do not exist.
+				hexDigest := hex.EncodeToString(req.ContentMd5)
+				base32Digest := base32.StdEncoding.EncodeToString(req.ContentMd5)
+				if hexDigest != metadata.Etag() && base32Digest != metadata.Etag() && base64Digest != metadata.Etag() {
+					logger.Errorf(ctx, "File already exists at location [%v] but hashes do not match", knownLocation)
+					return nil, errors.NewFlyteAdminErrorf(codes.AlreadyExists, "file already exists at location [%v], specify a matching hash if you wish to rewrite", knownLocation)
+				}
+			} else if base64Digest != metadata.ContentMD5() {
+				logger.Errorf(ctx, "File already exists at location [%v] but hashes do not match", knownLocation)
 				return nil, errors.NewFlyteAdminErrorf(codes.AlreadyExists, "file already exists at location [%v], specify a matching hash if you wish to rewrite", knownLocation)
 			}
 			logger.Debugf(ctx, "File already exists at location [%v] but allowing rewrite", knownLocation)
@@ -105,7 +115,7 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 		req.Filename = rand.String(s.cfg.Upload.DefaultFileNameLength)
 	}
 
-	md5 := base64.StdEncoding.EncodeToString(req.ContentMd5)
+	base64digestMD5 := base64.StdEncoding.EncodeToString(req.ContentMd5)
 
 	var prefix string
 	if len(req.FilenameRoot) > 0 {
@@ -117,23 +127,27 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 	storagePath, err := createStorageLocation(ctx, s.dataStore, s.cfg.Upload,
 		req.Project, req.Domain, prefix, req.Filename)
 	if err != nil {
+		logger.Errorf(ctx, "failed to create shardedStorageLocation. Error %v", err)
 		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to create shardedStorageLocation, Error: %v", err)
 	}
 
 	resp, err := s.dataStore.CreateSignedURL(ctx, storagePath, storage.SignedURLProperties{
-		Scope:      stow.ClientMethodPut,
-		ExpiresIn:  req.ExpiresIn.AsDuration(),
-		ContentMD5: md5,
+		Scope:                 stow.ClientMethodPut,
+		ExpiresIn:             req.ExpiresIn.AsDuration(),
+		ContentMD5:            base64digestMD5,
+		AddContentMD5Metadata: req.AddContentMd5Metadata,
 	})
 
 	if err != nil {
+		logger.Errorf(ctx, "failed to create signed url. Error:", err)
 		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to create a signed url. Error: %v", err)
 	}
 
 	return &service.CreateUploadLocationResponse{
 		SignedUrl: resp.URL.String(),
 		NativeUrl: storagePath.String(),
 		ExpiresAt: timestamppb.New(time.Now().Add(req.ExpiresIn.AsDuration())),
+		Headers:   resp.RequiredRequestHeaders,
 	}, nil
 }
 

diff --git a/flyteadmin/go.mod b/flyteadmin/go.mod
@@ -17,7 +17,7 @@ require (
 	github.com/flyteorg/flyte/flyteplugins v0.0.0-00010101000000-000000000000
 	github.com/flyteorg/flyte/flytepropeller v0.0.0-00010101000000-000000000000
 	github.com/flyteorg/flyte/flytestdlib v0.0.0-00010101000000-000000000000
-	github.com/flyteorg/stow v0.3.8
+	github.com/flyteorg/stow v0.3.9
 	github.com/ghodss/yaml v1.0.0
 	github.com/go-gormigrate/gormigrate/v2 v2.1.1
 	github.com/golang-jwt/jwt v3.2.2+incompatible

diff --git a/flyteadmin/go.sum b/flyteadmin/go.sum
@@ -237,8 +237,8 @@ github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/felixge/httpsnoop v1.0.3 h1:s/nj+GCswXYzN5v2DpNMuMQYe+0DDwt5WVCU6CWBdXk=
 github.com/felixge/httpsnoop v1.0.3/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
-github.com/flyteorg/stow v0.3.8 h1:4a6BtfgDR86fUwa48DkkZTcp6WK4oQXSfewPd/kN0Z4=
-github.com/flyteorg/stow v0.3.8/go.mod h1:fArjMpsYJNWkp/hyDKKdbcv07gxbuLmKFcb7YT1aSOM=
+github.com/flyteorg/stow v0.3.9 h1:oA9tKMSvdCBnTnVpx8vncwrjRf8h/w8Ks9c4VGm2Img=
+github.com/flyteorg/stow v0.3.9/go.mod h1:fArjMpsYJNWkp/hyDKKdbcv07gxbuLmKFcb7YT1aSOM=
 github.com/fogleman/gg v1.2.1-0.20190220221249-0403632d5b90/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
 github.com/fogleman/gg v1.3.0/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
 github.com/fortytw2/leaktest v1.3.0 h1:u8491cBMTQ8ft8aeV+adlcytMZylmA5nnwwkRZjI8vw=

diff --git a/flyteadmin/pkg/data/implementations/noop_remote_url_test.go b/flyteadmin/pkg/data/implementations/noop_remote_url_test.go
@@ -14,6 +14,10 @@ const noopFileSize = int64(1256)
 
 type MockMetadata struct{}
 
+func (m MockMetadata) ContentMD5() string {
+	return ""
+}
+
 func (m MockMetadata) Exists() bool {
 	return true
 }

diff --git a/flytecopilot/go.mod b/flytecopilot/go.mod
@@ -39,7 +39,7 @@ require (
 	github.com/emicklei/go-restful/v3 v3.9.0 // indirect
 	github.com/evanphx/json-patch/v5 v5.6.0 // indirect
 	github.com/fatih/color v1.13.0 // indirect
-	github.com/flyteorg/stow v0.3.8 // indirect
+	github.com/flyteorg/stow v0.3.9 // indirect
 	github.com/go-logr/logr v1.2.4 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-openapi/jsonpointer v0.19.6 // indirect

diff --git a/flytecopilot/go.sum b/flytecopilot/go.sum
@@ -103,8 +103,8 @@ github.com/evanphx/json-patch/v5 v5.6.0 h1:b91NhWfaz02IuVxO9faSllyAtNXHMPkC5J8sJ
 github.com/evanphx/json-patch/v5 v5.6.0/go.mod h1:G79N1coSVB93tBe7j6PhzjmR3/2VvlbKOFpnXhI9Bw4=
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
-github.com/flyteorg/stow v0.3.8 h1:4a6BtfgDR86fUwa48DkkZTcp6WK4oQXSfewPd/kN0Z4=
-github.com/flyteorg/stow v0.3.8/go.mod h1:fArjMpsYJNWkp/hyDKKdbcv07gxbuLmKFcb7YT1aSOM=
+github.com/flyteorg/stow v0.3.9 h1:oA9tKMSvdCBnTnVpx8vncwrjRf8h/w8Ks9c4VGm2Img=
+github.com/flyteorg/stow v0.3.9/go.mod h1:fArjMpsYJNWkp/hyDKKdbcv07gxbuLmKFcb7YT1aSOM=
 github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4HY=
 github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
 github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=

diff --git a/flyteidl/gen/pb-es/flyteidl/service/dataproxy_pb.ts b/flyteidl/gen/pb-es/flyteidl/service/dataproxy_pb.ts