From 7f510ae9c9d15b42c621b6ccb6ef90a6ca94c12b Mon Sep 17 00:00:00 2001
From: Troels Thomsen <troels@thomsen.io>
Date: Mon, 27 Mar 2017 22:04:00 +0200
Subject: [PATCH] Support session token

Signed-off-by: Troels Thomsen <troels@thomsen.io>
---
 registry/storage/driver/s3-aws/s3.go           | 5 +++++
 registry/storage/driver/s3-aws/s3_test.go      | 2 ++
 registry/storage/driver/s3-aws/s3_v2_signer.go | 3 +++
 3 files changed, 10 insertions(+)

diff --git a/registry/storage/driver/s3-aws/s3.go b/registry/storage/driver/s3-aws/s3.go
index 469e3997..19407d80 100644
--- a/registry/storage/driver/s3-aws/s3.go
+++ b/registry/storage/driver/s3-aws/s3.go
@@ -98,6 +98,7 @@ type DriverParameters struct {
 	StorageClass                string
 	UserAgent                   string
 	ObjectACL                   string
+	SessionToken                string
 }
 
 func init() {
@@ -331,6 +332,8 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 		objectACL = objectACLString
 	}
 
+	sessionToken := ""
+
 	params := DriverParameters{
 		fmt.Sprint(accessKey),
 		fmt.Sprint(secretKey),
@@ -349,6 +352,7 @@ func FromParameters(parameters map[string]interface{}) (*Driver, error) {
 		storageClass,
 		fmt.Sprint(userAgent),
 		objectACL,
+		fmt.Sprint(sessionToken),
 	}
 
 	return New(params)
@@ -398,6 +402,7 @@ func New(params DriverParameters) (*Driver, error) {
 			Value: credentials.Value{
 				AccessKeyID:     params.AccessKey,
 				SecretAccessKey: params.SecretKey,
+				SessionToken:    params.SessionToken,
 			},
 		},
 		&credentials.EnvProvider{},
diff --git a/registry/storage/driver/s3-aws/s3_test.go b/registry/storage/driver/s3-aws/s3_test.go
index eb7ee519..363a22eb 100644
--- a/registry/storage/driver/s3-aws/s3_test.go
+++ b/registry/storage/driver/s3-aws/s3_test.go
@@ -36,6 +36,7 @@ func init() {
 	objectACL := os.Getenv("S3_OBJECT_ACL")
 	root, err := ioutil.TempDir("", "driver-")
 	regionEndpoint := os.Getenv("REGION_ENDPOINT")
+	sessionToken := os.Getenv("AWS_SESSION_TOKEN")
 	if err != nil {
 		panic(err)
 	}
@@ -84,6 +85,7 @@ func init() {
 			storageClass,
 			driverName + "-test",
 			objectACL,
+			sessionToken,
 		}
 
 		return New(parameters)
diff --git a/registry/storage/driver/s3-aws/s3_v2_signer.go b/registry/storage/driver/s3-aws/s3_v2_signer.go
index 7cabe07e..cb801087 100644
--- a/registry/storage/driver/s3-aws/s3_v2_signer.go
+++ b/registry/storage/driver/s3-aws/s3_v2_signer.go
@@ -137,6 +137,9 @@ func (v2 *signer) Sign() error {
 	host, canonicalPath := parsedURL.Host, parsedURL.Path
 	v2.Request.Header["Host"] = []string{host}
 	v2.Request.Header["date"] = []string{v2.Time.In(time.UTC).Format(time.RFC1123)}
+	if credValue.SessionToken != "" {
+		v2.Request.Header["x-amz-security-token"] = []string{credValue.SessionToken}
+	}
 
 	smap = make(map[string]string)
 	for k, v := range headers {