From 2a542b05fe7655102324caf0031cd0e0bd958266 Mon Sep 17 00:00:00 2001
From: jibin jose <jibinjose.in@gmail.com>
Date: Mon, 13 Apr 2020 16:44:33 +0530
Subject: [PATCH] Fix: Add X-Amz-Security-Token to presigned url (#4)

Include the x-amz-security-token if signed with temporary credentials.

Ref:https://stackoverflow.com/questions/49306621/how-secret-is-the-session-token-in-aws-temporary-security-credentials/49344291#49344291
---
 presigned.go      |  4 ++++
 presigned_test.go | 24 ++++++++++++++++++++++++
 simples3.go       | 11 +++++++++++
 3 files changed, 39 insertions(+)

diff --git a/presigned.go b/presigned.go
index ca5f3f2..511448a 100644
--- a/presigned.go
+++ b/presigned.go
@@ -85,6 +85,10 @@ func (s3 *S3) GeneratePresignedURL(in PresignedInput) string {
 		"X-Amz-Date":       amzdate,
 		"X-Amz-Expires":    strconv.Itoa(in.ExpirySeconds),
 	}
+	//  include the x-amz-security-token incase we are using IAM role or AWS STS
+	if s3.Token != "" {
+		queryString["X-Amz-Security-Token"] = s3.Token
+	}
 
 	// We need to have a sorted order,
 	// for QueryStrings and SignedHeaders
diff --git a/presigned_test.go b/presigned_test.go
index c03ff6b..1bd7c16 100644
--- a/presigned_test.go
+++ b/presigned_test.go
@@ -29,6 +29,30 @@ func TestS3_GeneratePresignedURL(t *testing.T) {
 	})
 }
 
