visionA/visionA-backend/internal/oidctest/server.go

// Package oidctest 提供測試專用的 fake OpenID Connect Identity Provider，
// 模擬 Innovedus Member Center 的對外行為（OIDC discovery / JWKS / token exchange /
// authorize redirect），讓 visionA-backend 的 OIDC client 與 BFF 流程可以在純
// in-process 環境完成 end-to-end 整合測試 — 不需要 docker、不需要 docker-compose、
// 不需要真的 Member Center。
//
// # 設計理由
//
// OB1（internal/oidc）已經為自己的 unit test 寫過一份 fake OIDC server
// （internal/oidc/provider_test.go 裡的 fakeOIDC）。OT1 把它「再寫一次成可重用的
// 公開 package」而不是直接 export 那個 fake，理由有二：
//
//  1. 邊界乾淨：OB1 的 fakeOIDC 是 unit test 用的（檔名 *_test.go 不會出現在
//     production binary 也不會出現在其他 package 的 import path），刻意不拿來當
//     公開 fixture 是為了讓 unit test 自成一格、不被外部依賴牽動。
//
//  2. API 形狀不同：unit test 的 fake 暴露很多 hook（snapshot/withState/skipIDToken
//     …）給「驗測 OIDC client 的錯誤分類」這種白箱測試用；e2e 整合測試需要的是
//     「黑箱模擬完整 BFF flow」— issuer URL、ExchangeCode / authorize-redirect 的
//     整體行為。兩邊的 API 形狀一旦混用反而綁手綁腳。
//
// 實作仍刻意對齊 OB1 的 fakeOIDC（同樣 RS256 / 同樣 endpoint paths /
// 同樣 PKCE 與 nonce 處理），如果未來雙方有差異要對齊，更新本 package 即可。
//
// # 對齊文件
//
//   - oidc-tdd.md §3 BFF Flow 詳細時序圖
//   - oidc-tdd.md §6 PKCE 實作細節
//   - oidc-tdd.md §7 id_token 驗證
//   - adr-010-oidc-bff.md
//
// # 使用範例
//
//	srv := oidctest.NewServer(t,
//	    oidctest.WithClientCredentials("visiona-backend-test", "test-secret"),
//	)
//	defer srv.Close()
//
//	// 把 srv.URL 當作 IssuerURL 給 visionA-backend 的 OIDC provider。
//	provider, _ := oidc.NewProvider(ctx, oidc.ProviderConfig{
//	    IssuerURL:    srv.URL,
//	    ClientID:     srv.ClientID,
//	    ClientSecret: srv.ClientSecret,
//	    RedirectURL:  "http://localhost:8080/api/auth/callback",
//	})
//
//	// 預先告知 fake server：下一個 ExchangeCode 之後要簽發的 id_token claims
//	srv.SetNextIDTokenClaims(map[string]any{
//	    "sub":   "user-from-mc",
//	    "email": "alice@innovedus.com",
//	    "name":  "Alice",
//	})
package oidctest

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sync"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"
)

