key-cape/src/internal/server/oidc/handoff_test.go

package oidc_test

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"keycape/internal/domain"
	"keycape/internal/server/oidc"
)

func boolPtr(v bool) *bool { return &v }

func isolationClients() map[string]*domain.Client {
	return map[string]*domain.Client{
		"coulomb-social": {
			ClientID:        "coulomb-social",
			DisplayName:     "coulomb.social",
			RedirectURIs:    []string{"https://coulomb.social/auth/callback/"},
			AllowedScopes:   []string{"openid", "profile"},
			ClientType:      "public",
			MFARequired:     boolPtr(false),
			RegistrationURL: "https://users.example.com/register",
			EnrollmentURL:   "https://users.example.com/enroll",
		},
		"openbao-console": {
			ClientID:      "openbao-console",
			DisplayName:   "OpenBao",
			RedirectURIs:  []string{"https://bao.example.com/oidc/callback"},
			AllowedScopes: []string{"openid", "profile"},
			ClientType:    "public",
		},
	}
}

func isolationHandler(auth domain.AuthProvider, mfa domain.MFAProvider) *oidc.AuthorizeHandler {
	return &oidc.AuthorizeHandler{
		ClientConfig: isolationClients(),
		Auth:         auth,
		MFA:          mfa,
		Sessions:     oidc.NewSessionStore(),
		Logins:       oidc.NewLoginSessionStore(),
		Handoffs:     oidc.NewHandoffStore(),
		Emitter:      &captureEmitter{},
	}
}

func TestHandoff_PromptCreate_EligibleClientRedirectsToAllowList(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{authorizeURL: "https://authelia.example/auth"}, &mockMFAProvider{})
	params := url.Values{
		"client_id":             {"coulomb-social"},
		"redirect_uri":          {"https://coulomb.social/auth/callback/"},
		"response_type":         {"code"},
		"scope":                 {"openid profile"},
		"state":                 {"app-state"},
		"code_challenge":        {"E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"},
		"code_challenge_method": {"S256"},
		"prompt":                {"create"},
		"tenant_hint":           {"tenant:coulomb"},
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/authorize?"+params.Encode(), nil))
	if rec.Code != http.StatusFound {
		t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
	}
	loc, err := url.Parse(rec.Header().Get("Location"))
	if err != nil {
		t.Fatal(err)
	}
	if loc.Host != "users.example.com" || loc.Path != "/register" {
		t.Fatalf("expected allow-listed registration URL, got %s", loc)
	}
	if loc.Query().Get("kc_handoff") == "" {
		t.Fatal("expected kc_handoff on registration redirect")
	}
}

func TestHandoff_PromptCreate_IneligibleClientRejected(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{}, &mockMFAProvider{})
	params := url.Values{
		"client_id":             {"openbao-console"},
		"redirect_uri":          {"https://bao.example.com/oidc/callback"},
		"response_type":         {"code"},
		"scope":                 {"openid profile"},
		"state":                 {"app-state"},
		"code_challenge":        {"E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"},
		"code_challenge_method": {"S256"},
		"prompt":                {"create"},
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/authorize?"+params.Encode(), nil))
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
	}
}

func TestHandoff_ReturnRestartsAuthorizeWithoutToken(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{authorizeURL: "https://authelia.example/auth"}, &mockMFAProvider{})
	params := url.Values{
		"client_id":             {"coulomb-social"},
		"redirect_uri":          {"https://coulomb.social/auth/callback/"},
		"response_type":         {"code"},
		"scope":                 {"openid profile"},
		"state":                 {"app-state"},
		"code_challenge":        {"E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"},
		"code_challenge_method": {"S256"},
		"prompt":                {"create"},
	}
	start := httptest.NewRecorder()
	h.ServeHTTP(start, httptest.NewRequest(http.MethodGet, "/authorize?"+params.Encode(), nil))
	token := mustQuery(t, start.Header().Get("Location"), "kc_handoff")

	ret := httptest.NewRecorder()
	h.ServeHTTP(ret, httptest.NewRequest(http.MethodGet, "/authorize/return?kc_handoff="+url.QueryEscape(token), nil))
	if ret.Code != http.StatusFound {
		t.Fatalf("return status=%d body=%s", ret.Code, ret.Body.String())
	}
	loc, err := url.Parse(ret.Header().Get("Location"))
	if err != nil {
		t.Fatal(err)
	}
	if loc.Path != "/authorize" {
		t.Fatalf("return must restart /authorize, got %s", loc)
	}
	if loc.Query().Get("code") != "" {
		t.Fatal("handoff return must not mint a token or code")
	}
	if loc.Query().Get("client_id") != "coulomb-social" {
		t.Fatalf("client_id not preserved: %s", loc)
	}
}

