visionA/visionA-backend/internal/usersession/manager_test.go

package usersession

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

func newManager() *Manager {
	return NewManager(NewInMemoryStore(), newCookieCfg())
}

func TestManager_StartSession_WritesCookieAndStoresSession(t *testing.T) {
	m := newManager()
	w := httptest.NewRecorder()

	sess, err := m.StartSession(context.Background(), w)
	if err != nil {
		t.Fatalf("StartSession: %v", err)
	}
	if sess.ID == "" {
		t.Fatalf("session ID empty")
	}

	// 應該寫了一個 cookie
	cookies := w.Result().Cookies()
	if len(cookies) != 1 {
		t.Fatalf("expected 1 cookie, got %d", len(cookies))
	}
	c := cookies[0]
	if c.Name != DefaultCookieName {
		t.Fatalf("cookie name mismatch: %q", c.Name)
	}
	// cookie value 應該能解出 sessionID
	got, err := DecodeCookieValue(c.Value, testKey)
	if err != nil {
		t.Fatalf("DecodeCookieValue: %v", err)
	}
	if got != sess.ID {
		t.Fatalf("decoded sid mismatch: got %q want %q", got, sess.ID)
	}
}

// emptyIDStore 是 stub Store：Create() 回傳 sessionID 為空字串的 Session，
// 並紀錄所有 Delete 呼叫。
//
// 用途：NewManager 已強制 SigningKey ≥ 32 bytes，原本「無 SigningKey 觸發 WriteCookie
// 失敗」的路徑不再可達。此 stub 改用「空 sessionID」觸發 WriteCookie 的 ErrInvalidCookie 路徑，
// 確保 StartSession 的 rollback 邏輯（建 session 失敗 → 主動 Delete）仍被測到。
type emptyIDStore struct {
	createdIDs []string // 真實 Store 會給的 ID（非空）
	deleted    []string // 被 Delete 的 ID
}

func (s *emptyIDStore) Create(ctx context.Context) (*Session, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	// 對外回傳的 Session.ID 為空，觸發 Manager.StartSession → WriteCookie 失敗路徑。
	// 但內部記錄一個「假裝建立過」的 ID 供驗證 rollback。
	s.createdIDs = append(s.createdIDs, "fake-id-rollback")
	return &Session{ID: ""}, nil
}

func (s *emptyIDStore) Get(_ context.Context, _ string) (*Session, error) {
	return nil, ErrNoSession
}

func (s *emptyIDStore) Update(_ context.Context, _ *Session) error { return nil }

func (s *emptyIDStore) Delete(_ context.Context, id string) error {
	s.deleted = append(s.deleted, id)
	return nil
}

func (s *emptyIDStore) CleanupExpired(_ context.Context, _, _ time.Duration) (int, error) {
	return 0, nil
}

func TestManager_StartSession_RollbackOnCookieFailure(t *testing.T) {
	// 觸發 WriteCookie 失敗（sessionID 空）→ Manager 必須呼叫 Store.Delete 做 rollback，
	// 即使呼叫對象 ID 是空字串（stub Store 仍記錄到該次呼叫）。
	store := &emptyIDStore{}
	m := NewManager(store, newCookieCfg())
	w := httptest.NewRecorder()

	_, err := m.StartSession(context.Background(), w)
	if !errors.Is(err, ErrInvalidCookie) {
		t.Fatalf("expected ErrInvalidCookie (empty sessionID → WriteCookie fail), got %v", err)
	}
	if len(store.deleted) != 1 {
		t.Fatalf("StartSession failure must rollback (call store.Delete exactly once); deleted=%v", store.deleted)
	}
}

func TestManager_GetSession_Roundtrip(t *testing.T) {
	m := newManager()
	w := httptest.NewRecorder()
	sess, err := m.StartSession(context.Background(), w)
	if err != nil {
		t.Fatalf("StartSession: %v", err)
	}

	// 把 Set-Cookie 轉成 request cookie
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(w.Result().Cookies()[0])

	got, err := m.GetSession(context.Background(), r)
	if err != nil {
		t.Fatalf("GetSession: %v", err)
	}
	if got.ID != sess.ID {
		t.Fatalf("ID mismatch: got %q want %q", got.ID, sess.ID)
	}
}

