easyai-ai-gateway/apps/api/internal/httpapi/oidc_session_test.go

package httpapi

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/oidcsession"
)

func TestStartOIDCLoginSetsEncryptedLaxTransactionAndRedirectsWithPKCE(t *testing.T) {
	cipher, _ := oidcsession.NewCipher(bytes.Repeat([]byte{3}, 32))
	client := &fakeOIDCClient{authorizationURL: "https://auth.example.com/authorize?request=redacted"}
	server := &Server{
		auth: &auth.Authenticator{OIDCVerifier: &auth.OIDCVerifier{}}, oidcClient: client,
		oidcSessions: &fakeOIDCSessions{}, oidcSessionCipher: cipher,
		identityTestCookieSecure: true, logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	request := httptest.NewRequest(http.MethodGet, "/api/v1/auth/oidc/login?returnTo=%2Fworkspace%3Ftab%3Dwallet", nil)
	recorder := httptest.NewRecorder()
	server.startOIDCLogin(recorder, request)
	response := recorder.Result()
	defer response.Body.Close()
	if response.StatusCode != http.StatusSeeOther || response.Header.Get("Location") != client.authorizationURL {
		t.Fatalf("login status=%d location=%q", response.StatusCode, response.Header.Get("Location"))
	}
	cookies := response.Cookies()
	if len(cookies) != 1 || cookies[0].Name != oidcsession.LoginTransactionCookieName || !cookies[0].HttpOnly || !cookies[0].Secure || cookies[0].SameSite != http.SameSiteLaxMode || cookies[0].MaxAge != 600 {
		t.Fatalf("unsafe login transaction cookie: %#v", cookies)
	}
	transaction, err := cipher.DecodeLoginTransaction(cookies[0].Value, cookies[0].Expires.Add(-time.Minute))
	if err != nil || transaction.ReturnTo != "/workspace?tab=wallet" {
		t.Fatalf("transaction=%#v err=%v", transaction, err)
	}
	if client.state == "" || client.nonce == "" || client.challenge == "" {
		t.Fatal("authorization redirect omitted state, nonce or PKCE challenge")
	}
}

func TestStartOIDCLoginRejectsOpenRedirect(t *testing.T) {
	cipher, _ := oidcsession.NewCipher(bytes.Repeat([]byte{3}, 32))
	server := &Server{
		auth: &auth.Authenticator{OIDCVerifier: &auth.OIDCVerifier{}}, oidcClient: &fakeOIDCClient{},
		oidcSessions: &fakeOIDCSessions{}, oidcSessionCipher: cipher, logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	recorder := httptest.NewRecorder()
	server.startOIDCLogin(recorder, httptest.NewRequest(http.MethodGet, "/api/v1/auth/oidc/login?returnTo=https%3A%2F%2Fevil.example", nil))
	if recorder.Code != http.StatusBadRequest || recorder.Header().Get("Set-Cookie") != "" {
		t.Fatalf("open redirect status=%d cookie=%q", recorder.Code, recorder.Header().Get("Set-Cookie"))
	}
}

func TestCompleteOIDCLoginReportsSafeTransactionFailureReason(t *testing.T) {
	cipher, err := oidcsession.NewCipher(bytes.Repeat([]byte{3}, 32))
	if err != nil {
		t.Fatal(err)
	}
	transaction, err := oidcsession.NewLoginTransaction("/", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	encodedTransaction, err := cipher.EncodeLoginTransaction(transaction)
	if err != nil {
		t.Fatal(err)
	}

	for _, test := range []struct {
		name, cookie, state, code, wantReason string
	}{
		{name: "cookie missing", state: "sensitive-state-marker", code: "sensitive-code-marker", wantReason: "COOKIE_MISSING"},
		{name: "transaction invalid", cookie: "invalid-encrypted-cookie", state: "sensitive-state-marker", code: "sensitive-code-marker", wantReason: "TRANSACTION_INVALID"},
		{name: "state mismatch", cookie: encodedTransaction, state: "sensitive-state-marker", code: "sensitive-code-marker", wantReason: "STATE_MISMATCH"},
		{name: "authorization response missing", cookie: encodedTransaction, state: transaction.State, wantReason: "AUTHORIZATION_RESPONSE_MISSING"},
	} {
		t.Run(test.name, func(t *testing.T) {
			var logs bytes.Buffer
			server := &Server{
				cfg:  config.Config{WebBaseURL: "http://localhost:5178"},
				auth: &auth.Authenticator{OIDCVerifier: &auth.OIDCVerifier{}}, oidcClient: &fakeOIDCClient{},
				oidcSessions: &fakeOIDCSessions{}, oidcSessionCipher: cipher,
				logger: slog.New(slog.NewJSONHandler(&logs, nil)),
			}
			query := url.Values{}
			if test.state != "" {
				query.Set("state", test.state)
			}
			if test.code != "" {
				query.Set("code", test.code)
			}
			request := httptest.NewRequest(http.MethodGet, "/api/v1/auth/oidc/callback?"+query.Encode(), nil)
			if test.cookie != "" {
				request.AddCookie(&http.Cookie{Name: oidcsession.LoginTransactionCookieName, Value: test.cookie})
			}
			recorder := httptest.NewRecorder()

			server.completeOIDCLogin(recorder, request)

			if recorder.Code != http.StatusSeeOther {
				t.Fatalf("callback status=%d, want 303", recorder.Code)
			}
			location, err := url.Parse(recorder.Header().Get("Location"))
			if err != nil {
				t.Fatal(err)
			}
			if location.Query().Get("oidcError") != errorCodeOIDCLoginInvalid || location.Query().Get("oidcErrorReason") != test.wantReason {
				t.Fatalf("callback error=%q reason=%q", location.Query().Get("oidcError"), location.Query().Get("oidcErrorReason"))
			}
			diagnosticID := location.Query().Get("oidcDiagnosticId")
			if diagnosticID == "" || !strings.Contains(logs.String(), `"diagnosticId":"`+diagnosticID+`"`) || !strings.Contains(logs.String(), `"reason":"`+test.wantReason+`"`) {
				t.Fatalf("missing correlated safe diagnostic: location=%q logs=%s", location.RawQuery, logs.String())
			}
			for _, secretMarker := range []string{"sensitive-state-marker", "sensitive-code-marker", "invalid-encrypted-cookie"} {
				if strings.Contains(location.RawQuery, secretMarker) || strings.Contains(logs.String(), secretMarker) {
					t.Fatalf("OIDC diagnostic leaked callback security material")
				}
			}
		})
	}
}

func TestDeleteOIDCBrowserSessionIsIdempotentAndExpiresCookie(t *testing.T) {
	sessions := &fakeOIDCSessions{}
	server := &Server{oidcSessions: sessions, identityTestCookieSecure: true}
	request := httptest.NewRequest(http.MethodDelete, "/api/v1/auth/oidc/session", nil)
	request.AddCookie(&http.Cookie{Name: auth.OIDCSessionCookieName, Value: "opaque-session"})
	recorder := httptest.NewRecorder()
	server.deleteOIDCBrowserSession(recorder, request)
	response := recorder.Result()
	defer response.Body.Close()
	if response.StatusCode != http.StatusNoContent || sessions.deleted != "opaque-session" {
		t.Fatalf("delete status=%d session=%q", response.StatusCode, sessions.deleted)
	}
	cookies := response.Cookies()
	if len(cookies) != 1 || cookies[0].Name != auth.OIDCSessionCookieName || cookies[0].MaxAge >= 0 || !cookies[0].HttpOnly || cookies[0].SameSite != http.SameSiteStrictMode {
		t.Fatalf("OIDC session cookie was not safely expired: %#v", cookies)
	}
}

func TestOIDCSessionCSRFAcceptsOnlyAllowedOriginForCookieWrites(t *testing.T) {
	server := &Server{
		cfg:                        config.Config{CORSAllowedOrigin: "https://gateway.example.com"},
		identityTestRevision:       identity.Revision{WebBaseURL: "https://gateway.example.com"},
		identityTestBrowserEnabled: true,
	}
	next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNoContent) })
	handler := server.protectOIDCSessionCookie(next)
	for _, test := range []struct {
		name, method, origin string
		bearer               bool
		wantStatus           int
	}{
		{name: "missing origin", method: http.MethodPost, wantStatus: http.StatusForbidden},
		{name: "foreign origin", method: http.MethodDelete, origin: "https://evil.example", wantStatus: http.StatusForbidden},
		{name: "foreign origin cannot read cookie authenticated data", method: http.MethodGet, origin: "https://evil.example", wantStatus: http.StatusForbidden},
		{name: "allowed origin", method: http.MethodPatch, origin: "https://gateway.example.com", wantStatus: http.StatusNoContent},
		{name: "safe request", method: http.MethodGet, wantStatus: http.StatusNoContent},
		{name: "explicit bearer bypasses cookie csrf", method: http.MethodPost, bearer: true, wantStatus: http.StatusNoContent},
	} {
		t.Run(test.name, func(t *testing.T) {
			request := httptest.NewRequest(test.method, "/api/workspace/tasks", nil)
			request.AddCookie(&http.Cookie{Name: auth.OIDCSessionCookieName, Value: "session-token"})
			request.Header.Set("Origin", test.origin)
			if test.bearer {
				request.Header.Set("Authorization", "Bearer explicit-token")
			}
			recorder := httptest.NewRecorder()
			handler.ServeHTTP(recorder, request)
			if recorder.Code != test.wantStatus {
				t.Fatalf("status = %d, want %d", recorder.Code, test.wantStatus)
			}
		})
	}
}

