From d7f2c4640b616f1e3a0193284e94d6667511c420 Mon Sep 17 00:00:00 2001
From: Ara Jermakyan <ara.jermakyan@gmail.com>
Date: Sat, 30 Oct 2021 19:15:06 -0700
Subject: [PATCH] Update Decode Segment and add test case to parser

---
 parser_test.go | 10 ++++++++++
 token.go       |  8 ++++----
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/parser_test.go b/parser_test.go
index 6d1b7c71..599d25dd 100644
--- a/parser_test.go
+++ b/parser_test.go
@@ -57,6 +57,16 @@ var jwtTestData = []struct {
 		nil,
 		jwt.SigningMethodRS256,
 	},
+	{
+		"basic with padding",
+		"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJmb28iOiJwYWRkZWRiYXIifQ==.20kGGJaYekGTRFf8b0TwhuETcR8lv5z2363X5jf7G1yTWVTwOmte5Ii8L8_OQbYwPoiVHmZY6iJPbt_DhCN42AeFY74BcsUhR-BVrYUVhKK0RppuzEcSlILDNeQsJDLEL035CPm1VO6Jrgk7enQPIctVxUesRgswP71OpGvJxy3j1k_J8p0WzZvRZTe1D_2Misa0UDGwnEIHhmr97fIpMSZjFxlcygQw8QN34IHLHIXMaTY1eiCf4CCr6rOS9wUeu7P3CPkmFq9XhxBT_LLCmIMhHnxP5x27FUJE_JZlfek0MmARcrhpsZS2sFhHAiWrjxjOE27jkDtv1nEwn65wMw==",
+		defaultKeyFunc,
+		jwt.MapClaims{"foo": "paddedbar"},
+		true,
+		0,
+		nil,
+		jwt.SigningMethodRS256,
+	},
 	{
 		"basic expired",
 		"", // autogen
diff --git a/token.go b/token.go
index a5e25904..35a21652 100644
--- a/token.go
+++ b/token.go
@@ -104,7 +104,7 @@ func ParseWithClaims(tokenString string, claims Claims, keyFunc Keyfunc, options
 // Deprecated: In a future release, we will demote this function to a non-exported function, since it
 // should only be used internally
 func EncodeSegment(seg []byte) string {
-	return strings.TrimRight(base64.RawURLEncoding.EncodeToString(seg), "=")
+	return base64.RawURLEncoding.EncodeToString(seg)
 }
 
 // DecodeSegment decodes a JWT specific base64url encoding with padding stripped
@@ -112,9 +112,9 @@ func EncodeSegment(seg []byte) string {
 // Deprecated: In a future release, we will demote this function to a non-exported function, since it
 // should only be used internally
 func DecodeSegment(seg string) ([]byte, error) {
-	if l := len(seg) % 4; l > 0 {
-		seg += strings.Repeat("=", 4-l)
+	if strings.Contains(seg, "=") {
+		return base64.URLEncoding.DecodeString(seg)
 	}
 
-	return base64.URLEncoding.DecodeString(seg)
+	return base64.RawURLEncoding.DecodeString(seg)
 }