Add ContentMD5 to fast register upload request (flyteorg#303)

* Pass file hash Signed-off-by: Haytham Abuelfutuh <[email protected]> * Send MD5 as []byte Signed-off-by: Haytham Abuelfutuh <[email protected]> * Update deps Signed-off-by: Haytham Abuelfutuh <[email protected]> * Fix compile error and unit tests Signed-off-by: Haytham Abuelfutuh <[email protected]>
austin362667 · May 7, 2024 · 5d627bd · 5d627bd
1 parent 0162ceb
commit 5d627bd
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 16 deletions.
diff --git a/flytectl/cmd/register/register_util.go b/flytectl/cmd/register/register_util.go
@@ -4,6 +4,8 @@ import (
 	"archive/tar"
 	"compress/gzip"
 	"context"
+	"crypto/md5" //#nosec
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
@@ -705,7 +707,10 @@ func uploadFastRegisterArtifact(ctx context.Context, project, domain, sourceCode
 		return "", err
 	}
 
-	size, err := getTotalSize(dataRefReaderCloser)
+	/* #nosec */
+	hash := md5.New()
+	/* #nosec */
+	size, err := io.Copy(hash, dataRefReaderCloser)
 	if err != nil {
 		return "", err
 	}
@@ -720,12 +725,14 @@ func uploadFastRegisterArtifact(ctx context.Context, project, domain, sourceCode
 		return "", err
 	}
 
+	h := hash.Sum(nil)
 	remotePath := storage.DataReference(deprecatedSourceUploadPath)
 	_, fileName := filepath.Split(sourceCodeFilePath)
 	resp, err := dataProxyClient.CreateUploadLocation(ctx, &service.CreateUploadLocationRequest{
-		Project: project,
-		Domain:  domain,
-		Suffix:  strings.Join([]string{version, fileName}, "/"),
+		Project:    project,
+		Domain:     domain,
+		Filename:   fileName,
+		ContentMd5: h,
 	})
 
 	if err != nil {
@@ -737,7 +744,7 @@ func uploadFastRegisterArtifact(ctx context.Context, project, domain, sourceCode
 	}
 
 	if resp != nil && len(resp.SignedUrl) > 0 {
-		return storage.DataReference(resp.NativeUrl), DirectUpload(resp.SignedUrl, size, dataRefReaderCloser)
+		return storage.DataReference(resp.NativeUrl), DirectUpload(resp.SignedUrl, h, size, dataRefReaderCloser)
 	}
 
 	dataStore, err := getStorageClient(ctx)
@@ -764,14 +771,15 @@ func uploadFastRegisterArtifact(ctx context.Context, project, domain, sourceCode
 	return remotePath, nil
 }
 
-func DirectUpload(url string, size int64, data io.Reader) error {
+func DirectUpload(url string, contentMD5 []byte, size int64, data io.Reader) error {
 	req, err := http.NewRequest(http.MethodPut, url, data)
 	if err != nil {
 		return err
 	}
 
 	req.ContentLength = size
 	req.Header.Set("Content-Length", strconv.FormatInt(size, 10))
+	req.Header.Set("Content-MD5", base64.StdEncoding.EncodeToString(contentMD5))
 
 	client := &http.Client{}
 	res, err := client.Do(req)
@@ -785,7 +793,7 @@ func DirectUpload(url string, size int64, data io.Reader) error {
 			return fmt.Errorf("received response code [%v]. Failed to read response body. Error: %w", res.StatusCode, err)
 		}
 
-		return fmt.Errorf("bad status: %s : %s", res.Status, string(raw))
+		return fmt.Errorf("failed uploading to [%v]. bad status: %s: %s", url, res.Status, string(raw))
 	}
 
 	return nil

diff --git a/flytectl/cmd/register/register_util_test.go b/flytectl/cmd/register/register_util_test.go
@@ -384,9 +384,10 @@ func TestUploadFastRegisterArtifact(t *testing.T) {
 		assert.Nil(t, err)
 		Client = store
 		s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(s.Ctx, &service.CreateUploadLocationRequest{
-			Project: "flytesnacks",
-			Domain:  "development",
-			Suffix:  "/flytesnacks-core.tgz",
+			Project:    "flytesnacks",
+			Domain:     "development",
+			Filename:   "flytesnacks-core.tgz",
+			ContentMd5: []uint8{0x19, 0x72, 0x39, 0xcd, 0x85, 0x2d, 0xf1, 0x79, 0x8f, 0x6b, 0x3, 0xb3, 0xa9, 0x6c, 0xec, 0xa0},
 		}).Return(&service.CreateUploadLocationResponse{}, nil)
 		_, err = uploadFastRegisterArtifact(s.Ctx, "flytesnacks", "development", "testdata/flytesnacks-core.tgz", "", s.MockClient.DataProxyClient(), rconfig.DefaultFilesConfig.DeprecatedSourceUploadPath)
 		assert.Nil(t, err)
@@ -401,9 +402,10 @@ func TestUploadFastRegisterArtifact(t *testing.T) {
 		assert.Nil(t, err)
 		Client = store
 		s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(s.Ctx, &service.CreateUploadLocationRequest{
-			Project: "flytesnacks",
-			Domain:  "development",
-			Suffix:  "/flytesnacks-core.tgz",
+			Project:    "flytesnacks",
+			Domain:     "development",
+			Filename:   "flytesnacks-core.tgz",
+			ContentMd5: []uint8{0x19, 0x72, 0x39, 0xcd, 0x85, 0x2d, 0xf1, 0x79, 0x8f, 0x6b, 0x3, 0xb3, 0xa9, 0x6c, 0xec, 0xa0},
 		}).Return(&service.CreateUploadLocationResponse{}, nil)
 		_, err = uploadFastRegisterArtifact(context.Background(), "flytesnacks", "development", "testdata/flytesnacks-core.tgz", "", s.MockClient.DataProxyClient(), rconfig.DefaultFilesConfig.DeprecatedSourceUploadPath)
 		assert.Nil(t, err)

diff --git a/flytectl/go.mod b/flytectl/go.mod
@@ -9,7 +9,7 @@ require (
 	github.com/docker/docker v20.10.7+incompatible
 	github.com/docker/go-connections v0.4.0
 	github.com/enescakir/emoji v1.0.0
-	github.com/flyteorg/flyteidl v0.24.10
+	github.com/flyteorg/flyteidl v0.24.15
 	github.com/flyteorg/flytestdlib v0.4.16
 	github.com/ghodss/yaml v1.0.0
 	github.com/go-ozzo/ozzo-validation/v4 v4.3.0

diff --git a/flytectl/go.sum b/flytectl/go.sum
@@ -364,8 +364,8 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
 github.com/fatih/color v1.10.0 h1:s36xzo75JdqLaaWoiEHk767eHiwo0598uUxyfiPkDsg=
 github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
-github.com/flyteorg/flyteidl v0.24.10 h1:fCYpfp5fxKbhRMSkP0Hdw5lOPBTItLU1A3WA1Lc7sEU=
-github.com/flyteorg/flyteidl v0.24.10/go.mod h1:vHSugApgS3hRITIafzQDU8DZD/W8wFRfFcgaFU35Dww=
+github.com/flyteorg/flyteidl v0.24.15 h1:Iqbwx3w1a4Dh6byRZrZMlHsKPKoOZbBiS9vR0iXzacY=
+github.com/flyteorg/flyteidl v0.24.15/go.mod h1:vHSugApgS3hRITIafzQDU8DZD/W8wFRfFcgaFU35Dww=
 github.com/flyteorg/flytestdlib v0.3.13/go.mod h1:Tz8JCECAbX6VWGwFT6cmEQ+RJpZ/6L9pswu3fzWs220=
 github.com/flyteorg/flytestdlib v0.4.16 h1:r4dCPUOqoE9xCAhOw9KDB7O6cBoCxyEtepIWYcj93H0=
 github.com/flyteorg/flytestdlib v0.4.16/go.mod h1:WA5Y4hrcgD0ybGOKJVOQ4sP8q7NLRV+S5SWOlH0axgM=