From 2f0e9add62078527821828c76865661aa7718a84 Mon Sep 17 00:00:00 2001
From: Michael Fridman <mfridman@buf.build>
Date: Fri, 21 Mar 2025 16:42:51 -0400
Subject: [PATCH] Backporting 0951d18 to v4

---
 jwt_test.go | 89 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 parser.go   | 36 ++++++++++++++++++++--
 2 files changed, 122 insertions(+), 3 deletions(-)
 create mode 100644 jwt_test.go

diff --git a/jwt_test.go b/jwt_test.go
new file mode 100644
index 00000000..b01e899d
--- /dev/null
+++ b/jwt_test.go
@@ -0,0 +1,89 @@
+package jwt
+
+import (
+	"testing"
+)
+
+func TestSplitToken(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		input    string
+		expected []string
+		isValid  bool
+	}{
+		{
+			name:     "valid token with three parts",
+			input:    "header.claims.signature",
+			expected: []string{"header", "claims", "signature"},
+			isValid:  true,
+		},
+		{
+			name:     "invalid token with two parts only",
+			input:    "header.claims",
+			expected: nil,
+			isValid:  false,
+		},
+		{
+			name:     "invalid token with one part only",
+			input:    "header",
+			expected: nil,
+			isValid:  false,
+		},
+		{
+			name:     "invalid token with extra delimiter",
+			input:    "header.claims.signature.extra",
+			expected: nil,
+			isValid:  false,
+		},
+		{
+			name:     "invalid empty token",
+			input:    "",
+			expected: nil,
+			isValid:  false,
+		},
+		{
+			name:     "valid token with empty parts",
+			input:    "..signature",
+			expected: []string{"", "", "signature"},
+			isValid:  true,
+		},
+		{
+			// We are just splitting the token into parts, so we don't care about the actual values.
+			// It is up to the caller to validate the parts.
+			name:     "valid token with all parts empty",
+			input:    "..",
+			expected: []string{"", "", ""},
+			isValid:  true,
+		},
+		{
+			name:     "invalid token with just delimiters and extra part",
+			input:    "...",
+			expected: nil,
+			isValid:  false,
+		},
+		{
+			name:     "invalid token with many delimiters",
+			input:    "header.claims.signature..................",
+			expected: nil,
+			isValid:  false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			parts, ok := splitToken(tt.input)
+			if ok != tt.isValid {
+				t.Errorf("expected %t, got %t", tt.isValid, ok)
+			}
+			if ok {
+				for i, part := range tt.expected {
+					if parts[i] != part {
+						t.Errorf("expected %s, got %s", part, parts[i])
+					}
+				}
+			}
+		})
+	}
+}
diff --git a/parser.go b/parser.go
index 9dd36e5a..0fc510a0 100644
--- a/parser.go
+++ b/parser.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 )
 
+const tokenDelimiter = "."
+
 type Parser struct {
 	// If populated, only these methods will be considered valid.
 	//
@@ -122,9 +124,10 @@ func (p *Parser) ParseWithClaims(tokenString string, claims Claims, keyFunc Keyf
 // It's only ever useful in cases where you know the signature is valid (because it has
 // been checked previously in the stack) and you want to extract values from it.
 func (p *Parser) ParseUnverified(tokenString string, claims Claims) (token *Token, parts []string, err error) {
-	parts = strings.Split(tokenString, ".")
-	if len(parts) != 3 {
-		return nil, parts, NewValidationError("token contains an invalid number of segments", ValidationErrorMalformed)
+	var ok bool
+	parts, ok = splitToken(tokenString)
+	if !ok {
+		return nil, nil, NewValidationError("token contains an invalid number of segments", ValidationErrorMalformed)
 	}
 
 	token = &Token{Raw: tokenString}
@@ -174,3 +177,30 @@ func (p *Parser) ParseUnverified(tokenString string, claims Claims) (token *Toke
 
 	return token, parts, nil
 }
+
+// splitToken splits a token string into three parts: header, claims, and signature. It will only
+// return true if the token contains exactly two delimiters and three parts. In all other cases, it
+// will return nil parts and false.
+func splitToken(token string) ([]string, bool) {
+	parts := make([]string, 3)
+	header, remain, ok := strings.Cut(token, tokenDelimiter)
+	if !ok {
+		return nil, false
+	}
+	parts[0] = header
+	claims, remain, ok := strings.Cut(remain, tokenDelimiter)
+	if !ok {
+		return nil, false
+	}
+	parts[1] = claims
+	// One more cut to ensure the signature is the last part of the token and there are no more
+	// delimiters. This avoids an issue where malicious input could contain additional delimiters
+	// causing unecessary overhead parsing tokens.
+	signature, _, unexpected := strings.Cut(remain, tokenDelimiter)
+	if unexpected {
+		return nil, false
+	}
+	parts[2] = signature
+
+	return parts, true
+}