+func TestS3_GeneratePresignedURL_Token(t *testing.T) {
+	// Params based on
+	// https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-query-string-auth.html
+	var time, _ = time.Parse(time.RFC1123, "Fri, 24 May 2013 00:00:00 GMT")
+	t.Run("Test", func(t *testing.T) {
+		s := New(
+			"us-east-1",
+			"AKIAIOSFODNN7EXAMPLE",
+			"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+		)
+		s.SetToken("IQoJb3JpT2luX2VjEPP%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEaCmFwLXNvdXRoLTEiRzBFAiABaeeW0LZZaqVyQVx8EHfCY9KTLsR0hnw1nDae%2F%2BVDbwIhAKrGP4RYkoPv8x0qFScsp%2FQZZXAYWbspMOMpVEBa1%2FQ3Kr8DCPz%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEQARoMOTMyNjk0MjUxNzI3IgxHyURIpz%2FBVH7V0ikqkwMTy9uf3umf7OWghmeDE8fpS7KxXYlTCQdVyC6tHcTQZdZ13qziy0ZgImvJEUz4lFNCszdQWR2jaDjgNGvWEUJ1ODAir7F1gTb%2BSx0PpH8o18yrrTJYCwZe7ZKtViCN2yDKHAk8DN9Ke77fYEl2W%2FLWV3VH9oqwEwUzCh4f6JrluiLW6HaxHcDqu7K6Qk8bhgTVlW5eHBzlyRJtrlmy232auL1m8XAoR01sjnpoCwE0ra1L3QuK7XmC9BIR5bRwMdZFcL0Ai0vzCyX9kd15hhDBRgzKrTNSrBFDaRJ9N%2FV3bZ61RAd%2FkwfQEDBiwUcTdm%2BVDLvxIUfVNmtQj628ZCWi%2BztUAe8Yz8IKpY50nEXr%2BHHX4wtVF2MZQPSOr%2B%2FON3OJYCl6TwVTGWoVGapn9y%2Bj9JOcdnnDuFUJMoJERRWnMNPCadZT68%2B3t30IgmXU4hcSX51olExLeGMSMtfK6LC7YCvMlGG8YxIJAeW5qznc2d9u%2BX7nXjqhvPCyc9hXMv4hXS4rowWnR6gaz6xZuY9fb8TMIK4v%2FQFOusBpv3m9H7b45zUr3o6xYh28GyB5%2F9zW%2FPkfm%2FpysDbwfz3r3G0WLchyE0t4%2BH8YZibj0KwY8rJyAV26u2DzIlp0bmJ%2F7Aaq4wUo%2BgUbhz7NMFUpWuR2ywszf28pdgsRQ4SHAlVQ4rOhx5XGqMREzjFPJo7jRW6uMCSJ8LvrQU38VTpZyrm7yQDCBK2lHwU00O8xTWSDhFXmrqFrCL9P76ZYXh2dCCJm6gPiSU3eGyqGBKDBWFt20lRHLWCyXwiyhGRULg3WLoLDVsjJDRO8xZta8nVxALUZLcteEv%2BE1QGCxVSg1W1WSAGLz8FQ%3D%3D")
+		want := "https://examplebucket.s3.amazonaws.com/test.txt?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIOSFODNN7EXAMPLE%2F20130524%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20130524T000000Z&X-Amz-Expires=86400&X-Amz-Security-Token=IQoJb3JpT2luX2VjEPP%252F%252F%252F%252F%252F%252F%252F%252F%252F%252FwEaCmFwLXNvdXRoLTEiRzBFAiABaeeW0LZZaqVyQVx8EHfCY9KTLsR0hnw1nDae%252F%252BVDbwIhAKrGP4RYkoPv8x0qFScsp%252FQZZXAYWbspMOMpVEBa1%252FQ3Kr8DCPz%252F%252F%252F%252F%252F%252F%252F%252F%252F%252FwEQARoMOTMyNjk0MjUxNzI3IgxHyURIpz%252FBVH7V0ikqkwMTy9uf3umf7OWghmeDE8fpS7KxXYlTCQdVyC6tHcTQZdZ13qziy0ZgImvJEUz4lFNCszdQWR2jaDjgNGvWEUJ1ODAir7F1gTb%252BSx0PpH8o18yrrTJYCwZe7ZKtViCN2yDKHAk8DN9Ke77fYEl2W%252FLWV3VH9oqwEwUzCh4f6JrluiLW6HaxHcDqu7K6Qk8bhgTVlW5eHBzlyRJtrlmy232auL1m8XAoR01sjnpoCwE0ra1L3QuK7XmC9BIR5bRwMdZFcL0Ai0vzCyX9kd15hhDBRgzKrTNSrBFDaRJ9N%252FV3bZ61RAd%252FkwfQEDBiwUcTdm%252BVDLvxIUfVNmtQj628ZCWi%252BztUAe8Yz8IKpY50nEXr%252BHHX4wtVF2MZQPSOr%252B%252FON3OJYCl6TwVTGWoVGapn9y%252Bj9JOcdnnDuFUJMoJERRWnMNPCadZT68%252B3t30IgmXU4hcSX51olExLeGMSMtfK6LC7YCvMlGG8YxIJAeW5qznc2d9u%252BX7nXjqhvPCyc9hXMv4hXS4rowWnR6gaz6xZuY9fb8TMIK4v%252FQFOusBpv3m9H7b45zUr3o6xYh28GyB5%252F9zW%252FPkfm%252FpysDbwfz3r3G0WLchyE0t4%252BH8YZibj0KwY8rJyAV26u2DzIlp0bmJ%252F7Aaq4wUo%252BgUbhz7NMFUpWuR2ywszf28pdgsRQ4SHAlVQ4rOhx5XGqMREzjFPJo7jRW6uMCSJ8LvrQU38VTpZyrm7yQDCBK2lHwU00O8xTWSDhFXmrqFrCL9P76ZYXh2dCCJm6gPiSU3eGyqGBKDBWFt20lRHLWCyXwiyhGRULg3WLoLDVsjJDRO8xZta8nVxALUZLcteEv%252BE1QGCxVSg1W1WSAGLz8FQ%253D%253D&X-Amz-SignedHeaders=host&X-Amz-Signature=29d003f449ae4106d1c4cabaeebf84fc47960ee127e98f1b9132261852250cb4"
+		if got := s.GeneratePresignedURL(PresignedInput{
+			Bucket:        "examplebucket",
+			ObjectKey:     "test.txt",
+			Method:        "GET",
+			Timestamp:     time,
+			ExpirySeconds: 86400,
+		}); got != want {
+			t.Errorf("S3.GeneratePresignedURL() = %v, want %v", got, want)
+		}
+	})
+}
+
 func TestS3_GeneratePresignedURL_Personal(t *testing.T) {
 	t.Run("Test", func(t *testing.T) {
 		s := New(
diff --git a/simples3.go b/simples3.go
index eed5eec..3e506e8 100644
--- a/simples3.go
+++ b/simples3.go
@@ -30,6 +30,7 @@ type S3 struct {
 	Region    string
 	Client    *http.Client
 
+	Token     string
 	Endpoint  string
 	URIFormat string
 }
@@ -139,6 +140,7 @@ func newUsingIAMImpl(baseURL, region string) (*S3, error) {
 		Region:    region,
 		AccessKey: jsonResp.AccessKeyID,
 		SecretKey: jsonResp.SecretAccessKey,
+		Token:     jsonResp.Token,
 
 		URIFormat: "https://s3.%s.amazonaws.com/%s",
 	}, nil
@@ -173,6 +175,15 @@ func (s3 *S3) SetEndpoint(uri string) *S3 {
 	return s3
 }
 
+// SetToken can be used to set a Temporary Security Credential token obtained from
+// using an IAM role or AWS STS.
+func (s3 *S3) SetToken(token string) *S3 {
+	if token != "" {
+		s3.Token = token
+	}
+	return s3
+}
+
 func detectFileSize(body io.Seeker) (int64, error) {
 	pos, err := body.Seek(0, 1)
 	if err != nil {