Fix base64 error when input is multiple of 4 bytes

itsdangerous.go

@@ -9,8 +9,6 @@ package itsdangerous
 
 import (
 	"encoding/base64"
-	"fmt"
-	"strings"
 	"time"
 )
 
@@ -19,26 +17,12 @@ const EPOCH = 1293840000
 
 // Encodes a single string. The resulting string is safe for putting into URLs.
 func base64Encode(src []byte) string {
-	s := base64.URLEncoding.EncodeToString(src)
-	return strings.Trim(s, "=")
+	return base64.RawURLEncoding.EncodeToString(src)
 }
 
 // Decodes a single string.
 func base64Decode(s string) ([]byte, error) {
-	var padLen int
-
-	if l := len(s) % 4; l > 0 {
-		padLen = 4 - l
-	} else {
-		padLen = 1
-	}
-
-	b, err := base64.URLEncoding.DecodeString(s + strings.Repeat("=", padLen))
-	if err != nil {
-		fmt.Println(s)
-		return []byte(""), err
-	}
-	return b, nil
+	return base64.RawURLEncoding.DecodeString(s)
 }
 
 // Returns the current timestamp.  This implementation returns the

itsdangerous_test.go

@@ -0,0 +1,37 @@
+package itsdangerous
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestBase64(t *testing.T) {
+
+	tests := []struct {
+		value   []byte
+		encoded string
+	}{
+		{value: []byte("a"), encoded: "YQ"},
+		{value: []byte("ab"), encoded: "YWI"},
+		{value: []byte("abc"), encoded: "YWJj"},
+		{value: []byte("abcd"), encoded: "YWJjZA"},
+		{value: []byte("abcde"), encoded: "YWJjZGU"},
+		{value: []byte("abcdef"), encoded: "YWJjZGVm"},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(string(test.value), func(t *testing.T) {
+			actualEncoded := base64Encode(test.value)
+			if actualEncoded != test.encoded {
+				t.Errorf("base64Encode(%v) got %s; want %s", test.value, actualEncoded, test.encoded)
+			}
+
+			decoded, err := base64Decode(test.encoded)
+			if err != nil {
+				t.Errorf("base64Decode(%s) returned error: %s", test.encoded, err)
+			} else if !reflect.DeepEqual(decoded, test.value) {
+				t.Errorf("base64Decode(%s) got %v; want %v", test.encoded, decoded, test.value)
+			}
+		})
+	}
+}