func TestOIDCSessionCSRFMalformedAuthorizationCannotBypassForeignOrigin(t *testing.T) {
	authenticator := auth.New("local-jwt-secret", "", "")
	authenticator.OIDCSessionResolver = func(context.Context, string) (*auth.User, error) {
		return &auth.User{ID: "manager", Source: "gateway", Roles: []string{"manager"}}, nil
	}
	server := &Server{
		auth:                       authenticator,
		identityTestRevision:       identity.Revision{WebBaseURL: "https://gateway.example.com"},
		identityTestBrowserEnabled: true,
	}
	called := false
	handler := server.protectOIDCSessionCookie(server.requireAdmin(auth.PermissionManager, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		called = true
		w.WriteHeader(http.StatusNoContent)
	})))
	request := httptest.NewRequest(http.MethodPost, "/api/admin/system/identity/disable", nil)
	request.Header.Set("Origin", "https://evil.example.com")
	request.Header.Set("Authorization", "malformed")
	request.AddCookie(&http.Cookie{Name: auth.OIDCSessionCookieName, Value: "manager-session"})
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)
	if recorder.Code != http.StatusForbidden || called {
		t.Fatalf("malformed Authorization CSRF status=%d handler_called=%t", recorder.Code, called)
	}
}

func TestAdminRouteRejectsManagerJWTInQueryAndAcceptsAuthorizationHeader(t *testing.T) {
	authenticator := auth.New("local-jwt-secret", "", "")
	managerToken, err := authenticator.SignJWT(&auth.User{ID: "manager", Source: "gateway", Roles: []string{"manager"}}, time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	server := &Server{auth: authenticator}
	handler := server.requireAdmin(auth.PermissionManager, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))

	queryRequest := httptest.NewRequest(http.MethodPost, "/api/admin/system/identity/disable?key="+managerToken, nil)
	queryRecorder := httptest.NewRecorder()
	handler.ServeHTTP(queryRecorder, queryRequest)
	if queryRecorder.Code != http.StatusUnauthorized {
		t.Fatalf("query manager JWT status=%d, want 401", queryRecorder.Code)
	}

	headerRequest := httptest.NewRequest(http.MethodPost, "/api/admin/system/identity/disable", nil)
	headerRequest.Header.Set("Authorization", "Bearer "+managerToken)
	headerRecorder := httptest.NewRecorder()
	handler.ServeHTTP(headerRecorder, headerRequest)
	if headerRecorder.Code != http.StatusNoContent {
		t.Fatalf("Authorization manager JWT status=%d, want 204", headerRecorder.Code)
	}
}

