From 51c0c8148ab75fae72f60957b3b9f246e421b6f5 Mon Sep 17 00:00:00 2001
From: Kirat Singh <kirat.singh@wsq.io>
Date: Fri, 20 Jan 2017 12:12:37 -0500
Subject: [PATCH] Add new parameter s3accelerate to S3 storage driver.

If s3accelerate is set to true then we turn on S3 Transfer
Acceleration via the AWS SDK.  It defaults to false since this is an
opt-in feature on the S3 bucket.

Signed-off-by: Kirat Singh <kirat.singh@wsq.io>
Signed-off-by: Simone Locci <simonelocci88@gmail.com>
---
 docs/configuration.md                     |  2 ++
 registry/storage/driver/s3-aws/s3.go      | 23 +++++++++++++++++++++++
 registry/storage/driver/s3-aws/s3_test.go | 11 +++++++++++
 3 files changed, 36 insertions(+)

diff --git a/docs/configuration.md b/docs/configuration.md
index 603181950..037d2912b 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -118,6 +118,7 @@ storage:
     secretkey: awssecretkey
     region: us-west-1
     regionendpoint: http://myobjects.local
+    s3accelerate: false
     bucket: bucketname
     encrypt: true
     keyid: mykeyid
@@ -422,6 +423,7 @@ storage:
     secretkey: awssecretkey
     region: us-west-1
     regionendpoint: http://myobjects.local
+    s3accelerate: false
     bucket: bucketname
     encrypt: true
     keyid: mykeyid
diff --git a/registry/storage/driver/s3-aws/s3.go b/registry/storage/driver/s3-aws/s3.go
index 238c5d80a..3b1c9eb53 100644
--- a/registry/storage/driver/s3-aws/s3.go
+++ b/registry/storage/driver/s3-aws/s3.go
@@ -104,6 +104,7 @@ type DriverParameters struct {
 	ObjectACL                   string
 	SessionToken                string
 	UseDualStack                bool
+	S3Accelerate                bool
 }
 
 func init() {
@@ -377,6 +378,23 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 
 	sessionToken := ""
 
+	s3accelerateBool := false
+	s3accelerate := parameters["s3accelerate"]
+	switch s3accelerate := s3accelerate.(type) {
+	case string:
+		b, err := strconv.ParseBool(s3accelerate)
+		if err != nil {
+			return nil, fmt.Errorf("The s3accelerate parameter should be a boolean")
+		}
+		s3accelerateBool = b
+	case bool:
+		s3accelerateBool = s3accelerate
+	case nil:
+		// do nothing
+	default:
+		return nil, fmt.Errorf("The s3accelerate parameter should be a boolean")
+	}
+
 	params := DriverParameters{
 		fmt.Sprint(accessKey),
 		fmt.Sprint(secretKey),
@@ -399,6 +417,7 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 		objectACL,
 		fmt.Sprint(sessionToken),
 		useDualStackBool,
+		s3accelerateBool,
 	}
 
 	return New(params)
@@ -458,6 +477,10 @@ func New(params DriverParameters) (*Driver, error) {
 		awsConfig.WithEndpoint(params.RegionEndpoint)
 	}
 
+	if params.S3Accelerate {
+		awsConfig.WithS3UseAccelerate(true)
+	}
+
 	awsConfig.WithRegion(params.Region)
 	awsConfig.WithDisableSSL(!params.Secure)
 	if params.UseDualStack {
diff --git a/registry/storage/driver/s3-aws/s3_test.go b/registry/storage/driver/s3-aws/s3_test.go
index 31662b175..0314eaaf6 100644
--- a/registry/storage/driver/s3-aws/s3_test.go
+++ b/registry/storage/driver/s3-aws/s3_test.go
@@ -44,6 +44,7 @@ func init() {
 	sessionToken := os.Getenv("AWS_SESSION_TOKEN")
 	useDualStack := os.Getenv("S3_USE_DUALSTACK")
 	combineSmallPart := os.Getenv("MULTIPART_COMBINE_SMALL_PART")
+	s3accelerate := os.Getenv("S3_ACCELERATE")
 	if err != nil {
 		panic(err)
 	}
@@ -86,6 +87,7 @@ func init() {
 		if useDualStack != "" {
 			useDualStackBool, err = strconv.ParseBool(useDualStack)
 		}
+
 		multipartCombineSmallPart := true
 		if combineSmallPart != "" {
 			multipartCombineSmallPart, err = strconv.ParseBool(combineSmallPart)
@@ -94,6 +96,14 @@ func init() {
 			}
 		}
 
+		s3accelerateBool := true
+		if s3accelerate != "" {
+			s3accelerateBool, err = strconv.ParseBool(s3accelerate)
+			if err != nil {
+				return nil, err
+			}
+		}
+
 		parameters := DriverParameters{
 			accessKey,
 			secretKey,
@@ -116,6 +126,7 @@ func init() {
 			objectACL,
 			sessionToken,
 			useDualStackBool,
+			s3accelerateBool,
 		}
 
 		return New(parameters)