func TestManager_GetSession_NoCookie(t *testing.T) {
	m := newManager()
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	_, err := m.GetSession(context.Background(), r)
	if !errors.Is(err, ErrNoSession) {
		t.Fatalf("expected ErrNoSession, got %v", err)
	}
}

func TestManager_GetSession_TamperedCookie(t *testing.T) {
	m := newManager()
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(&http.Cookie{Name: DefaultCookieName, Value: "fake.sig"})
	_, err := m.GetSession(context.Background(), r)
	if !errors.Is(err, ErrNoSession) {
		t.Fatalf("expected ErrNoSession for tampered cookie, got %v", err)
	}
}

func TestManager_GetSession_DeletedSession(t *testing.T) {
	m := newManager()

	// Start → 拿到 cookie
	w := httptest.NewRecorder()
	sess, _ := m.StartSession(context.Background(), w)
	cookie := w.Result().Cookies()[0]

	// 從 store 直接刪掉
	if err := m.Store.Delete(context.Background(), sess.ID); err != nil {
		t.Fatalf("Delete: %v", err)
	}

	// 此時拿著 cookie 來 Get → ErrNoSession
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(cookie)
	_, err := m.GetSession(context.Background(), r)
	if !errors.Is(err, ErrNoSession) {
		t.Fatalf("expected ErrNoSession after Delete, got %v", err)
	}
}

func TestManager_UpdateSession(t *testing.T) {
	m := newManager()
	w := httptest.NewRecorder()
	sess, _ := m.StartSession(context.Background(), w)

	sess.UserID = "u-1"
	sess.Email = "alice@example.com"
	if err := m.UpdateSession(context.Background(), sess); err != nil {
		t.Fatalf("UpdateSession: %v", err)
	}

	got, err := m.Store.Get(context.Background(), sess.ID)
	if err != nil {
		t.Fatalf("Get: %v", err)
	}
	if got.UserID != "u-1" || got.Email != "alice@example.com" {
		t.Fatalf("UpdateSession did not persist: %+v", got)
	}
}

func TestManager_UpdateSession_Nil(t *testing.T) {
	m := newManager()
	if err := m.UpdateSession(context.Background(), nil); !errors.Is(err, ErrNoSession) {
		t.Fatalf("expected ErrNoSession for nil session, got %v", err)
	}
}

func TestManager_EndSession_DeletesAndClearsCookie(t *testing.T) {
	m := newManager()

	// 先建立 session 拿 cookie
	w1 := httptest.NewRecorder()
	sess, _ := m.StartSession(context.Background(), w1)
	cookie := w1.Result().Cookies()[0]

	// EndSession：帶著 cookie 進來
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(cookie)
	w2 := httptest.NewRecorder()
	if err := m.EndSession(context.Background(), w2, r); err != nil {
		t.Fatalf("EndSession: %v", err)
	}

	// store 應該找不到
	if _, err := m.Store.Get(context.Background(), sess.ID); !errors.Is(err, ErrNoSession) {
		t.Fatalf("session should be gone, got %v", err)
	}

	// 應該寫了一個 expiration cookie
	cookies := w2.Result().Cookies()
	if len(cookies) != 1 || cookies[0].MaxAge >= 0 {
		t.Fatalf("EndSession should write expiration cookie, got %+v", cookies)
	}

	// 後續 GetSession 用同一個 cookie → ErrNoSession
	r2 := httptest.NewRequest(http.MethodGet, "/", nil)
	r2.AddCookie(cookie)
	if _, err := m.GetSession(context.Background(), r2); !errors.Is(err, ErrNoSession) {
		t.Fatalf("after EndSession, GetSession should fail, got %v", err)
	}
}

