fix: forward leader api auth issue

Author: Code2Life

internal/server/router/assign_host_port.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/NexusGPU/tensor-fusion/internal/constants"
 	"github.com/NexusGPU/tensor-fusion/internal/portallocator"
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	"github.com/gin-gonic/gin"
@@ -26,9 +25,8 @@ func NewAssignHostPortRouter(ctx context.Context, allocator *portallocator.PortA
 
 func (r *AssignHostPortRouter) AssignHostPort(ctx *gin.Context) {
 	podName := ctx.Query("podName")
-	token := ctx.Request.Header.Get(constants.AuthorizationHeader)
-
-	if token == "" {
+	token, ok := utils.ExtractBearerToken(ctx)
+	if !ok {
 		log.FromContext(ctx).Error(nil, "assigned host port failed, missing token", "podName", podName)
 		ctx.String(http.StatusUnauthorized, "missing authorization header")
 		return

internal/server/router/assign_index.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/NexusGPU/tensor-fusion/internal/constants"
 	"github.com/NexusGPU/tensor-fusion/internal/indexallocator"
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	"github.com/gin-gonic/gin"
@@ -26,9 +25,8 @@ func NewAssignIndexRouter(ctx context.Context, allocator *indexallocator.IndexAl
 
 func (r *AssignIndexRouter) AssignIndex(ctx *gin.Context) {
 	podName := ctx.Query("podName")
-	token := ctx.Request.Header.Get(constants.AuthorizationHeader)
-
-	if token == "" {
+	token, ok := utils.ExtractBearerToken(ctx)
+	if !ok {
 		log.FromContext(ctx).Error(nil, "assigned index failed, missing token", "podName", podName)
 		ctx.String(http.StatusUnauthorized, "missing authorization header")
 		return
@@ -47,7 +45,8 @@ func (r *AssignIndexRouter) AssignIndex(ctx *gin.Context) {
 		return
 	}
 	if !tokenReview.Status.Authenticated || tokenReview.Status.User.Username != utils.GetSelfServiceAccountNameFull() {
-		log.FromContext(ctx).Error(nil, "assigned index failed, token invalid", "podName", podName)
+		log.FromContext(ctx).Error(nil, "assigned index failed, token invalid", "podName", podName,
+			"authPassed", tokenReview.Status.Authenticated, "username", tokenReview.Status.User.Username)
 		ctx.String(http.StatusUnauthorized, "token authentication failed")
 		return
 	}

internal/utils/config.go

@@ -29,6 +29,11 @@ const (
 var selfServiceAccountName string
 
 func InitServiceAccountConfig() {
+	if os.Getenv("IMPERSONATE_SERVICE_ACCOUNT") != "" {
+		selfServiceAccountName = os.Getenv("IMPERSONATE_SERVICE_ACCOUNT")
+		ctrl.Log.Info("impersonate service account mode detected", "name", selfServiceAccountName)
+		return
+	}
 	data, err := os.ReadFile(ServiceAccountTokenPath)
 	if err != nil {
 		ctrl.Log.Info("service account token not found, run outside of Kubernetes cluster")

internal/utils/svr.go

@@ -0,0 +1,26 @@
+package utils
+
+import (
+	"strings"
+
+	"github.com/NexusGPU/tensor-fusion/internal/constants"
+	"github.com/gin-gonic/gin"
+)
+
+const BearerPrefix = "Bearer "
+
+// ExtractBearerToken extracts the authorization token from the gin context.
+// It handles both cases: token with "Bearer " prefix and token without prefix.
+// Returns the token string (with Bearer prefix stripped if present) and true if token exists.
+// Returns empty string and false if token is missing.
+func ExtractBearerToken(ctx *gin.Context) (string, bool) {
+	token := ctx.Request.Header.Get(constants.AuthorizationHeader)
+	if token == "" {
+		return "", false
+	}
+
+	// Strip Bearer prefix if present
+	token = strings.TrimPrefix(token, BearerPrefix)
+
+	return token, true
+}