easyai-ai-gateway/apps/api/internal/auth/oidc_session_cookie_test.go

package auth

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

func TestAuthenticateResolvesOpaqueOIDCSessionCookie(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	var resolved string
	authenticator.OIDCSessionResolver = func(_ context.Context, raw string) (*User, error) {
		resolved = raw
		return &User{ID: "platform-subject", Source: "oidc", Roles: []string{"user"}}, nil
	}

	request := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "opaque-session-id"})
	user, err := authenticator.Authenticate(request)
	if err != nil {
		t.Fatalf("authenticate OIDC session cookie: %v", err)
	}
	if user.ID != "platform-subject" || user.Source != "oidc" {
		t.Fatalf("unexpected session user: %#v", user)
	}
	if resolved != "opaque-session-id" {
		t.Fatalf("session resolver received %q", resolved)
	}
}

func TestAuthenticateBearerTakesPrecedenceOverOIDCSessionCookie(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	localToken, err := authenticator.SignJWT(&User{ID: "local-user", Source: "gateway", Roles: []string{"user"}}, time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	request := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	request.Header.Set("Authorization", "Bearer "+localToken)
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "not-a-token"})

	user, err := authenticator.Authenticate(request)
	if err != nil {
		t.Fatalf("authenticate bearer token: %v", err)
	}
	if user.ID != "local-user" || user.Source != "gateway" {
		t.Fatalf("cookie overrode explicit bearer credentials: %#v", user)
	}
}

func TestAuthenticateRejectsMalformedExplicitCredentialInsteadOfFallingBackToCookie(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	var resolved bool
	authenticator.OIDCSessionResolver = func(context.Context, string) (*User, error) {
		resolved = true
		return &User{ID: "cookie-manager", Source: "gateway", Roles: []string{"manager"}}, nil
	}
	request := httptest.NewRequest(http.MethodPost, "/api/admin/system/identity/disable", nil)
	request.Header.Set("Authorization", "not-a-bearer-credential")
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "valid-cookie-session"})

	if _, err := authenticator.Authenticate(request); !errors.Is(err, ErrUnauthorized) {
		t.Fatalf("malformed explicit credential error=%v, want unauthorized", err)
	}
	if resolved {
		t.Fatal("malformed Authorization header fell back to the OIDC session cookie")
	}
}

func TestAuthenticateRejectsManagerJWTInQueryWithoutCookieFallback(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	managerToken, err := authenticator.SignJWT(&User{ID: "manager", Source: "gateway", Roles: []string{"manager"}}, time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	var resolved bool
	authenticator.OIDCSessionResolver = func(context.Context, string) (*User, error) {
		resolved = true
		return &User{ID: "cookie-manager", Source: "gateway", Roles: []string{"manager"}}, nil
	}
	request := httptest.NewRequest(http.MethodPost, "/api/admin/system/identity/disable?key="+managerToken, nil)
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "valid-cookie-session"})

	if _, err := authenticator.Authenticate(request); !errors.Is(err, ErrUnauthorized) {
		t.Fatalf("manager query JWT error=%v, want unauthorized", err)
	}
	if resolved {
		t.Fatal("rejected query credential fell back to the OIDC session cookie")
	}
}

func TestAuthenticateRejectsInvalidOIDCSessionCookie(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	request := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "not-a-token"})
	if _, err := authenticator.Authenticate(request); err == nil {
		t.Fatal("invalid OIDC session cookie was accepted")
	}
}

func TestAuthenticatorReadsOIDCSessionResolverAndLegacyPolicyDynamically(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	currentUser := &User{ID: "first", Roles: []string{"manager"}}
	authenticator.OIDCSessionResolverProvider = func(ctx context.Context, sessionID string) (*User, error) {
		return currentUser, nil
	}
	authenticator.LegacyJWTEnabledProvider = func() bool { return false }
	request := httptest.NewRequest(http.MethodGet, "/", nil)
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "opaque-session-id"})
	user, err := authenticator.Authenticate(request)
	if err != nil || user.ID != "first" {
		t.Fatalf("first runtime session resolution failed: user=%#v err=%v", user, err)
	}
	currentUser = &User{ID: "second", Roles: []string{"manager"}}
	user, err = authenticator.Authenticate(request)
	if err != nil || user.ID != "second" {
		t.Fatalf("swapped runtime session resolution failed: user=%#v err=%v", user, err)
	}
	if authenticator.legacyJWTEnabled() {
		t.Fatal("dynamic legacy JWT policy was ignored")
	}
}

func TestAuthenticateKeepsOnlySignedBreakGlassManagerWhenLegacyJWTDisabled(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	authenticator.LegacyJWTEnabledProvider = func() bool { return false }

	managerToken, err := authenticator.SignJWT(&User{ID: "manager", Source: "gateway", Roles: []string{"manager"}}, time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	managerRequest := httptest.NewRequest(http.MethodGet, "/api/admin/system/identity/configuration", nil)
	managerRequest.Header.Set("Authorization", "Bearer "+managerToken)
	manager, err := authenticator.Authenticate(managerRequest)
	if err != nil || manager.ID != "manager" {
		t.Fatalf("signed break-glass manager was rejected: user=%#v err=%v", manager, err)
	}

	userToken, err := authenticator.SignJWT(&User{ID: "user", Source: "gateway", Roles: []string{"user"}}, time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	userRequest := httptest.NewRequest(http.MethodGet, "/api/v1/me", nil)
	userRequest.Header.Set("Authorization", "Bearer "+userToken)
	if _, err := authenticator.Authenticate(userRequest); !errors.Is(err, ErrUnauthorized) {
		t.Fatalf("ordinary local JWT error = %v, want unauthorized", err)
	}

	legacyManager := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
		"sub": "legacy-manager", "source": "gateway", "role": []string{"manager"},
		"iat": time.Now().Unix(), "exp": time.Now().Add(time.Hour).Unix(),
	})
	legacyManagerToken, err := legacyManager.SignedString([]byte(authenticator.JWTSecret))
	if err != nil {
		t.Fatal(err)
	}
	legacyRequest := httptest.NewRequest(http.MethodGet, "/api/admin/system/identity/configuration", nil)
	legacyRequest.Header.Set("Authorization", "Bearer "+legacyManagerToken)
	if _, err := authenticator.Authenticate(legacyRequest); !errors.Is(err, ErrUnauthorized) {
		t.Fatalf("legacy manager without token purpose error = %v, want unauthorized", err)
	}
}

func TestPublicRouteIgnoresExpiredOptionalOIDCSession(t *testing.T) {
	authenticator := New("local-jwt-secret", "", "")
	authenticator.OIDCSessionResolver = func(context.Context, string) (*User, error) {
		return nil, &RequestAuthError{Status: http.StatusUnauthorized, Code: "OIDC_SESSION_EXPIRED", Message: "expired"}
	}
	request := httptest.NewRequest(http.MethodGet, "/api/v1/public/catalog/providers", nil)
	request.AddCookie(&http.Cookie{Name: OIDCSessionCookieName, Value: "opaque-session-id"})
	recorder := httptest.NewRecorder()
	authenticator.Require(PermissionPublic, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})).ServeHTTP(recorder, request)
	if recorder.Code != http.StatusNoContent {
		t.Fatalf("public route status = %d, want %d", recorder.Code, http.StatusNoContent)
	}
}