From 50803c27881302a74bfd536d6a3eacb9da4402bb Mon Sep 17 00:00:00 2001
From: Skolli <tanusuch@gmail.com>
Date: Tue, 31 Mar 2026 13:24:47 -0400
Subject: [PATCH] fix(statusserver): improve bearer token parsing and add
 helper tests

Signed-off-by: Skolli <tanusuch@gmail.com>
---
 pkg/statusserver/auth.go       |  4 +--
 pkg/statusserver/auth_test.go  | 61 +++++++++++++++++++++++++++++++
 pkg/statusserver/utils_test.go | 65 ++++++++++++++++++++++++++++++++++
 3 files changed, 128 insertions(+), 2 deletions(-)
 create mode 100644 pkg/statusserver/auth_test.go
 create mode 100644 pkg/statusserver/utils_test.go

diff --git a/pkg/statusserver/auth.go b/pkg/statusserver/auth.go
index 603cf76b91..f7835038d3 100644
--- a/pkg/statusserver/auth.go
+++ b/pkg/statusserver/auth.go
@@ -114,9 +114,9 @@ func (p *projectedServiceAccountTokenAuthorizer) Authorize(ctx context.Context,
 }
 
 func extractRawToken(authHeader string) string {
-	parts := strings.Split(authHeader, " ")
+	parts := strings.Fields(authHeader)
 
-	if len(parts) != 2 || parts[0] != "Bearer" {
+	if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
 		return ""
 	}
 
diff --git a/pkg/statusserver/auth_test.go b/pkg/statusserver/auth_test.go
new file mode 100644
index 0000000000..ac0d94776a
--- /dev/null
+++ b/pkg/statusserver/auth_test.go
@@ -0,0 +1,61 @@
+package statusserver
+
+import "testing"
+
+func TestExtractRawToken(t *testing.T) {
+	tests := []struct {
+		name       string
+		authHeader string
+		expected   string
+	}{
+		{
+			name:       "valid bearer token",
+			authHeader: "Bearer abc123",
+			expected:   "abc123",
+		},
+		{
+			name:       "empty header",
+			authHeader: "",
+			expected:   "",
+		},
+		{
+			name:       "missing token",
+			authHeader: "Bearer",
+			expected:   "",
+		},
+		{
+			name:       "wrong auth scheme",
+			authHeader: "Basic abc123",
+			expected:   "",
+		},
+		{
+			name:       "too many parts",
+			authHeader: "Bearer abc def",
+			expected:   "",
+		},
+		{
+			name:       "multiple spaces",
+			authHeader: "Bearer    abc123",
+			expected:   "abc123",
+		},
+		{
+			name:       "lowercase bearer",
+			authHeader: "bearer abc123",
+			expected:   "abc123",
+		},
+		{
+			name:       "leading and trailing spaces",
+			authHeader: "   Bearer abc123   ",
+			expected:   "abc123",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := extractRawToken(tt.authHeader)
+			if got != tt.expected {
+				t.Fatalf("extractRawToken(%q) = %q, want %q", tt.authHeader, got, tt.expected)
+			}
+		})
+	}
+}
\ No newline at end of file
diff --git a/pkg/statusserver/utils_test.go b/pkg/statusserver/utils_test.go
new file mode 100644
index 0000000000..40d2dec94b
--- /dev/null
+++ b/pkg/statusserver/utils_test.go
@@ -0,0 +1,65 @@
+package statusserver
+
+import "testing"
+
+func TestTokenAudience(t *testing.T) {
+	tests := []struct {
+		name      string
+		namespace string
+		trainJob  string
+		expected  string
+	}{
+		{
+			name:      "default namespace",
+			namespace: "default",
+			trainJob:  "mnist",
+			expected:  "trainer.kubeflow.org/v1alpha1/namespaces/default/trainjobs/mnist/status",
+		},
+		{
+			name:      "kubeflow namespace",
+			namespace: "kubeflow",
+			trainJob:  "llm-job",
+			expected:  "trainer.kubeflow.org/v1alpha1/namespaces/kubeflow/trainjobs/llm-job/status",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := TokenAudience(tt.namespace, tt.trainJob)
+			if got != tt.expected {
+				t.Fatalf("TokenAudience(%q, %q) = %q, want %q", tt.namespace, tt.trainJob, got, tt.expected)
+			}
+		})
+	}
+}
+
+func TestStatusUrl(t *testing.T) {
+	tests := []struct {
+		name      string
+		namespace string
+		trainJob  string
+		expected  string
+	}{
+		{
+			name:      "default namespace",
+			namespace: "default",
+			trainJob:  "mnist",
+			expected:  "/apis/trainer.kubeflow.org/v1alpha1/namespaces/default/trainjobs/mnist/status",
+		},
+		{
+			name:      "kubeflow namespace",
+			namespace: "kubeflow",
+			trainJob:  "llm-job",
+			expected:  "/apis/trainer.kubeflow.org/v1alpha1/namespaces/kubeflow/trainjobs/llm-job/status",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := StatusUrl(tt.namespace, tt.trainJob)
+			if got != tt.expected {
+				t.Fatalf("StatusUrl(%q, %q) = %q, want %q", tt.namespace, tt.trainJob, got, tt.expected)
+			}
+		})
+	}
+}
\ No newline at end of file