// TestNewManager_RejectsShortSigningKey 確認 startup-time 強制 SigningKey ≥ 32 bytes。
//
// 用 panic 而非 return error，目的是讓任何 misconfig 在啟動瞬間就掛掉，
// 而不是讓 cookie 被簽出去之後才在 logs 裡出現 entropy 警告。
func TestNewManager_RejectsShortSigningKey(t *testing.T) {
	cases := map[string][]byte{
		"nil":         nil,
		"empty":       {},
		"too_short_1": []byte("x"),
		"31_bytes":    []byte("0123456789012345678901234567890"), // 31 bytes
	}
	for name, key := range cases {
		t.Run(name, func(t *testing.T) {
			cfg := newCookieCfg()
			cfg.SigningKey = key
			defer func() {
				r := recover()
				if r == nil {
					t.Fatalf("expected NewManager to panic for SigningKey len=%d", len(key))
				}
			}()
			_ = NewManager(NewInMemoryStore(), cfg)
		})
	}
}

// TestNewManager_AcceptsExactly32Bytes 邊界值測試：剛好 32 bytes 應通過。
func TestNewManager_AcceptsExactly32Bytes(t *testing.T) {
	cfg := newCookieCfg()
	cfg.SigningKey = []byte("01234567890123456789012345678901") // 正好 32 bytes
	defer func() {
		if r := recover(); r != nil {
			t.Fatalf("NewManager should accept 32-byte key, got panic: %v", r)
		}
	}()
	m := NewManager(NewInMemoryStore(), cfg)
	if m == nil {
		t.Fatalf("NewManager returned nil")
	}
}

func TestManager_EndSession_NoCookie_StillWritesClearCookie(t *testing.T) {
	// logout 應該是冪等的：即使 request 沒帶 cookie，也應寫一個 expiration cookie 確保 browser 端清乾淨
	m := newManager()
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	w := httptest.NewRecorder()
	if err := m.EndSession(context.Background(), w, r); err != nil {
		t.Fatalf("EndSession with no cookie should still succeed, got %v", err)
	}
	cookies := w.Result().Cookies()
	if len(cookies) != 1 || cookies[0].MaxAge >= 0 {
		t.Fatalf("EndSession (no cookie) should still write expiration cookie")
	}
}

// ---------------------------------------------------------------------
// Fix-A1：RotateSessionID（session fixation 防護）測試
// ---------------------------------------------------------------------