// Server 是用 httptest.Server 包起來的 fake OIDC IdP。
//
// 提供以下端點：
//   - GET  /.well-known/openid-configuration  → discovery doc
//   - GET  /jwks                              → JWKS（含 1 把 RSA public key）
//   - POST /oauth/token                       → form-encoded code exchange，回 id_token + access_token
//   - GET  /authorize                         → 自動「同意登入」並 redirect 回 client redirect_uri 帶 code
//
// Server 是 goroutine-safe（內部 RWMutex 保護所有 mutable state），
// 但 caller 仍應在「設定下一輪 token 行為 → 觸發 ExchangeCode」之間以同步方式進行，
// 否則 race 行為可預期但難以推理。
type Server struct {
	// URL 是 httptest.Server 的 URL，同時也是 OIDC discovery 的 issuer。
	// caller 直接拿這個當 ProviderConfig.IssuerURL。
	URL string

	// Issuer 等同 URL；保留欄位是為了測試「issuer mismatch」場景時能暫時 override。
	Issuer string

	// ClientID 與 ClientSecret 是 fake server 認可的 confidential client 憑證。
	// /oauth/token 會驗 client_secret；不符直接 401。
	ClientID     string
	ClientSecret string

	// PrivateKey 是用來簽 id_token 的 RSA key；JWKS endpoint 公開對應的 public key。
	// caller 一般不需要直接用，但 IssueIDToken 暴露給「自定 token claims」場景。
	PrivateKey *rsa.PrivateKey
	KeyID      string

	httpServer *httptest.Server

	// ─── mutable state（test 之間用 SetXXX 改寫；rwmu 保護） ───
	rwmu sync.RWMutex

	// nextIDTokenClaims：若非 nil，下一個 /oauth/token response 的 id_token 用這份 claims 簽發。
	// caller 通常在每個 test 開頭 Set 一次；handleToken 用完不會自動清空，方便同個
	// fake server 在多次 ExchangeCode 中重複簽發同一個使用者的 token。
	nextIDTokenClaims map[string]any

	// nextAccessToken：若非空字串，下一個 /oauth/token response 的 access_token 用這個值。
	// 預設為 "fake-access-token"。
	nextAccessToken string

	// 觀測欄位：記錄最後一次 /authorize 與 /oauth/token 收到的關鍵參數，
	// e2e test 可用來驗 BFF 是否把 PKCE / state / nonce 正確帶過來。
	lastAuthorizeQuery url.Values
	lastTokenForm      url.Values

	// codeStore 是 authorization code → 該 code 對應的 PKCE challenge / nonce 的暫存。
	// 模擬真 IdP 會把 code 與當時的 PKCE challenge 綁定，token endpoint 才能驗 PKCE proof。
	codeStore map[string]issuedCode
}

// issuedCode 是 SimulateAuthorizationFlow / IssueAuthCode 簽發 code 時記下的元資料。
// 之後 /oauth/token 用 code 反查出當時的 challenge 比對 PKCE。
type issuedCode struct {
	CodeChallenge       string
	CodeChallengeMethod string
	Nonce               string
	RedirectURI         string
	ClientID            string
	IssuedAt            time.Time
}

// Option 是 NewServer 的功能選項。
type Option func(*Server)

// WithClientCredentials 設定 fake server 認可的 OAuth client_id / client_secret。
// 不呼叫此 option 則使用預設值（visiona-backend-test / test-secret）。
func WithClientCredentials(clientID, clientSecret string) Option {
	return func(s *Server) {
		s.ClientID = clientID
		s.ClientSecret = clientSecret
	}
}

// WithIssuer 強制 override discovery doc 的 issuer claim。
// 一般場景請勿使用；此選項只給「測 issuer mismatch」這種對抗性測試用。
func WithIssuer(issuer string) Option {
	return func(s *Server) {
		s.Issuer = issuer
	}
}

// NewServer 啟動一個 fake OIDC server。
//
// 會立即在 t.Cleanup 註冊關閉動作，caller 不必自己呼叫 Close
// （但保留 Close 公開方法供「同個 test 內提早關閉以驗錯誤情境」使用）。
func NewServer(t *testing.T, opts ...Option) *Server {
	t.Helper()

	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("oidctest: rsa key gen failed: %v", err)
	}

	s := &Server{
		PrivateKey:      priv,
		KeyID:           "oidctest-key-1",
		ClientID:        "visiona-backend-test",
		ClientSecret:    "test-secret",
		nextAccessToken: "fake-access-token",
		codeStore:       make(map[string]issuedCode),
	}
	for _, opt := range opts {
		opt(s)
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", s.handleDiscovery)
	mux.HandleFunc("/jwks", s.handleJWKS)
	mux.HandleFunc("/oauth/token", s.handleToken)
	mux.HandleFunc("/authorize", s.handleAuthorize)

	s.httpServer = httptest.NewServer(mux)
	s.URL = s.httpServer.URL
	if s.Issuer == "" {
		s.Issuer = s.URL
	}

	t.Cleanup(s.Close)
	return s
}

