From aabf07ec657fdca1138783edf9f4700d36993af8 Mon Sep 17 00:00:00 2001
From: Bouke van der Bijl <me@bou.ke>
Date: Mon, 2 Sep 2019 12:31:30 +0000
Subject: [PATCH] s3: accept S3 parameter

---
 registry/storage/driver/s3-aws/s3.go      | 110 ++++++++++++----------
 registry/storage/driver/s3-aws/s3_test.go |   1 +
 2 files changed, 60 insertions(+), 51 deletions(-)

diff --git a/registry/storage/driver/s3-aws/s3.go b/registry/storage/driver/s3-aws/s3.go
index b14ef084a..cdc588abf 100644
--- a/registry/storage/driver/s3-aws/s3.go
+++ b/registry/storage/driver/s3-aws/s3.go
@@ -84,6 +84,10 @@ var validObjectACLs = map[string]struct{}{}
 
 //DriverParameters A struct that encapsulates all of the driver parameters after all values have been set
 type DriverParameters struct {
+	// S3 is an optional parameter. If specified, it will use the existing session
+	// to construct the Driver.
+	S3 *s3.S3
+
 	AccessKey                   string
 	SecretKey                   string
 	Bucket                      string
@@ -342,6 +346,7 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 	sessionToken := ""
 
 	params := DriverParameters{
+		nil,
 		fmt.Sprint(accessKey),
 		fmt.Sprint(secretKey),
 		fmt.Sprint(bucket),
@@ -398,66 +403,69 @@ func getParameterAsInt64(parameters map[string]interface{}, name string, default
 // New constructs a new Driver with the given AWS credentials, region, encryption flag, and
 // bucketName
 func New(params DriverParameters) (*Driver, error) {
-	if !params.V4Auth &&
-		(params.RegionEndpoint == "" ||
-			strings.Contains(params.RegionEndpoint, "s3.amazonaws.com")) {
-		return nil, fmt.Errorf("on Amazon S3 this storage driver can only be used with v4 authentication")
-	}
+	s3obj := params.S3
+	if s3obj == nil {
+		if !params.V4Auth &&
+			(params.RegionEndpoint == "" ||
+				strings.Contains(params.RegionEndpoint, "s3.amazonaws.com")) {
+			return nil, fmt.Errorf("on Amazon S3 this storage driver can only be used with v4 authentication")
+		}
 
-	awsConfig := aws.NewConfig()
-	sess, err := session.NewSession()
-	if err != nil {
-		return nil, fmt.Errorf("failed to create new session: %v", err)
-	}
-	creds := credentials.NewChainCredentials([]credentials.Provider{
-		&credentials.StaticProvider{
-			Value: credentials.Value{
-				AccessKeyID:     params.AccessKey,
-				SecretAccessKey: params.SecretKey,
-				SessionToken:    params.SessionToken,
+		awsConfig := aws.NewConfig()
+		sess, err := session.NewSession()
+		if err != nil {
+			return nil, fmt.Errorf("failed to create new session: %v", err)
+		}
+		creds := credentials.NewChainCredentials([]credentials.Provider{
+			&credentials.StaticProvider{
+				Value: credentials.Value{
+					AccessKeyID:     params.AccessKey,
+					SecretAccessKey: params.SecretKey,
+					SessionToken:    params.SessionToken,
+				},
 			},
-		},
-		&credentials.EnvProvider{},
-		&credentials.SharedCredentialsProvider{},
-		&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess)},
-	})
+			&credentials.EnvProvider{},
+			&credentials.SharedCredentialsProvider{},
+			&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess)},
+		})
 
-	if params.RegionEndpoint != "" {
-		awsConfig.WithS3ForcePathStyle(true)
-		awsConfig.WithEndpoint(params.RegionEndpoint)
-	}
+		if params.RegionEndpoint != "" {
+			awsConfig.WithS3ForcePathStyle(true)
+			awsConfig.WithEndpoint(params.RegionEndpoint)
+		}
 
-	awsConfig.WithCredentials(creds)
-	awsConfig.WithRegion(params.Region)
-	awsConfig.WithDisableSSL(!params.Secure)
+		awsConfig.WithCredentials(creds)
+		awsConfig.WithRegion(params.Region)
+		awsConfig.WithDisableSSL(!params.Secure)
 
-	if params.UserAgent != "" || params.SkipVerify {
-		httpTransport := http.DefaultTransport
-		if params.SkipVerify {
-			httpTransport = &http.Transport{
-				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		if params.UserAgent != "" || params.SkipVerify {
+			httpTransport := http.DefaultTransport
+			if params.SkipVerify {
+				httpTransport = &http.Transport{
+					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+				}
+			}
+			if params.UserAgent != "" {
+				awsConfig.WithHTTPClient(&http.Client{
+					Transport: transport.NewTransport(httpTransport, transport.NewHeaderRequestModifier(http.Header{http.CanonicalHeaderKey("User-Agent"): []string{params.UserAgent}})),
+				})
+			} else {
+				awsConfig.WithHTTPClient(&http.Client{
+					Transport: transport.NewTransport(httpTransport),
+				})
 			}
 		}
-		if params.UserAgent != "" {
-			awsConfig.WithHTTPClient(&http.Client{
-				Transport: transport.NewTransport(httpTransport, transport.NewHeaderRequestModifier(http.Header{http.CanonicalHeaderKey("User-Agent"): []string{params.UserAgent}})),
-			})
-		} else {
-			awsConfig.WithHTTPClient(&http.Client{
-				Transport: transport.NewTransport(httpTransport),
-			})
+
+		sess, err = session.NewSession(awsConfig)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create new session with aws config: %v", err)
 		}
-	}
+		s3obj = s3.New(sess)
 
-	sess, err = session.NewSession(awsConfig)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create new session with aws config: %v", err)
-	}
-	s3obj := s3.New(sess)
-
-	// enable S3 compatible signature v2 signing instead
-	if !params.V4Auth {
-		setv2Handlers(s3obj)
+		// enable S3 compatible signature v2 signing instead
+		if !params.V4Auth {
+			setv2Handlers(s3obj)
+		}
 	}
 
 	// TODO Currently multipart uploads have no timestamps, so this would be unwise
diff --git a/registry/storage/driver/s3-aws/s3_test.go b/registry/storage/driver/s3-aws/s3_test.go
index be02772e9..d8235f657 100644
--- a/registry/storage/driver/s3-aws/s3_test.go
+++ b/registry/storage/driver/s3-aws/s3_test.go
@@ -77,6 +77,7 @@ func init() {
 		}
 
 		parameters := DriverParameters{
+			nil,
 			accessKey,
 			secretKey,
 			bucket,