Update proto, add file extension check

Author: ddl-rliu

flytecopilot/cmd/download_test.go

@@ -162,6 +162,9 @@ func TestDownloadOptions_Download(t *testing.T) {
 				"csv_blob": {
 					Type: &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{Dimensionality: core.BlobType_SINGLE, Format: "csv", FileExtension: "csv"}}},
 				},
+				"invalid_extension_blob": {
+					Type: &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{Dimensionality: core.BlobType_SINGLE, Format: "xyz", FileExtension: "/../../../invalid"}}},
+				},
 			},
 		}
 		d, err := proto.Marshal(iface)
@@ -207,10 +210,26 @@ func TestDownloadOptions_Download(t *testing.T) {
 						},
 					},
 				}},
+				"invalid_extension_blob": {Value: &core.Literal_Scalar{
+					Scalar: &core.Scalar{
+						Value: &core.Scalar_Blob{
+							Blob: &core.Blob{
+								Uri: blobLoc.String(),
+								Metadata: &core.BlobMetadata{
+									Type: &core.BlobType{
+										Dimensionality: core.BlobType_SINGLE,
+										Format:         "xyz",
+										FileExtension:  "/../../../invalid",
+									},
+								},
+							},
+						},
+					},
+				}},
 			},
 		}))
 		assert.NoError(t, dopts.Download(ctx), "Download Operation failed")
-		assert.ElementsMatch(t, []string{"inputs.json", "inputs.pb", "x", "y", "blob.xyz", "csv_blob.csv"}, collectFile(tmpDir))
+		assert.ElementsMatch(t, []string{"inputs.json", "inputs.pb", "x", "y", "blob.xyz", "csv_blob.csv", "invalid_extension_blob"}, collectFile(tmpDir))
 	})
 
 	t.Run("primitiveAndMissingBlobInputs", func(t *testing.T) {

flytecopilot/data/download.go

@@ -9,6 +9,7 @@ import (
 	"path"
 	"path/filepath"
 	"reflect"
+	"regexp"
 	"strconv"
 	"strings"
 	"sync"
@@ -26,6 +27,8 @@ import (
 	"github.com/flyteorg/flyte/flytestdlib/storage"
 )
 
+var validFileExtensionRe = regexp.MustCompile(`^[a-zA-Z0-9]+(\.[a-zA-Z0-9]+)*$`)
+
 type Downloader struct {
 	format core.DataLoadingConfig_LiteralMapFormat
 	store  *storage.DataStore
@@ -42,16 +45,20 @@ type Downloader struct {
 // However, an input blob
 // `data: Annotated[FlyteFile["csv"], FileDownloadConfig(file_extension="csv")]`
 // should be written to "inputs/data.csv" (when FileExtension="csv" - new behavior).
-func resolveVarFilenames(vars *core.VariableMap) map[string]string {
+func resolveVarFilenames(ctx context.Context, vars *core.VariableMap) map[string]string {
 	varFilenames := make(map[string]string, len(vars.GetVariables()))
 	for varName, variable := range vars.GetVariables() {
 		varType := variable.GetType()
 		switch varType.GetType().(type) {
 		case *core.LiteralType_Blob:
-			if varType.GetBlob().GetFileExtension() == "" {
+			ext := varType.GetBlob().GetFileExtension()
+			if ext == "" {
+				varFilenames[varName] = varName
+			} else if !validFileExtensionRe.MatchString(ext) {
+				logger.Warnf(ctx, "invalid file extension for variable %q [%q], ignoring...", varName, ext)
 				varFilenames[varName] = varName
 			} else {
-				varFilenames[varName] = varName + "." + varType.GetBlob().GetFileExtension()
+				varFilenames[varName] = varName + "." + ext
 			}
 		default:
 			varFilenames[varName] = varName
@@ -565,7 +572,7 @@ func (d Downloader) DownloadInputs(ctx context.Context, vars *core.VariableMap,
 		return errors.Wrapf(err, "failed to download input metadata message from remote store")
 	}
 
-	varFilenames := resolveVarFilenames(vars)
+	varFilenames := resolveVarFilenames(ctx, vars)
 	varMap, lMap, err := d.RecursiveDownload(ctx, inputs, outputDir, varFilenames, true)
 	if err != nil {
 		return errors.Wrapf(err, "failed to download input variable from remote store")