func TestCORSUsesOnlyCurrentActiveIdentityWebOriginWithoutRestart(t *testing.T) {
	server := &Server{
		cfg: config.Config{CORSAllowedOrigin: "https://bootstrap.example.com"},
		identityTestRevision: identity.Revision{
			ID: "active-revision", State: identity.RevisionActive, WebBaseURL: "https://gateway.example.com",
		},
	}
	handler := server.cors(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNoContent) }))
	request := func(origin string) *httptest.ResponseRecorder {
		r := httptest.NewRequest(http.MethodOptions, "/api/admin/system/identity/configuration", nil)
		r.Header.Set("Origin", origin)
		r.Header.Set("Access-Control-Request-Method", http.MethodGet)
		w := httptest.NewRecorder()
		handler.ServeHTTP(w, r)
		return w
	}

	active := request("https://gateway.example.com")
	if active.Header().Get("Access-Control-Allow-Origin") != "https://gateway.example.com" || active.Header().Get("Access-Control-Allow-Credentials") != "true" {
		t.Fatalf("active Web origin headers=%v", active.Header())
	}
	if evil := request("https://evil.example.com"); evil.Header().Get("Access-Control-Allow-Origin") != "" {
		t.Fatalf("evil origin was allowed: headers=%v", evil.Header())
	}

	server.identityTestRevision = identity.Revision{}
	if disabled := request("https://gateway.example.com"); disabled.Header().Get("Access-Control-Allow-Origin") != "" {
		t.Fatalf("disabled Revision retained dynamic origin: headers=%v", disabled.Header())
	}
	if bootstrap := request("https://bootstrap.example.com"); bootstrap.Header().Get("Access-Control-Allow-Origin") != "https://bootstrap.example.com" {
		t.Fatalf("deployment bootstrap origin stopped working: headers=%v", bootstrap.Header())
	}
}

