From 657dfc925fb011d561fd6c85cc718a57b77968c8 Mon Sep 17 00:00:00 2001
From: Iaroslav Ciupin <iaroslav@union.ai>
Date: Thu, 25 Jul 2024 09:19:02 +0300
Subject: [PATCH] CreateDownloadLink: Head before signing

Signed-off-by: Andrew Dye <andrewwdye@gmail.com>
---
 flyteadmin/dataproxy/service.go      | 12 +++++-
 flyteadmin/dataproxy/service_test.go | 59 ++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/flyteadmin/dataproxy/service.go b/flyteadmin/dataproxy/service.go
index cb20075a2e..c02fa3699f 100644
--- a/flyteadmin/dataproxy/service.go
+++ b/flyteadmin/dataproxy/service.go
@@ -182,7 +182,17 @@ func (s Service) CreateDownloadLink(ctx context.Context, req *service.CreateDown
 		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "no deckUrl found for request [%+v]", req)
 	}
 
-	signedURLResp, err := s.dataStore.CreateSignedURL(ctx, storage.DataReference(nativeURL), storage.SignedURLProperties{
+	ref := storage.DataReference(nativeURL)
+	meta, err := s.dataStore.Head(ctx, ref)
+	if err != nil {
+		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to head object before signing url. Error: %v", err)
+	}
+
+	if !meta.Exists() {
+		return nil, errors.NewFlyteAdminErrorf(codes.NotFound, "object not found")
+	}
+
+	signedURLResp, err := s.dataStore.CreateSignedURL(ctx, ref, storage.SignedURLProperties{
 		Scope:     stow.ClientMethodGet,
 		ExpiresIn: req.ExpiresIn.AsDuration(),
 	})
diff --git a/flyteadmin/dataproxy/service_test.go b/flyteadmin/dataproxy/service_test.go
index fb7a956a27..4c3f3ea720 100644
--- a/flyteadmin/dataproxy/service_test.go
+++ b/flyteadmin/dataproxy/service_test.go
@@ -4,12 +4,15 @@ import (
 	"bytes"
 	"context"
 	"crypto/md5" // #nosec
+	"fmt"
 	"net/url"
 	"testing"
 	"time"
 
 	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/durationpb"
 
 	commonMocks "github.com/flyteorg/flyte/flyteadmin/pkg/common/mocks"
@@ -157,6 +160,15 @@ func TestCreateUploadLocationMore(t *testing.T) {
 	})
 }
 
+type testMetadata struct {
+	storage.Metadata
+	exists bool
+}
+
+func (t testMetadata) Exists() bool {
+	return t.exists
+}
+
 func TestCreateDownloadLink(t *testing.T) {
 	dataStore := commonMocks.GetMockStorageClient()
 	nodeExecutionManager := &mocks.MockNodeExecutionManager{}
@@ -179,7 +191,30 @@ func TestCreateDownloadLink(t *testing.T) {
 		assert.Error(t, err)
 	})
 
+	t.Run("item not found", func(t *testing.T) {
+		dataStore.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb = func(ctx context.Context, ref storage.DataReference) (storage.Metadata, error) {
+			return testMetadata{exists: false}, nil
+		}
+
+		_, err = s.CreateDownloadLink(context.Background(), &service.CreateDownloadLinkRequest{
+			ArtifactType: service.ArtifactType_ARTIFACT_TYPE_DECK,
+			Source: &service.CreateDownloadLinkRequest_NodeExecutionId{
+				NodeExecutionId: &core.NodeExecutionIdentifier{},
+			},
+			ExpiresIn: durationpb.New(time.Hour),
+		})
+
+		st, ok := status.FromError(err)
+		assert.True(t, ok)
+		assert.Equal(t, codes.NotFound, st.Code())
+		assert.Equal(t, "object not found", st.Message())
+	})
+
 	t.Run("valid config", func(t *testing.T) {
+		dataStore.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb = func(ctx context.Context, ref storage.DataReference) (storage.Metadata, error) {
+			return testMetadata{exists: true}, nil
+		}
+
 		_, err = s.CreateDownloadLink(context.Background(), &service.CreateDownloadLinkRequest{
 			ArtifactType: service.ArtifactType_ARTIFACT_TYPE_DECK,
 			Source: &service.CreateDownloadLinkRequest_NodeExecutionId{
@@ -187,10 +222,34 @@ func TestCreateDownloadLink(t *testing.T) {
 			},
 			ExpiresIn: durationpb.New(time.Hour),
 		})
+
 		assert.NoError(t, err)
 	})
 
+	t.Run("head failed", func(t *testing.T) {
+		dataStore.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb = func(ctx context.Context, ref storage.DataReference) (storage.Metadata, error) {
+			return testMetadata{}, fmt.Errorf("head fail")
+		}
+
+		_, err = s.CreateDownloadLink(context.Background(), &service.CreateDownloadLinkRequest{
+			ArtifactType: service.ArtifactType_ARTIFACT_TYPE_DECK,
+			Source: &service.CreateDownloadLinkRequest_NodeExecutionId{
+				NodeExecutionId: &core.NodeExecutionIdentifier{},
+			},
+			ExpiresIn: durationpb.New(time.Hour),
+		})
+
+		st, ok := status.FromError(err)
+		assert.True(t, ok)
+		assert.Equal(t, codes.Internal, st.Code())
+		assert.Equal(t, "failed to head object before signing url. Error: head fail", st.Message())
+	})
+
 	t.Run("use default ExpiresIn", func(t *testing.T) {
+		dataStore.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb = func(ctx context.Context, ref storage.DataReference) (storage.Metadata, error) {
+			return testMetadata{exists: true}, nil
+		}
+
 		_, err = s.CreateDownloadLink(context.Background(), &service.CreateDownloadLinkRequest{
 			ArtifactType: service.ArtifactType_ARTIFACT_TYPE_DECK,
 			Source: &service.CreateDownloadLinkRequest_NodeExecutionId{