// Close 停掉內部 httptest.Server。重複呼叫安全。
func (s *Server) Close() {
	if s.httpServer != nil {
		s.httpServer.Close()
		s.httpServer = nil
	}
}

// SetNextIDTokenClaims 設定下一次 /oauth/token 回應裡 id_token 的 claims。
//
// 注意：傳入的 map 不會被 merge，而是「整份覆蓋」預設值（除了 iss/aud/exp 由 server 補足）。
// caller 可以放 sub / email / name / nonce 等自定 claim；若漏傳 nonce，handleToken
// 會用 lastAuthorizeQuery 收到的 nonce 補上（模擬真 IdP 行為：authorize 收到的 nonce 會回灌到 id_token）。
//
// 若傳 nil 等同呼叫 ResetIDTokenClaims，回到預設 sub。
func (s *Server) SetNextIDTokenClaims(claims map[string]any) {
	s.rwmu.Lock()
	defer s.rwmu.Unlock()
	if claims == nil {
		s.nextIDTokenClaims = nil
		return
	}
	cp := make(map[string]any, len(claims))
	for k, v := range claims {
		cp[k] = v
	}
	s.nextIDTokenClaims = cp
}

// ResetIDTokenClaims 把「下一輪 id_token claims」回到預設值。
func (s *Server) ResetIDTokenClaims() { s.SetNextIDTokenClaims(nil) }

// SetNextAccessToken 改下一次 /oauth/token response 的 access_token 字串。
// 主要供「驗 backend 是否正確存了 access_token」這種測試。
func (s *Server) SetNextAccessToken(tok string) {
	s.rwmu.Lock()
	defer s.rwmu.Unlock()
	s.nextAccessToken = tok
}

// LastAuthorizeQuery 回傳上一次 /authorize 收到的 query string（複製，可安全修改）。
// e2e test 通常用來驗 BFF 是否正確產 PKCE / state / nonce。
func (s *Server) LastAuthorizeQuery() url.Values {
	s.rwmu.RLock()
	defer s.rwmu.RUnlock()
	return cloneValues(s.lastAuthorizeQuery)
}

// LastTokenForm 回傳上一次 /oauth/token 收到的 form value（複製）。
// 用來驗 ExchangeCode 是否正確帶 client_secret / code_verifier。
func (s *Server) LastTokenForm() url.Values {
	s.rwmu.RLock()
	defer s.rwmu.RUnlock()
	return cloneValues(s.lastTokenForm)
}

// IssueIDToken 直接用 fake server 的 RSA private key 簽一個 id_token，回傳 raw JWT 字串。
//
// 用途：少數場景需要「跳過 token endpoint，直接拿 id_token 餵給 VerifyIDToken」測試
// （例如測 backend 對「不正確 issuer 的 id_token」的拒絕行為）。
func (s *Server) IssueIDToken(claims map[string]any) (string, error) {
	return signJWT(s.PrivateKey, s.KeyID, jose.RS256, claims)
}

// IssueAuthCode 預先簽發一個 authorization code，並把對應的 PKCE challenge / nonce
// 記在 codeStore 中。後續 /oauth/token 收到此 code + 正確 code_verifier 才會放行。
//
// 主要供「不走完整 redirect 流程、直接構造 callback」的測試用。
// 如果你只是要 e2e 跑完整 flow，呼叫 SimulateAuthorizationFlow 即可（會自動 issue code）。
func (s *Server) IssueAuthCode(challenge, challengeMethod, nonce, redirectURI string) (string, error) {
	code, err := randomURLToken(24)
	if err != nil {
		return "", err
	}
	s.rwmu.Lock()
	defer s.rwmu.Unlock()
	s.codeStore[code] = issuedCode{
		CodeChallenge:       challenge,
		CodeChallengeMethod: challengeMethod,
		Nonce:               nonce,
		RedirectURI:         redirectURI,
		ClientID:            s.ClientID,
		IssuedAt:            time.Now(),
	}
	return code, nil
}

