Aws cli compatibility improvement

Author: bogyo210

cache-cli/pkg/storage/s3.go

@@ -9,6 +9,8 @@ import (
 	"github.com/aws/aws-sdk-go-v2/credentials"
 	"github.com/aws/aws-sdk-go-v2/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -26,6 +28,30 @@ type S3StorageOptions struct {
 	Config  StorageConfig
 }
 
+type removeS3OperationIDMiddleware struct{}
+
+func (m *removeS3OperationIDMiddleware) ID() string {
+	return "RemoveS3OperationIDMiddleware"
+}
+
+func (m *removeS3OperationIDMiddleware) HandleBuild(ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler) (
+	out middleware.BuildOutput,
+	metadata middleware.Metadata,
+	err error,
+) {
+	if request, ok := in.Request.(*smithyhttp.Request); ok {
+		query := request.URL.Query()
+		query.Del("x-id")
+		request.URL.RawQuery = query.Encode()
+	}
+
+	return next.HandleBuild(ctx, in)
+}
+
+func removeS3OperationID(stack *middleware.Stack) error {
+	return stack.Build.Add(&removeS3OperationIDMiddleware{}, middleware.After)
+}
+
 func NewS3Storage(options S3StorageOptions) (*S3Storage, error) {
 	if options.URL != "" {
 		return createS3StorageUsingEndpoint(options.Bucket, options.Project, options.URL, options.Config)
@@ -52,7 +78,9 @@ func createDefaultS3Storage(s3Bucket, project string, storageConfig StorageConfi
 		}
 
 		return &S3Storage{
-			Client:        s3.NewFromConfig(config),
+			Client: s3.NewFromConfig(config, func(o *s3.Options) {
+				o.APIOptions = append(o.APIOptions, removeS3OperationID)
+			}),
 			Bucket:        s3Bucket,
 			Project:       project,
 			StorageConfig: storageConfig,
@@ -69,7 +97,9 @@ func createDefaultS3Storage(s3Bucket, project string, storageConfig StorageConfi
 		}
 
 		return &S3Storage{
-			Client:        s3.NewFromConfig(config),
+			Client: s3.NewFromConfig(config, func(o *s3.Options) {
+				o.APIOptions = append(o.APIOptions, removeS3OperationID)
+			}),
 			Bucket:        s3Bucket,
 			Project:       project,
 			StorageConfig: storageConfig,
@@ -83,20 +113,28 @@ func createDefaultS3Storage(s3Bucket, project string, storageConfig StorageConfi
 	}
 
 	return &S3Storage{
-		Client:        s3.NewFromConfig(config),
+		Client: s3.NewFromConfig(config, func(o *s3.Options) {
+			o.APIOptions = append(o.APIOptions, removeS3OperationID)
+		}),
 		Bucket:        s3Bucket,
 		Project:       project,
 		StorageConfig: storageConfig,
 	}, nil
 }
 
 func createS3StorageUsingEndpoint(s3Bucket, project, s3Url string, storageConfig StorageConfig) (*S3Storage, error) {
+	region := os.Getenv("SEMAPHORE_CACHE_S3_REGION")
+	if region == "" {
+		region = "auto"
+	}
+
 	options := []func(*awsConfig.LoadOptions) error{
-		awsConfig.WithRegion("auto"),
+		awsConfig.WithRegion(region),
 		awsConfig.WithEndpointResolverWithOptions(
 			aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) {
 				return aws.Endpoint{
-					URL: s3Url,
+					URL:           s3Url,
+					SigningRegion: region,
 				}, nil
 			}),
 		),
@@ -125,6 +163,7 @@ func createS3StorageUsingEndpoint(s3Bucket, project, s3Url string, storageConfig
 		StorageConfig: storageConfig,
 		Client: s3.NewFromConfig(cfg, func(o *s3.Options) {
 			o.UsePathStyle = true
+			o.APIOptions = append(o.APIOptions, removeS3OperationID)
 		}),
 	}, nil
 }

cache-cli/pkg/storage/s3_restore.go

@@ -3,30 +3,89 @@ package storage
 import (
 	"context"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 
-	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 )
 
+const defaultS3DownloadPartSize = int64(8 * 1024 * 1024)
+
 func (s *S3Storage) Restore(key string) (*os.File, error) {
 	tempFile, err := ioutil.TempFile(os.TempDir(), fmt.Sprintf("%s-*", key))
 	if err != nil {
 		return nil, err
 	}
 
 	bucketKey := fmt.Sprintf("%s/%s", s.Project, key)
-	downloader := manager.NewDownloader(s.Client)
-	_, err = downloader.Download(context.TODO(), tempFile, &s3.GetObjectInput{
+	headOutput, err := s.Client.HeadObject(context.TODO(), &s3.HeadObjectInput{
 		Bucket: &s.Bucket,
 		Key:    &bucketKey,
 	})
-
 	if err != nil {
 		_ = tempFile.Close()
 		return nil, err
 	}
 
+	// Download in 8 MiB ranges to match aws-cli style segmented GET requests.
+	if headOutput.ContentLength > 0 {
+		for start := int64(0); start < headOutput.ContentLength; start += defaultS3DownloadPartSize {
+			end := start + defaultS3DownloadPartSize - 1
+			if end >= headOutput.ContentLength {
+				end = headOutput.ContentLength - 1
+			}
+
+			byteRange := fmt.Sprintf("bytes=%d-%d", start, end)
+			input := &s3.GetObjectInput{
+				Bucket: &s.Bucket,
+				Key:    &bucketKey,
+				Range:  &byteRange,
+			}
+			if headOutput.ETag != nil {
+				input.IfMatch = headOutput.ETag
+			}
+
+			output, err := s.Client.GetObject(context.TODO(), input)
+			if err != nil {
+				_ = tempFile.Close()
+				return nil, err
+			}
+
+			_, err = io.Copy(tempFile, output.Body)
+			closeErr := output.Body.Close()
+			if err != nil {
+				_ = tempFile.Close()
+				return nil, err
+			}
+
+			if closeErr != nil {
+				_ = tempFile.Close()
+				return nil, closeErr
+			}
+		}
+	} else {
+		output, err := s.Client.GetObject(context.TODO(), &s3.GetObjectInput{
+			Bucket: &s.Bucket,
+			Key:    &bucketKey,
+		})
+		if err != nil {
+			_ = tempFile.Close()
+			return nil, err
+		}
+
+		_, err = io.Copy(tempFile, output.Body)
+		closeErr := output.Body.Close()
+		if err != nil {
+			_ = tempFile.Close()
+			return nil, err
+		}
+
+		if closeErr != nil {
+			_ = tempFile.Close()
+			return nil, closeErr
+		}
+	}
+
 	return tempFile, tempFile.Close()
 }

cache-cli/pkg/storage/s3_store.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"os"
 
-	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	log "github.com/sirupsen/logrus"
 )
@@ -18,8 +17,7 @@ func (s *S3Storage) Store(key, path string) error {
 	}
 
 	destination := fmt.Sprintf("%s/%s", s.Project, key)
-	uploader := manager.NewUploader(s.Client)
-	_, err = uploader.Upload(context.TODO(), &s3.PutObjectInput{
+	_, err = s.Client.PutObject(context.TODO(), &s3.PutObjectInput{
 		Bucket: &s.Bucket,
 		Key:    &destination,
 		Body:   file,