From b59e077a56ccc467b56b67045ea54a1c50903a3b Mon Sep 17 00:00:00 2001
From: ICHINOSE Shogo <shogo82148@gmail.com>
Date: Wed, 20 Dec 2023 23:52:55 +0900
Subject: [PATCH] simplify base64 encoding of Cloudfront signing

---
 service/cloudfront/sign/policy.go      | 28 ++++++--------------------
 service/cloudfront/sign/policy_test.go | 16 +++------------
 2 files changed, 9 insertions(+), 35 deletions(-)

diff --git a/service/cloudfront/sign/policy.go b/service/cloudfront/sign/policy.go
index 27c3ed318d4..f8664821e87 100644
--- a/service/cloudfront/sign/policy.go
+++ b/service/cloudfront/sign/policy.go
@@ -16,6 +16,8 @@ import (
 	"unicode"
 )
 
+var customBase64Encoding = base64.NewEncoding("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-~").WithPadding('_')
+
 // An AWSEpochTime wraps a time value providing JSON serialization needed for
 // AWS Policy epoch time fields.
 type AWSEpochTime struct {
@@ -102,14 +104,12 @@ func (p *Policy) Sign(privKey *rsa.PrivateKey) (b64Signature, b64Policy []byte,
 	if err != nil {
 		return nil, nil, err
 	}
-	awsEscapeEncoded(b64Policy)
 
 	// Build and escape the signature
 	b64Signature, err = signEncodedPolicy(randReader, jsonPolicy, privKey)
 	if err != nil {
 		return nil, nil, err
 	}
-	awsEscapeEncoded(b64Signature)
 
 	return b64Signature, b64Policy, nil
 }
@@ -190,8 +190,8 @@ func encodePolicy(p *Policy) (b64Policy, jsonPolicy []byte, err error) {
 	// whitespace within the encoding.
 	jsonPolicy = bytes.TrimSpace(jsonPolicy)
 
-	b64Policy = make([]byte, base64.StdEncoding.EncodedLen(len(jsonPolicy)))
-	base64.StdEncoding.Encode(b64Policy, jsonPolicy)
+	b64Policy = make([]byte, customBase64Encoding.EncodedLen(len(jsonPolicy)))
+	customBase64Encoding.Encode(b64Policy, jsonPolicy)
 	return b64Policy, jsonPolicy, nil
 }
 
@@ -207,27 +207,11 @@ func signEncodedPolicy(randReader io.Reader, jsonPolicy []byte, privKey *rsa.Pri
 		return nil, fmt.Errorf("failed to sign policy, %s", err.Error())
 	}
 
-	b64Sig := make([]byte, base64.StdEncoding.EncodedLen(len(sig)))
-	base64.StdEncoding.Encode(b64Sig, sig)
+	b64Sig := make([]byte, customBase64Encoding.EncodedLen(len(sig)))
+	customBase64Encoding.Encode(b64Sig, sig)
 	return b64Sig, nil
 }
 
-// special characters to be replaced with awsEscapeEncoded
-var invalidEncodedChar = map[byte]byte{
-	'+': '-',
-	'=': '_',
-	'/': '~',
-}
-
-// awsEscapeEncoded will replace base64 encoding's special characters to be URL safe.
-func awsEscapeEncoded(b []byte) {
-	for i, v := range b {
-		if r, ok := invalidEncodedChar[v]; ok {
-			b[i] = r
-		}
-	}
-}
-
 func isASCII(u string) bool {
 	for _, c := range u {
 		if c > unicode.MaxASCII {
diff --git a/service/cloudfront/sign/policy_test.go b/service/cloudfront/sign/policy_test.go
index 2da5d8dc4c7..0fa0531bc2a 100644
--- a/service/cloudfront/sign/policy_test.go
+++ b/service/cloudfront/sign/policy_test.go
@@ -5,7 +5,6 @@ import (
 	"crypto"
 	"crypto/rsa"
 	"crypto/sha1"
-	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"math/rand"
@@ -103,7 +102,7 @@ var testTime = time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
 func TestEncodePolicy(t *testing.T) {
 	const (
 		expectedJSONPolicy = `{"Statement":[{"Resource":"https://example.com/a","Condition":{"DateLessThan":{"AWS:EpochTime":1257894000}}}]}`
-		expectedB64Policy  = `eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9leGFtcGxlLmNvbS9hIiwiQ29uZGl0aW9uIjp7IkRhdGVMZXNzVGhhbiI6eyJBV1M6RXBvY2hUaW1lIjoxMjU3ODk0MDAwfX19XX0=`
+		expectedB64Policy  = `eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9leGFtcGxlLmNvbS9hIiwiQ29uZGl0aW9uIjp7IkRhdGVMZXNzVGhhbiI6eyJBV1M6RXBvY2hUaW1lIjoxMjU3ODk0MDAwfX19XX0_`
 	)
 	p := NewCannedPolicy("https://example.com/a", testTime)
 
@@ -124,7 +123,7 @@ func TestEncodePolicy(t *testing.T) {
 func TestEncodePolicyWithQueryParams(t *testing.T) {
 	const (
 		expectedJSONPolicy = `{"Statement":[{"Resource":"https://example.com/a?b=1&c=2","Condition":{"DateLessThan":{"AWS:EpochTime":1257894000}}}]}`
-		expectedB64Policy  = `eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9leGFtcGxlLmNvbS9hP2I9MSZjPTIiLCJDb25kaXRpb24iOnsiRGF0ZUxlc3NUaGFuIjp7IkFXUzpFcG9jaFRpbWUiOjEyNTc4OTQwMDB9fX1dfQ==`
+		expectedB64Policy  = `eyJTdGF0ZW1lbnQiOlt7IlJlc291cmNlIjoiaHR0cHM6Ly9leGFtcGxlLmNvbS9hP2I9MSZjPTIiLCJDb25kaXRpb24iOnsiRGF0ZUxlc3NUaGFuIjp7IkFXUzpFcG9jaFRpbWUiOjEyNTc4OTQwMDB9fX1dfQ__`
 	)
 	p := NewCannedPolicy("https://example.com/a?b=1&c=2", testTime)
 
@@ -166,7 +165,7 @@ func TestSignEncodedPolicy(t *testing.T) {
 		t.Fatalf("Unexpected hash error, %#v", err)
 	}
 
-	decodedSig, err := base64.StdEncoding.DecodeString(string(b64Signature))
+	decodedSig, err := customBase64Encoding.DecodeString(string(b64Signature))
 	if err != nil {
 		t.Fatalf("Unexpected base64 decode signature, %#v", err)
 	}
@@ -175,12 +174,3 @@ func TestSignEncodedPolicy(t *testing.T) {
 		t.Fatalf("Unable to verify signature, %#v", err)
 	}
 }
-
-func TestAWSEscape(t *testing.T) {
-	expect := "a-b_c~"
-	actual := []byte("a+b=c/")
-	awsEscapeEncoded(actual)
-	if string(actual) != expect {
-		t.Errorf("expect: %s, actual: %s", expect, string(actual))
-	}
-}