Improve S3 remote URL handling

Author: amisevsk

pkg/artifact/kitfile.go

@@ -137,6 +137,9 @@ func (kf *KitFile) LoadModel(kitfileContent io.ReadCloser) error {
 	if err := decoder.Decode(kf); err != nil {
 		return err
 	}
+	if err := kf.Validate(); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -210,6 +213,10 @@ func (kf *KitFile) Validate() error {
 			if dataset.RemoteHash == "" {
 				return fmt.Errorf("remoteHash is required when remote dataset paths are used (%s)", dataset.RemotePath)
 			}
+		} else {
+			if dataset.RemoteHash != "" {
+				return fmt.Errorf("remote hash is only applicable when remotePath is set")
+			}
 		}
 	}
 	for _, doc := range kf.Docs {

pkg/lib/external/s3api/api.go

@@ -65,6 +65,13 @@ func ParseS3ObjectReference(ref string, hash string) (*S3ObjectReference, error)
 	bucketKey := strings.TrimPrefix(s3url.Path, "/")
 	version := s3url.Query().Get("versionId")
 
+	if bucketName == "" {
+		return nil, fmt.Errorf("bucket name is required for S3 references")
+	}
+	if bucketKey == "" {
+		return nil, fmt.Errorf("bucket key is required for S3 references")
+	}
+
 	return &S3ObjectReference{
 		Bucket:  bucketName,
 		Key:     bucketKey,
@@ -104,6 +111,7 @@ func DownloadObject(ctx context.Context, client S3ClientAPI, ref *S3ObjectRefere
 	if err != nil {
 		return fmt.Errorf("failed to get object from S3 bucket: %w", err)
 	}
+	defer obj.Body.Close()
 
 	outfile, err := os.Create(outputPath)
 	if err != nil {
@@ -143,6 +151,9 @@ func SetUpClient(ctx context.Context) (S3ClientAPI, error) {
 		// For now, use path style URLs (endpoint.com/bucket/key) instead of virtual-hosted style (bucket.endpoint.com/key) when
 		// an alternate endpoint is specified; this is used/supported by most S3-compatible APIs
 		clientOpts = append(clientOpts, func(o *s3.Options) { o.UsePathStyle = true })
+	} else {
+		// If we're using the default endpoint, enable UseARNRegion to allow finding buckets across regions
+		clientOpts = append(clientOpts, func(o *s3.Options) { o.UseARNRegion = true })
 	}
 
 	s3cfg, err := s3config.LoadDefaultConfig(ctx, cfgOpts...)

pkg/lib/filesystem/unpack/core.go

@@ -176,17 +176,20 @@ func unpackRecursive(ctx context.Context, opts *UnpackOptions, visitedRefs []str
 
 		case mediatype.DatasetBaseType:
 			// Since some datasets may be remote, we need to search the Kitfile for the next non-remote dataset
-			var entry artifact.DataSet
+			var entry *artifact.DataSet
 			for idx := datasetIdx; idx < len(config.DataSets); idx++ {
 				dataset := config.DataSets[idx]
 				if dataset.RemotePath != "" {
 					continue
 				}
-				entry = dataset
+				entry = &dataset
 				datasetIdx = idx + 1
 				break
 			}
-			if !shouldUnpackLayer(entry, opts.FilterConfs) {
+			if entry == nil {
+				continue
+			}
+			if !shouldUnpackLayer(*entry, opts.FilterConfs) {
 				continue
 			}
 			layerInfo, layerPath = entry.LayerInfo, entry.Path
@@ -249,15 +252,30 @@ func unpackRecursive(ctx context.Context, opts *UnpackOptions, visitedRefs []str
 				return err
 			}
 			for path, s3Ref := range remoteFiles {
+				_, relPath, err := filesystem.VerifySubpath(opts.UnpackDir, path)
+				if err != nil {
+					return fmt.Errorf("error verifying path %s for remote reference: %w", path, err)
+				}
+
 				output.Debugf("Downloading remote dataset: Bucket: %s, Key: %s", s3Ref.Bucket, s3Ref.Key)
-				if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) && !opts.Overwrite {
-					return fmt.Errorf("failed to unpack remote dataset: path '%s' already exists", path)
+				if fi, exists := filesystem.PathExists(relPath); exists {
+					if opts.IgnoreExisting {
+						output.Debugf("File %s already exists; skipping", path)
+						continue
+					}
+					if !opts.Overwrite {
+						return fmt.Errorf("failed to unpack remote dataset: path '%s' already exists", path)
+					}
+					if !fi.Mode().IsRegular() {
+						return fmt.Errorf("failed to unpack remote dataset: path '%s' already exists and is not a regular file", path)
+					}
 				}
-				pathDir := filepath.Dir(path)
+
+				pathDir := filepath.Dir(relPath)
 				if err := os.MkdirAll(pathDir, 0755); err != nil {
 					return fmt.Errorf("failed to create directory %s: %w", pathDir, err)
 				}
-				if err := s3api.DownloadObject(ctx, client, &s3Ref, path); err != nil {
+				if err := s3api.DownloadObject(ctx, client, &s3Ref, relPath); err != nil {
 					return fmt.Errorf("failed to download remote dataset for path %s: %w", path, err)
 				}
 				output.Infof("Downloaded remote S3 dataset for path %s", path)