// TestManager_RotateSessionID_HappyPath 驗證：
//   - Rotate 後舊 session ID 在 store 取不到
//   - 新 session ID 能在 store 取到
//   - Cookie 真的被改寫成新的（與舊 cookie value 不同）
//   - 舊 session 上的所有欄位（OIDC pending state + Extra）都複製到新 session
func TestManager_RotateSessionID_HappyPath(t *testing.T) {
	m := newManager()

	// 建立 pending session 並塞滿可能的欄位
	w1 := httptest.NewRecorder()
	oldSess, err := m.StartSession(context.Background(), w1)
	if err != nil {
		t.Fatalf("StartSession: %v", err)
	}
	oldSess.OIDCState = "state-xyz"
	oldSess.OIDCNonce = "nonce-xyz"
	oldSess.OIDCCodeVerifier = "verifier-xyz"
	oldSess.Extra = map[string]any{"return_to": "/dashboard"}
	if err := m.UpdateSession(context.Background(), oldSess); err != nil {
		t.Fatalf("UpdateSession: %v", err)
	}
	oldCookie := w1.Result().Cookies()[0]
	oldCookieValue := oldCookie.Value

	// Rotate：模擬「callback 驗 id_token 成功後」
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(oldCookie)
	w2 := httptest.NewRecorder()
	newSess, err := m.RotateSessionID(context.Background(), w2, r)
	if err != nil {
		t.Fatalf("RotateSessionID: %v", err)
	}

	// 新 session ID 必須與舊不同
	if newSess.ID == oldSess.ID {
		t.Fatalf("new session ID should differ from old; got same %q", newSess.ID)
	}

	// 舊 session 在 store 應已被刪
	if _, err := m.Store.Get(context.Background(), oldSess.ID); !errors.Is(err, ErrNoSession) {
		t.Fatalf("old session should be deleted; Get returned %v", err)
	}

	// 新 session 在 store 應存在
	got, err := m.Store.Get(context.Background(), newSess.ID)
	if err != nil {
		t.Fatalf("new session should exist in store; Get returned %v", err)
	}

	// 欄位都應複製過去
	if got.OIDCState != "state-xyz" || got.OIDCNonce != "nonce-xyz" || got.OIDCCodeVerifier != "verifier-xyz" {
		t.Fatalf("OIDC pending fields not copied: %+v", got)
	}
	if v, ok := got.Extra["return_to"]; !ok || v != "/dashboard" {
		t.Fatalf("Extra[return_to] not copied; Extra=%+v", got.Extra)
	}

	// 確保 Extra 是「複本」而非共享 map：改新 session 的 Extra 不影響舊 session 殘片
	got.Extra["return_to"] = "/changed"
	if v := oldSess.Extra["return_to"]; v != "/dashboard" {
		t.Fatalf("Extra map was shared (mutation leaked back to oldSess.Extra=%v)", v)
	}

	// Cookie 應被改寫
	newCookies := w2.Result().Cookies()
	if len(newCookies) != 1 {
		t.Fatalf("expected 1 new cookie, got %d", len(newCookies))
	}
	if newCookies[0].Value == oldCookieValue {
		t.Fatalf("new cookie value should differ from old; got same value")
	}

	// 用新 cookie 應能取到新 session
	r2 := httptest.NewRequest(http.MethodGet, "/", nil)
	r2.AddCookie(newCookies[0])
	got2, err := m.GetSession(context.Background(), r2)
	if err != nil {
		t.Fatalf("GetSession with new cookie: %v", err)
	}
	if got2.ID != newSess.ID {
		t.Fatalf("GetSession returned ID %q, want %q", got2.ID, newSess.ID)
	}

	// 用舊 cookie 應拿不到（session 已刪）
	r3 := httptest.NewRequest(http.MethodGet, "/", nil)
	r3.AddCookie(oldCookie)
	if _, err := m.GetSession(context.Background(), r3); !errors.Is(err, ErrNoSession) {
		t.Fatalf("GetSession with old cookie should return ErrNoSession, got %v", err)
	}
}

// TestManager_RotateSessionID_NoCookie 驗證：沒 cookie → ErrNoSession
func TestManager_RotateSessionID_NoCookie(t *testing.T) {
	m := newManager()
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	w := httptest.NewRecorder()
	_, err := m.RotateSessionID(context.Background(), w, r)
	if !errors.Is(err, ErrNoSession) {
		t.Fatalf("expected ErrNoSession, got %v", err)
	}
}

// TestManager_RotateSessionID_OldSessionDeletedFromStore 驗證：cookie 有效但 store 中
// 找不到 session（store 重啟、或 race 中被刪）→ ErrNoSession
func TestManager_RotateSessionID_OldSessionDeletedFromStore(t *testing.T) {
	m := newManager()
	w1 := httptest.NewRecorder()
	oldSess, _ := m.StartSession(context.Background(), w1)
	cookie := w1.Result().Cookies()[0]

	// 偷偷把 store 中的 session 刪掉（模擬 race）
	_ = m.Store.Delete(context.Background(), oldSess.ID)

	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(cookie)
	w2 := httptest.NewRecorder()
	_, err := m.RotateSessionID(context.Background(), w2, r)
	if !errors.Is(err, ErrNoSession) {
		t.Fatalf("expected ErrNoSession when old session missing from store, got %v", err)
	}
	// 不應寫新 cookie
	if cookies := w2.Result().Cookies(); len(cookies) != 0 {
		t.Fatalf("should not write cookie when rotation fails, got %+v", cookies)
	}
}

