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

package httpapi

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"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/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{
		cfg:  config.Config{OIDCEnabled: true, OIDCBrowserSessionEnabled: true, OIDCSessionCookieSecure: true},
		auth: &auth.Authenticator{OIDCVerifier: &auth.OIDCVerifier{}}, oidcClient: client,
		oidcSessions: &fakeOIDCSessions{}, oidcSessionCipher: cipher,
		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{
		cfg:  config.Config{OIDCEnabled: true, OIDCBrowserSessionEnabled: true},
		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 TestDeleteOIDCBrowserSessionIsIdempotentAndExpiresCookie(t *testing.T) {
	sessions := &fakeOIDCSessions{}
	server := &Server{cfg: config.Config{OIDCSessionCookieSecure: true}, oidcSessions: sessions}
	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{OIDCEnabled: true, OIDCBrowserSessionEnabled: true, CORSAllowedOrigin: "https://gateway.example.com"}}
	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: "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 TestOIDCSessionCSRFIsInactiveWhenOIDCIsDisabled(t *testing.T) {
	server := &Server{cfg: config.Config{OIDCEnabled: false, OIDCBrowserSessionEnabled: true}}
	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) 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 }