From 7e81a9365c6fe9f1a4691742792dcd2e3d8a1fa6 Mon Sep 17 00:00:00 2001
From: Will McCutchen <will@mccutch.org>
Date: Fri, 5 May 2023 17:39:53 -0400
Subject: [PATCH] fix: correctly handle padded base64 input (#119)

Also, switch to URL-safe base64 instead of standard base64, for
compatibility with original httpbin.

Fixes #118.
---
 httpbin/handlers_test.go | 24 ++++++++++++++++++++++++
 httpbin/helpers.go       | 10 ++++------
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/httpbin/handlers_test.go b/httpbin/handlers_test.go
index 26d4ea9c..1fd1c2b4 100644
--- a/httpbin/handlers_test.go
+++ b/httpbin/handlers_test.go
@@ -2828,6 +2828,24 @@ func TestBase64(t *testing.T) {
 			"/base64/encode/valid_base64_encoded_string",
 			"dmFsaWRfYmFzZTY0X2VuY29kZWRfc3RyaW5n",
 		},
+		{
+			// make sure we correctly handle padding
+			// https://github.com/mccutchen/go-httpbin/issues/118
+			"/base64/dGVzdC1pbWFnZQ==",
+			"test-image",
+		},
+		{
+			// URL-safe base64 is used for decoding (note the - instead of + in
+			// encoded input string)
+			"/base64/decode/YWJjMTIzIT8kKiYoKSctPUB-",
+			"abc123!?$*&()'-=@~",
+		},
+		{
+			// URL-safe base64 is used for encoding (note the - instead of + in
+			// encoded output string)
+			"/base64/encode/abc123%21%3F%24%2A%26%28%29%27-%3D%40~",
+			"YWJjMTIzIT8kKiYoKSctPUB-",
+		},
 	}
 
 	for _, test := range okTests {
@@ -2883,6 +2901,12 @@ func TestBase64(t *testing.T) {
 			"/base64/unknown/dmFsaWRfYmFzZTY0X2VuY29kZWRfc3RyaW5n",
 			"invalid operation: unknown",
 		},
+		{
+			// we only support URL-safe base64 encoded strings (note the +
+			// instead of - in encoded input string)
+			"/base64/decode/YWJjMTIzIT8kKiYoKSctPUB+",
+			"illegal base64 data",
+		},
 	}
 
 	for _, test := range errorTests {
diff --git a/httpbin/helpers.go b/httpbin/helpers.go
index b99f76ed..2e75e742 100644
--- a/httpbin/helpers.go
+++ b/httpbin/helpers.go
@@ -184,7 +184,7 @@ func parseBody(w http.ResponseWriter, r *http.Request, resp *bodyResponse) error
 
 // return provided string as base64 encoded data url, with the given content type
 func encodeData(body []byte, contentType string) string {
-	data := base64.StdEncoding.EncodeToString(body)
+	data := base64.URLEncoding.EncodeToString(body)
 
 	// If no content type is provided, default to application/octet-stream
 	if contentType == "" {
@@ -370,14 +370,12 @@ func newBase64Helper(path string) (*base64Helper, error) {
 
 // Encode - encode data as base64
 func (b *base64Helper) Encode() ([]byte, error) {
-	buff := make([]byte, base64.StdEncoding.EncodedLen(len(b.data)))
-	base64.StdEncoding.Encode(buff, []byte(b.data))
+	buff := make([]byte, base64.URLEncoding.EncodedLen(len(b.data)))
+	base64.URLEncoding.Encode(buff, []byte(b.data))
 	return buff, nil
 }
 
 // Decode - decode data from base64
 func (b *base64Helper) Decode() ([]byte, error) {
-	buff := make([]byte, base64.StdEncoding.DecodedLen(len(b.data)))
-	_, err := base64.StdEncoding.Decode(buff, []byte(b.data))
-	return buff, err
+	return base64.URLEncoding.DecodeString(b.data)
 }