// ───────────────────────── HTTP handlers ─────────────────────────

func (s *Server) handleDiscovery(w http.ResponseWriter, r *http.Request) {
	s.rwmu.RLock()
	issuer := s.Issuer
	s.rwmu.RUnlock()

	doc := map[string]any{
		"issuer":                                issuer,
		"authorization_endpoint":                s.URL + "/authorize",
		"token_endpoint":                        s.URL + "/oauth/token",
		"jwks_uri":                              s.URL + "/jwks",
		"response_types_supported":              []string{"code"},
		"id_token_signing_alg_values_supported": []string{"RS256"},
		"subject_types_supported":               []string{"public"},
		"scopes_supported":                      []string{"openid", "email", "profile"},
		"code_challenge_methods_supported":      []string{"S256"},
	}
	writeJSON(w, http.StatusOK, doc)
}

func (s *Server) handleJWKS(w http.ResponseWriter, r *http.Request) {
	jwks := map[string]any{
		"keys": []map[string]any{
			rsaPublicKeyToJWK(&s.PrivateKey.PublicKey, s.KeyID),
		},
	}
	writeJSON(w, http.StatusOK, jwks)
}

// handleAuthorize 模擬「使用者打開 /authorize → 同意登入 → IdP redirect 回 client redirect_uri 帶 code」。
//
// 為了讓測試可以「不開瀏覽器」就跑通整段，我們不顯示登入頁，
// 而是直接把 code 帶上 redirect_uri 立刻 302 回去（等同「使用者已存在 SSO 並自動同意」）。
//
// 我們同時把 code 與當時的 PKCE challenge / nonce 綁起來，
// 後續 /oauth/token 才能驗 PKCE proof，符合真 IdP 行為。
func (s *Server) handleAuthorize(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()

	s.rwmu.Lock()
	s.lastAuthorizeQuery = cloneValues(q)
	s.rwmu.Unlock()

	redirectURI := q.Get("redirect_uri")
	state := q.Get("state")
	challenge := q.Get("code_challenge")
	challengeMethod := q.Get("code_challenge_method")
	nonce := q.Get("nonce")
	clientID := q.Get("client_id")

	// 基本驗：redirect_uri / client_id 缺則 400。
	// 真 IdP 還會驗 redirect_uri 是否在註冊白名單；fake server 簡化掉，反正測試 caller 一定會
	// 帶正確的 redirect_uri。
	if redirectURI == "" || clientID == "" {
		http.Error(w, "missing redirect_uri or client_id", http.StatusBadRequest)
		return
	}

	// 簽發 code（記 challenge / nonce 進 codeStore）
	code, err := s.IssueAuthCode(challenge, challengeMethod, nonce, redirectURI)
	if err != nil {
		http.Error(w, "issue code failed: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 組 callback URL：redirect_uri?code=<code>&state=<state>
	cbURL, err := url.Parse(redirectURI)
	if err != nil {
		http.Error(w, "invalid redirect_uri: "+err.Error(), http.StatusBadRequest)
		return
	}
	rq := cbURL.Query()
	rq.Set("code", code)
	if state != "" {
		rq.Set("state", state)
	}
	cbURL.RawQuery = rq.Encode()

	http.Redirect(w, r, cbURL.String(), http.StatusFound)
}

// handleToken 處理 POST /oauth/token（authorization_code grant）。
//
// 流程：
//  1. 驗 client_id / client_secret
//  2. 驗 grant_type == authorization_code
//  3. 從 codeStore 取出 code 對應的 challenge / nonce
//  4. 驗 PKCE：sha256(code_verifier) == challenge（base64url 比對）
//  5. 簽 id_token（用 nextIDTokenClaims 或預設 claims；nonce 自動補入）
//  6. 回 token response
func (s *Server) handleToken(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeOAuthError(w, http.StatusBadRequest, "invalid_request", "parse form: "+err.Error())
		return
	}

	// HTTP Basic auth or form fields — 兩種都支援，跟真 OIDC IdP 一致
	clientID, clientSecret := extractClientCredentials(r)

	s.rwmu.Lock()
	s.lastTokenForm = cloneValues(r.Form)
	s.rwmu.Unlock()

	// ─── 1. client credentials ───
	if clientID != s.ClientID || clientSecret != s.ClientSecret {
		writeOAuthError(w, http.StatusUnauthorized, "invalid_client", "client credentials mismatch")
		return
	}

	// ─── 2. grant type ───
	if gt := r.Form.Get("grant_type"); gt != "authorization_code" {
		writeOAuthError(w, http.StatusBadRequest, "unsupported_grant_type", "got "+gt)
		return
	}

	// ─── 3. code 取對應 metadata ───
	code := r.Form.Get("code")
	if code == "" {
		writeOAuthError(w, http.StatusBadRequest, "invalid_grant", "missing code")
		return
	}
	s.rwmu.Lock()
	meta, ok := s.codeStore[code]
	if ok {
		// 真 IdP 的 code 是「一次性」— 用過就刪，避免 replay。
		delete(s.codeStore, code)
	}
	s.rwmu.Unlock()
	if !ok {
		writeOAuthError(w, http.StatusBadRequest, "invalid_grant", "unknown or already-used code")
		return
	}

	// ─── 4. PKCE ───
	verifier := r.Form.Get("code_verifier")
	if meta.CodeChallenge != "" {
		if verifier == "" {
			writeOAuthError(w, http.StatusBadRequest, "invalid_grant", "code_verifier required (PKCE was used)")
			return
		}
		if !verifyPKCE(verifier, meta.CodeChallenge, meta.CodeChallengeMethod) {
			writeOAuthError(w, http.StatusBadRequest, "invalid_grant", "PKCE verifier mismatch")
			return
		}
	}

	// ─── 5. 簽 id_token ───
	s.rwmu.RLock()
	customClaims := cloneClaims(s.nextIDTokenClaims)
	accessToken := s.nextAccessToken
	issuer := s.Issuer
	s.rwmu.RUnlock()

	now := time.Now()
	claims := map[string]any{
		"iss":   issuer,
		"aud":   s.ClientID,
		"iat":   now.Unix(),
		"exp":   now.Add(5 * time.Minute).Unix(),
		"nbf":   now.Unix(),
		"sub":   "sub-fake-user-001",
		"email": "fake-user@example.com",
		"name":  "Fake User",
	}
	// 把 caller 指定的 claims 蓋過預設（sub/email/name 等）
	for k, v := range customClaims {
		claims[k] = v
	}
	// 永遠把 authorize 收到的 nonce 灌回去（除非 caller 已經自行指定）
	if _, has := claims["nonce"]; !has && meta.Nonce != "" {
		claims["nonce"] = meta.Nonce
	}

	idToken, err := signJWT(s.PrivateKey, s.KeyID, jose.RS256, claims)
	if err != nil {
		writeOAuthError(w, http.StatusInternalServerError, "server_error", "sign id_token: "+err.Error())
		return
	}

	resp := map[string]any{
		"access_token": accessToken,
		"token_type":   "Bearer",
		"expires_in":   3600,
		"id_token":     idToken,
	}
	writeJSON(w, http.StatusOK, resp)
}

// ───────────────────────── helpers ─────────────────────────

// extractClientCredentials 從 HTTP Basic auth 或 form field 取出 client_id / client_secret。
// 真 IdP 通常兩種都接（RFC 6749 §2.3.1），fake server 也比照辦理。
func extractClientCredentials(r *http.Request) (string, string) {
	if cid, csec, ok := r.BasicAuth(); ok && cid != "" {
		return cid, csec
	}
	return r.Form.Get("client_id"), r.Form.Get("client_secret")
}

// verifyPKCE 對照 RFC 7636 §4.6 規定驗 code_verifier 對 challenge：
//
//	S256: BASE64URL(SHA256(verifier)) == challenge
//
// 不支援 plain（OAuth 2.1 已 deprecated；fake server 也只認 S256）。
func verifyPKCE(verifier, challenge, method string) bool {
	if method == "" {
		method = "S256" // 真 IdP 的預設可能不同，但 visionA 一律用 S256
	}
	if method != "S256" {
		return false
	}
	expected := pkceS256(verifier)
	return expected == challenge
}

// pkceS256：BASE64URL(SHA256(verifier))。重新實作而不 import internal/oidc，
// 保持 oidctest 不依賴 production package（避免循環依賴 + 確保 oidctest 可被
// internal/oidc 自己未來想用而不打死）。
func pkceS256(verifier string) string {
	sum := sha256.Sum256([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(sum[:])
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

// writeOAuthError 寫 RFC 6749 §5.2 規範的 token error response。
func writeOAuthError(w http.ResponseWriter, status int, code, desc string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(map[string]string{
		"error":             code,
		"error_description": desc,
	})
}

// cloneValues 複製 url.Values，避免 caller 改到我們存的觀測值。
func cloneValues(in url.Values) url.Values {
	if in == nil {
		return nil
	}
	out := make(url.Values, len(in))
	for k, vs := range in {
		cp := make([]string, len(vs))
		copy(cp, vs)
		out[k] = cp
	}
	return out
}

// cloneClaims 複製 claims map，避免並發 race。
func cloneClaims(in map[string]any) map[string]any {
	if in == nil {
		return nil
	}
	out := make(map[string]any, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

// rsaPublicKeyToJWK 把 RSA public key 編成 JWKS spec 的 key 物件。
// 與 OB1 的 fakeOIDC 寫法一致（base64url 無 padding；exponent 手動轉 byte slice）。
func rsaPublicKeyToJWK(pub *rsa.PublicKey, kid string) map[string]any {
	return map[string]any{
		"kty": "RSA",
		"alg": "RS256",
		"use": "sig",
		"kid": kid,
		"n":   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
		"e":   base64.RawURLEncoding.EncodeToString(bigIntBytes(pub.E)),
	}
}

func bigIntBytes(e int) []byte {
	out := []byte{}
	for e > 0 {
		out = append([]byte{byte(e & 0xff)}, out...)
		e >>= 8
	}
	if len(out) == 0 {
		out = []byte{0}
	}
	return out
}

// signJWT 用 RSA private key 簽出 RS256 JWT。
//
// 接受 map[string]any 而非 jwt.Claims struct，方便 caller 灌任意 claim
// （包含 OIDC 的 sub/email/name 與測試用的非標準欄位）。
func signJWT(priv *rsa.PrivateKey, kid string, alg jose.SignatureAlgorithm, claims map[string]any) (string, error) {
	signerOpts := (&jose.SignerOptions{}).WithType("JWT")
	signerOpts.WithHeader("kid", kid)
	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: alg, Key: priv}, signerOpts)
	if err != nil {
		return "", fmt.Errorf("oidctest: new signer: %w", err)
	}
	tok, err := jwt.Signed(signer).Claims(claims).Serialize()
	if err != nil {
		return "", fmt.Errorf("oidctest: sign jwt: %w", err)
	}
	return tok, nil
}

// randomURLToken 產生 base64url 編碼的隨機 token，給 authorization code 用。
func randomURLToken(nBytes int) (string, error) {
	b := make([]byte, nBytes)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}