Fixes requested by codyoss@

google/internal/externalaccount/aws.go

@@ -1,13 +1,16 @@
 // Copyright 2021 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
+
 package externalaccount
 
 import (
 	"crypto/hmac"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"path"
@@ -16,13 +19,13 @@ import (
 	"time"
 )
 
-// A utility class to sign http requests using a AWS V4 signature
+// RequestSigner is a utility class to sign http requests using a AWS V4 signature.
 type RequestSigner struct {
 	RegionName             string
 	AwsSecurityCredentials map[string]string
-	debugTimestamp         time.Time
 }
 
+// NewRequestSigner is a method to create a RequestSigner with the appropriately filled out fields.
 func NewRequestSigner(regionName string, awsSecurityCredentials map[string]string) *RequestSigner {
 	return &RequestSigner{
 		RegionName:             regionName,
@@ -30,32 +33,60 @@ func NewRequestSigner(regionName string, awsSecurityCredentials map[string]strin
 	}
 }
 
+const (
 // AWS Signature Version 4 signing algorithm identifier.
-const awsAlgorithm = "AWS4-HMAC-SHA256"
+	awsAlgorithm = "AWS4-HMAC-SHA256"
 
 // The termination string for the AWS credential scope value as defined in
 // https://docs.aws.amazon.com/general/latest/gr/sigv4-create-string-to-sign.html
-const awsRequestType = "aws4_request"
+	awsRequestType = "aws4_request"
 
 // The AWS authorization header name for the security session token if available.
-const awsSecurityTokenHeader = "x-amz-security-token"
+	awsSecurityTokenHeader = "x-amz-security-token"
 
 // The AWS authorization header name for the auto-generated date.
-const awsDateHeader = "x-amz-date"
+	awsDateHeader = "x-amz-date"
 
-const awsTimeFormatLong = "20060102T150405Z"
+	awsTimeFormatLong = "20060102T150405Z"
-const awsTimeFormatShort = "20060102"
+	awsTimeFormatShort = "20060102"
+)
 
-func getSha256(input []byte) string {
+func getSha256(input []byte) (string, error) {
 	hash := sha256.New()
-	hash.Write(input)
+	if _, err := hash.Write(input); err != nil {
-	return hex.EncodeToString(hash.Sum(nil))
+		return "", err
+	}
+	return hex.EncodeToString(hash.Sum(nil)), nil
 }
 
-func getHmacSha256(key, input []byte) []byte {
+func getHmacSha256(key, input []byte) ([]byte, error) {
 	hash := hmac.New(sha256.New, key)
-	hash.Write(input)
+	if _, err := hash.Write(input); err != nil {
-	return hash.Sum(nil)
+		return nil, err
+	}
+	return hash.Sum(nil), nil
+}
+
+func cloneRequest(r *http.Request) *http.Request {
+	r2 := new(http.Request)
+	*r2 = *r
+	if r.Header != nil {
+		r2.Header = make(http.Header, len(r.Header))
+
+		// Find total number of values.
+		headerCount := 0
+		for _, headerValues := range r.Header {
+			headerCount += len(headerValues)
+		}
+		copiedHeaders := make([]string, headerCount) // shared backing array for headers' values
+
+		for headerKey, headerValues := range r.Header {
+			headerCount = copy(copiedHeaders, headerValues)
+			r2.Header[headerKey] = copiedHeaders[:headerCount:headerCount]
+			copiedHeaders = copiedHeaders[headerCount:]
+		}
+	}
+	return r2
 }
 
 func canonicalPath(req *http.Request) string {
@@ -90,20 +121,31 @@ func canonicalHeaders(req *http.Request) (string, string) {
 	}
 	sort.Strings(headers)
 
-	var fullHeaders []string
+	var fullHeaders strings.Builder
 	for _, header := range headers {
 		headerValue := strings.Join(lowerCaseHeaders[header], ",")
-		fullHeaders = append(fullHeaders, header+":"+headerValue+"\n")
+		fullHeaders.WriteString(header)
+		fullHeaders.WriteRune(':')
+		fullHeaders.WriteString(headerValue)
+		fullHeaders.WriteRune('\n')
 	}
 
-	return strings.Join(headers, ";"), strings.Join(fullHeaders, "")
+	return strings.Join(headers, ";"), fullHeaders.String()
 }
 
-func requestDataHash(req *http.Request) string {
+func requestDataHash(req *http.Request) (string, error) {
-	requestData := []byte{}
+	var requestData []byte
 	if req.Body != nil {
-		requestBody, _ := req.GetBody()
+		requestBody, err := req.GetBody()
-		requestData, _ = ioutil.ReadAll(requestBody)
+		if err != nil {
+			return "", err
+		}
+		defer requestBody.Close()
+
+		requestData, err = ioutil.ReadAll(io.LimitReader(requestBody, 1<<20))
+		if err != nil {
+			return "", err
+		}
 	}
 
 	return getSha256(requestData)
@@ -116,63 +158,93 @@ func requestHost(req *http.Request) string {
 	return req.URL.Host
 }
 
-func canonicalRequest(req *http.Request, canonicalHeaderColumns, canonicalHeaderData string) string {
+func canonicalRequest(req *http.Request, canonicalHeaderColumns, canonicalHeaderData string) (string, error) {
+	dataHash, err := requestDataHash(req)
+	if err != nil {
+		return "", err
+	}
+
 	return strings.Join([]string{
 		req.Method,
 		canonicalPath(req),
 		canonicalQuery(req),
 		canonicalHeaderData,
 		canonicalHeaderColumns,
-		requestDataHash(req),
+		dataHash,
-	}, "\n")
+	}, "\n"), nil
 }
 
-func (rs *RequestSigner) SignedRequest(req *http.Request) *http.Request {
+// SignRequest adds the appropriate headers to an http.Request
-	timestamp := rs.debugTimestamp
+// or returns an error if something prevented this.
-	if timestamp.IsZero() {
+func (rs *RequestSigner) SignRequest(req *http.Request) error {
-		timestamp = time.Now()
+	signedRequest := cloneRequest(req)
-	}
+	timestamp := now()
-	signedRequest := req.Clone(req.Context())
 
 	signedRequest.Header.Add("host", requestHost(req))
 
-	securityToken, ok := rs.AwsSecurityCredentials["security_token"]
+	if securityToken, ok := rs.AwsSecurityCredentials["security_token"]; ok {
-	if ok {
+		signedRequest.Header.Add(awsSecurityTokenHeader, securityToken)
-		signedRequest.Header.Add("x-amz-security-token", securityToken)
 	}
 
 	if signedRequest.Header.Get("date") == "" {
-		signedRequest.Header.Add("x-amz-date", timestamp.Format(awsTimeFormatLong))
+		signedRequest.Header.Add(awsDateHeader, timestamp.Format(awsTimeFormatLong))
 	}
 
-	signedRequest.Header.Set("Authorization", rs.generateAuthentication(signedRequest, timestamp))
+	authorizationCode, err := rs.generateAuthentication(signedRequest, timestamp)
+	if err != nil {
+		return err
+	}
+	signedRequest.Header.Set("Authorization", authorizationCode)
 
-	return signedRequest
+	req.Header = signedRequest.Header
+	return nil
+}
+
+func (rs *RequestSigner) generateAuthentication(req *http.Request, timestamp time.Time) (string, error) {
+	secretAccessKey, ok := rs.AwsSecurityCredentials["secret_access_key"]
+	if !ok {
+		return "", errors.New("Missing Secret Access Key")
+	}
+	accessKeyId, ok := rs.AwsSecurityCredentials["access_key_id"]
+	if !ok {
+		return "", errors.New("Missing Access Key Id")
 	}
 
-func (rs *RequestSigner) generateAuthentication(req *http.Request, timestamp time.Time) string {
 	canonicalHeaderColumns, canonicalHeaderData := canonicalHeaders(req)
 
 	dateStamp := timestamp.Format(awsTimeFormatShort)
-	serviceName := strings.Split(requestHost(req), ".")[0]
+	serviceName := ""
+	if splitHost := strings.Split(requestHost(req), "."); len(splitHost) > 0 {
+		serviceName = splitHost[0]
+	}
 
-	credentialScope := strings.Join([]string{
+	credentialScope := fmt.Sprintf("%s/%s/%s/%s",dateStamp, rs.RegionName, serviceName, awsRequestType)
-		dateStamp, rs.RegionName, serviceName, awsRequestType,
+
-	}, "/")
+	requestString, err := canonicalRequest(req, canonicalHeaderColumns, canonicalHeaderData)
+	if err != nil {
+		return "", err
+	}
+	requestHash, err := getSha256([]byte(requestString))
+	if err != nil{
+		return "", err
+	}
 
 	stringToSign := strings.Join([]string{
 		awsAlgorithm,
 		timestamp.Format(awsTimeFormatLong),
 		credentialScope,
-		getSha256([]byte(canonicalRequest(req, canonicalHeaderColumns, canonicalHeaderData))),
+		requestHash,
 	}, "\n")
 
-	signingKey := []byte("AWS4" + rs.AwsSecurityCredentials["secret_access_key"])
+	signingKey := []byte("AWS4" + secretAccessKey)
 	for _, signingInput := range []string{
 		dateStamp, rs.RegionName, serviceName, awsRequestType, stringToSign,
 	} {
-		signingKey = getHmacSha256(signingKey, []byte(signingInput))
+		signingKey, err = getHmacSha256(signingKey, []byte(signingInput))
+		if err != nil{
+			return "", err
+		}
 	}
 
-	return fmt.Sprintf("%s Credential=%s/%s, SignedHeaders=%s, Signature=%s", awsAlgorithm, rs.AwsSecurityCredentials["access_key_id"], credentialScope, canonicalHeaderColumns, hex.EncodeToString(signingKey))
+	return fmt.Sprintf("%s Credential=%s/%s, SignedHeaders=%s, Signature=%s", awsAlgorithm, accessKeyId, credentialScope, canonicalHeaderColumns, hex.EncodeToString(signingKey)), nil
 }

google/internal/externalaccount/aws_test.go

@@ -14,6 +14,12 @@ import (
 var defaultTime = time.Date(2011, 9, 9, 23, 36, 0, 0, time.UTC)
 var secondDefaultTime = time.Date(2020, 8, 11, 6, 55, 22, 0, time.UTC)
 
+func setTime(testTime time.Time) func() time.Time {
+	return func() time.Time {
+		return testTime
+	}
+}
+
 var defaultRequestSigner = NewRequestSigner("us-east-1", map[string]string{
 	"access_key_id":     "AKIDEXAMPLE",
 	"secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
@@ -36,16 +42,19 @@ func setDefaultTime(req *http.Request) {
 }
 
 func testRequestSigner(t *testing.T, rs *RequestSigner, input, expectedOutput *http.Request) {
-	actualOutput := rs.SignedRequest(input)
+	err := rs.SignRequest(input)
+	if err != nil {
+		t.Errorf("unexpected error: %q", err.Error())
+	}
 
-	if got, want := actualOutput.URL.String(), expectedOutput.URL.String(); !reflect.DeepEqual(got, want) {
+	if got, want := input.URL.String(), expectedOutput.URL.String(); !reflect.DeepEqual(got, want) {
 		t.Errorf("url = %q, want %q", got, want)
 	}
-	if got, want := actualOutput.Method, expectedOutput.Method; !reflect.DeepEqual(got, want) {
+	if got, want := input.Method, expectedOutput.Method; !reflect.DeepEqual(got, want) {
 		t.Errorf("method = %q, want %q", got, want)
 	}
 	for header := range expectedOutput.Header {
-		if got, want := actualOutput.Header[header], expectedOutput.Header[header]; !reflect.DeepEqual(got, want) {
+		if got, want := input.Header[header], expectedOutput.Header[header]; !reflect.DeepEqual(got, want) {
 			t.Errorf("header[%q] = %q, want %q", header, got, want)
 		}
 	}
@@ -62,6 +71,10 @@ func TestAwsV4Signature_GetRequest(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=b27ccfbfa7df52a200ff74193ca6e32d4b48b8856fab7ebf1c595d0670a7e470"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -76,6 +89,10 @@ func TestAwsV4Signature_GetRequestWithRelativePath(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=b27ccfbfa7df52a200ff74193ca6e32d4b48b8856fab7ebf1c595d0670a7e470"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -90,6 +107,10 @@ func TestAwsV4Signature_GetRequestWithDotPath(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=b27ccfbfa7df52a200ff74193ca6e32d4b48b8856fab7ebf1c595d0670a7e470"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -104,6 +125,10 @@ func TestAwsV4Signature_GetRequestWithPointlessDotPath(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=910e4d6c9abafaf87898e1eb4c929135782ea25bb0279703146455745391e63a"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -118,6 +143,10 @@ func TestAwsV4Signature_GetRequestWithUtf8Path(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=8d6634c189aa8c75c2e51e106b6b5121bed103fdb351f7d7d4381c738823af74"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -132,6 +161,10 @@ func TestAwsV4Signature_GetRequestWithDuplicateQuery(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=be7148d34ebccdc6423b19085378aa0bee970bdc61d144bd1a8c48c33079ab09"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -146,6 +179,10 @@ func TestAwsV4Signature_GetRequestWithMisorderedQuery(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=feb926e49e382bec75c9d7dcb2a1b6dc8aa50ca43c25d2bc51143768c0875acc"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -160,6 +197,10 @@ func TestAwsV4Signature_GetRequestWithUtf8Query(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=6fb359e9a05394cc7074e0feb42573a2601abc0c869a953e8c5c12e4e01f1a8c"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -176,6 +217,10 @@ func TestAwsV4Signature_PostRequest(t *testing.T) {
 		"Zoo":           []string{"zoobar"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -192,6 +237,10 @@ func TestAwsV4Signature_PostRequestWithCapitalizedHeaderValue(t *testing.T) {
 		"Zoo":           []string{"ZOOBAR"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -208,6 +257,10 @@ func TestAwsV4Signature_PostRequestPhfft(t *testing.T) {
 		"P":             []string{"phfft"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -224,6 +277,10 @@ func TestAwsV4Signature_PostRequestWithBody(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=content-type;date;host, Signature=5a15b22cf462f047318703b92e6f4f38884e4a7ab7b1d6426ca46a8bd1c26cbc"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -238,6 +295,10 @@ func TestAwsV4Signature_PostRequestWithQueryString(t *testing.T) {
 		"Authorization": []string{"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20110909/us-east-1/host/aws4_request, SignedHeaders=date;host, Signature=b6e3b79003ce0743a491606ba1035a804593b0efb1e20a11cba83f8c25a57a92"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(defaultTime)
+
 	testRequestSigner(t, defaultRequestSigner, input, output)
 }
 
@@ -252,6 +313,10 @@ func TestAwsV4Signature_GetRequestWithSecurityToken(t *testing.T) {
 		"X-Amz-Security-Token": []string{securityToken},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(secondDefaultTime)
+
 	testRequestSigner(t, requestSignerWithToken, input, output)
 }
 
@@ -266,6 +331,10 @@ func TestAwsV4Signature_PostRequestWithSecurityToken(t *testing.T) {
 		"X-Amz-Security-Token": []string{securityToken},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(secondDefaultTime)
+
 	testRequestSigner(t, requestSignerWithToken, input, output)
 }
 
@@ -285,6 +354,10 @@ func TestAwsV4Signature_PostRequestWithSecurityTokenAndAdditionalHeaders(t *test
 		"X-Amz-Security-Token": []string{securityToken},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(secondDefaultTime)
+
 	testRequestSigner(t, requestSignerWithToken, input, output)
 }
 
@@ -293,7 +366,6 @@ func TestAwsV4Signature_PostRequestWithAmzDateButNoSecurityToken(t *testing.T) {
 		"access_key_id":     accessKeyId,
 		"secret_access_key": secretAccessKey,
 	})
-	requestSigner.debugTimestamp = secondDefaultTime
 
 	input, _ := http.NewRequest("POST", "https://sts.us-east-2.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15", nil)
 
@@ -304,10 +376,9 @@ func TestAwsV4Signature_PostRequestWithAmzDateButNoSecurityToken(t *testing.T) {
 		"X-Amz-Date":    []string{"20200811T065522Z"},
 	}
 
+	oldNow := now
+	defer func() { now = oldNow }()
+	now = setTime(secondDefaultTime)
+
 	testRequestSigner(t, requestSigner, input, output)
 }
-
-func init() {
-	defaultRequestSigner.debugTimestamp = defaultTime
-	requestSignerWithToken.debugTimestamp = secondDefaultTime
-}