func TestOIDCSessionCSRFIsInactiveWhenOIDCIsDisabled(t *testing.T) {
	server := &Server{}
	request := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", nil)
	request.AddCookie(&http.Cookie{Name: auth.OIDCSessionCookieName, Value: "irrelevant-cookie"})
	recorder := httptest.NewRecorder()
	server.protectOIDCSessionCookie(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNoContent) })).ServeHTTP(recorder, request)
	if recorder.Code != http.StatusNoContent {
		t.Fatalf("OIDC-disabled request status = %d, want 204", recorder.Code)
	}
}

type fakeOIDCClient struct {
	authorizationURL        string
	state, nonce, challenge string
}

func (f *fakeOIDCClient) AuthorizationURL(_ context.Context, state, nonce, challenge string) (string, error) {
	f.state, f.nonce, f.challenge = state, nonce, challenge
	return f.authorizationURL, nil
}
func (f *fakeOIDCClient) ExchangeCode(context.Context, string, string) (auth.OIDCTokenResponse, error) {
	return auth.OIDCTokenResponse{}, nil
}
func (f *fakeOIDCClient) VerifyIDToken(context.Context, string, string) (string, error) {
	return "", nil
}
func (f *fakeOIDCClient) Refresh(context.Context, string) (auth.OIDCTokenResponse, error) {
	return auth.OIDCTokenResponse{}, nil
}
func (f *fakeOIDCClient) RevokeRefreshToken(context.Context, string) error { return nil }
func (f *fakeOIDCClient) EndSessionURL(context.Context, string) (string, error) {
	return "https://gateway.example.com/", nil
}

type fakeOIDCSessions struct{ deleted string }

func (f *fakeOIDCSessions) Create(context.Context, oidcsession.TokenBundle, *auth.User) (string, error) {
	return "opaque-session", nil
}
func (f *fakeOIDCSessions) Resolve(context.Context, string) (*auth.User, error) { return nil, nil }
func (f *fakeOIDCSessions) Delete(_ context.Context, raw string) (oidcsession.TokenBundle, error) {
	f.deleted = raw
	return oidcsession.TokenBundle{}, nil
}
func (f *fakeOIDCSessions) Cleanup(context.Context) (int64, error) { return 0, nil }