// TestManager_RotateSessionID_PreservesUserFields 驗證已升級為 logged-in 後再 rotate
// 也會保留 UserID / Email / Name。
//
// 雖然正式流程中 rotate 發生在「填 user info 之前」，但 RotateSessionID 本身應該對
// 任何 session 狀態都安全（不假設 UserID 為空）。
func TestManager_RotateSessionID_PreservesUserFields(t *testing.T) {
	m := newManager()
	w1 := httptest.NewRecorder()
	oldSess, _ := m.StartSession(context.Background(), w1)
	oldSess.UserID = "u-123"
	oldSess.Email = "alice@example.com"
	oldSess.Name = "Alice"
	oldSess.AccessToken = "secret-token"
	oldSess.IDTokenRaw = "raw.jwt.here"
	if err := m.UpdateSession(context.Background(), oldSess); err != nil {
		t.Fatalf("UpdateSession: %v", err)
	}
	cookie := w1.Result().Cookies()[0]

	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.AddCookie(cookie)
	w2 := httptest.NewRecorder()
	newSess, err := m.RotateSessionID(context.Background(), w2, r)
	if err != nil {
		t.Fatalf("RotateSessionID: %v", err)
	}

	if newSess.UserID != "u-123" || newSess.Email != "alice@example.com" || newSess.Name != "Alice" {
		t.Fatalf("user fields lost: %+v", newSess)
	}
	if newSess.AccessToken != "secret-token" || newSess.IDTokenRaw != "raw.jwt.here" {
		t.Fatalf("token snapshot lost: %+v", newSess)
	}
}

// TestManager_RotateSessionID_StoreUpdateFailureRollsBack 驗證：Update 失敗時新建的 session 會被清掉。
//
// 用 stub Store：Create 成功但 Update 第一次回 error。Rotate 應 rollback 並回 error。
func TestManager_RotateSessionID_StoreUpdateFailureRollsBack(t *testing.T) {
	// 建一個會：Create 成功、Get 回有效 session、Update 永遠失敗的 store
	failStore := &updateFailingStore{
		inner: NewInMemoryStore(),
	}
	// 先用內部 store 真的 Create 一個 session 以模擬 pending session
	pending, err := failStore.inner.Create(context.Background())
	if err != nil {
		t.Fatalf("inner Create: %v", err)
	}

	m := NewManager(failStore, newCookieCfg())

	// 假裝 cookie 是這個 pending session 的
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	cookieValue := EncodeCookieValue(pending.ID, testKey)
	r.AddCookie(&http.Cookie{Name: DefaultCookieName, Value: cookieValue})

	w := httptest.NewRecorder()

	// 啟用 update 失敗模式
	failStore.failUpdate = true

	_, err = m.RotateSessionID(context.Background(), w, r)
	if err == nil {
		t.Fatalf("expected error when Update fails")
	}

	// rollback：新 session 不應殘留在 store 中。
	// 我們無法直接知道新 session ID（rotation 失敗了），但可以驗：
	//   - 舊 session 仍在（rotation 沒提早刪它）
	//   - store 內部除了「pending 那個」之外不該多出新 session
	if got := failStore.inner.Len(); got != 1 {
		t.Fatalf("rollback failed: store should have only the original pending session, got %d entries", got)
	}
}

// updateFailingStore 包裝一個正常 Store，但允許測試切換 Update 行為到 fail。
type updateFailingStore struct {
	inner      *InMemoryStore
	failUpdate bool
}

func (s *updateFailingStore) Create(ctx context.Context) (*Session, error) {
	return s.inner.Create(ctx)
}
func (s *updateFailingStore) Get(ctx context.Context, id string) (*Session, error) {
	return s.inner.Get(ctx, id)
}
func (s *updateFailingStore) Update(ctx context.Context, sess *Session) error {
	if s.failUpdate {
		return errors.New("simulated store update failure")
	}
	return s.inner.Update(ctx, sess)
}
func (s *updateFailingStore) Delete(ctx context.Context, id string) error {
	return s.inner.Delete(ctx, id)
}
func (s *updateFailingStore) CleanupExpired(ctx context.Context, idle, abs time.Duration) (int, error) {
	return s.inner.CleanupExpired(ctx, idle, abs)
}