func TestHandoff_ReplayRejected(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{authorizeURL: "https://authelia.example/auth"}, &mockMFAProvider{})
	params := url.Values{
		"client_id":             {"coulomb-social"},
		"redirect_uri":          {"https://coulomb.social/auth/callback/"},
		"response_type":         {"code"},
		"scope":                 {"openid"},
		"state":                 {"app-state"},
		"code_challenge":        {"E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"},
		"code_challenge_method": {"S256"},
		"prompt":                {"create"},
	}
	start := httptest.NewRecorder()
	h.ServeHTTP(start, httptest.NewRequest(http.MethodGet, "/authorize?"+params.Encode(), nil))
	token := mustQuery(t, start.Header().Get("Location"), "kc_handoff")

	first := httptest.NewRecorder()
	h.ServeHTTP(first, httptest.NewRequest(http.MethodGet, "/authorize/return?kc_handoff="+url.QueryEscape(token), nil))
	if first.Code != http.StatusFound {
		t.Fatalf("first return status=%d", first.Code)
	}
	replay := httptest.NewRecorder()
	h.ServeHTTP(replay, httptest.NewRequest(http.MethodGet, "/authorize/return?kc_handoff="+url.QueryEscape(token), nil))
	if replay.Code != http.StatusBadRequest {
		t.Fatalf("replay status=%d body=%s", replay.Code, replay.Body.String())
	}
}

func TestHandoff_TamperedEnvelopeRejected(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{authorizeURL: "https://authelia.example/auth"}, &mockMFAProvider{})
	params := url.Values{
		"client_id":             {"coulomb-social"},
		"redirect_uri":          {"https://coulomb.social/auth/callback/"},
		"response_type":         {"code"},
		"scope":                 {"openid"},
		"state":                 {"app-state"},
		"code_challenge":        {"E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"},
		"code_challenge_method": {"S256"},
		"prompt":                {"create"},
	}
	start := httptest.NewRecorder()
	h.ServeHTTP(start, httptest.NewRequest(http.MethodGet, "/authorize?"+params.Encode(), nil))
	token := mustQuery(t, start.Header().Get("Location"), "kc_handoff")
	tampered := token + "x"

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/authorize/return?kc_handoff="+url.QueryEscape(tampered), nil))
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("tampered status=%d body=%s", rec.Code, rec.Body.String())
	}
}

func TestHandoff_UnknownUserOffersSignupWithoutToken(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{callbackErr: domain.ErrAuthFailed}, &mockMFAProvider{})
	h.PendingStates().Store("s-unknown", &oidc.PendingState{
		ClientID:    "coulomb-social",
		RedirectURI: "https://coulomb.social/auth/callback/",
		State:       "s-unknown",
		ExpiresAt:   time.Now().Add(time.Minute),
	})
	rec := httptest.NewRecorder()
	h.ServeHTTPCallback(rec, httptest.NewRequest(http.MethodGet, "/authorize/callback?code=x&state=s-unknown", nil))
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
	}
	body := rec.Body.String()
	if !strings.Contains(body, "/authorize/register?state=s-unknown") {
		t.Fatalf("expected signup link, body=%s", body)
	}
	if strings.Contains(body, "code=") {
		t.Fatal("unknown-user page must not mint a code")
	}
}

func TestHandoff_UnknownUserIneligibleHasNoSignupLink(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{callbackErr: domain.ErrAuthFailed}, &mockMFAProvider{})
	h.PendingStates().Store("s-admin", &oidc.PendingState{
		ClientID:    "openbao-console",
		RedirectURI: "https://bao.example.com/oidc/callback",
		State:       "s-admin",
		ExpiresAt:   time.Now().Add(time.Minute),
	})
	rec := httptest.NewRecorder()
	h.ServeHTTPCallback(rec, httptest.NewRequest(http.MethodGet, "/authorize/callback?code=x&state=s-admin", nil))
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
	}
	if strings.Contains(rec.Body.String(), "/authorize/register") {
		t.Fatal("ineligible client must not receive a registration link")
	}
}

func TestAuthorizeCallback_ExpiredStateRejected(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{callbackResult: &domain.AuthResult{Username: "alice"}}, &mockMFAProvider{})
	h.PendingStates().Store("expired", &oidc.PendingState{
		ClientID:    "coulomb-social",
		RedirectURI: "https://coulomb.social/auth/callback/",
		State:       "expired",
		ExpiresAt:   time.Now().Add(-time.Minute),
	})
	rec := httptest.NewRecorder()
	h.ServeHTTPCallback(rec, httptest.NewRequest(http.MethodGet, "/authorize/callback?code=x&state=expired", nil))
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
	}
}

func TestAuthorizeCallback_ReplayAfterSuccessRejected(t *testing.T) {
	h := isolationHandler(&mockAuthProvider{callbackResult: &domain.AuthResult{Username: "alice"}}, &mockMFAProvider{required: true})
	h.PendingStates().Store("once", &oidc.PendingState{
		ClientID:    "coulomb-social",
		RedirectURI: "https://coulomb.social/auth/callback/",
		State:       "once",
		ExpiresAt:   time.Now().Add(time.Minute),
	})
	first := httptest.NewRecorder()
	h.ServeHTTPCallback(first, httptest.NewRequest(http.MethodGet, "/authorize/callback?code=x&state=once", nil))
	if first.Code != http.StatusFound {
		t.Fatalf("first status=%d body=%s", first.Code, first.Body.String())
	}
	second := httptest.NewRecorder()
	h.ServeHTTPCallback(second, httptest.NewRequest(http.MethodGet, "/authorize/callback?code=x&state=once", nil))
	if second.Code != http.StatusBadRequest {
		t.Fatalf("replay status=%d body=%s", second.Code, second.Body.String())
	}
}

func mustQuery(t *testing.T, raw, key string) string {
	t.Helper()
	u, err := url.Parse(raw)
	if err != nil {
		t.Fatal(err)
	}
	v := u.Query().Get(key)
	if v == "" {
		t.Fatalf("missing %s in %s", key, raw)
	}
	return v
}