Kiro/internal/auth/session_store_test.go

package auth

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestCookieSessionStore(t *testing.T) {
	store := NewCookieSessionStore("test-secret-key", 3600)

	t.Run("create and retrieve session", func(t *testing.T) {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("GET", "/", nil)

		userID := "test-user-123"
		err := store.CreateSession(w, r, userID)
		if err != nil {
			t.Fatalf("Failed to create session: %v", err)
		}

		// Get the cookie from the response
		cookies := w.Result().Cookies()
		if len(cookies) == 0 {
			t.Fatal("Expected session cookie, got none")
		}

		// Create a new request with the cookie
		r2 := httptest.NewRequest("GET", "/", nil)
		for _, cookie := range cookies {
			r2.AddCookie(cookie)
		}

		retrievedUserID, err := store.GetUserID(r2)
		if err != nil {
			t.Fatalf("Failed to get user ID: %v", err)
		}

		if retrievedUserID != userID {
			t.Errorf("Expected user ID %s, got %s", userID, retrievedUserID)
		}
	})

	t.Run("get user ID without session", func(t *testing.T) {
		r := httptest.NewRequest("GET", "/", nil)

		_, err := store.GetUserID(r)
		if err == nil {
			t.Error("Expected error when getting user ID without session")
		}
	})

	t.Run("destroy session", func(t *testing.T) {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("GET", "/", nil)

		userID := "test-user-456"
		err := store.CreateSession(w, r, userID)
		if err != nil {
			t.Fatalf("Failed to create session: %v", err)
		}

		// Get the cookie
		cookies := w.Result().Cookies()
		r2 := httptest.NewRequest("GET", "/", nil)
		for _, cookie := range cookies {
			r2.AddCookie(cookie)
		}

		// Destroy the session
		w2 := httptest.NewRecorder()
		err = store.DestroySession(w2, r2)
		if err != nil {
			t.Fatalf("Failed to destroy session: %v", err)
		}

		// Check that the cookie has MaxAge set to -1 (deletion marker)
		destroyCookies := w2.Result().Cookies()
		if len(destroyCookies) == 0 {
			t.Fatal("Expected cookie with MaxAge=-1 for deletion")
		}

		foundDeleteCookie := false
		for _, cookie := range destroyCookies {
			if cookie.MaxAge == -1 {
				foundDeleteCookie = true
				break
			}
		}

		if !foundDeleteCookie {
			t.Error("Expected cookie with MaxAge=-1 to indicate deletion")
		}
	})

	t.Run("update existing session", func(t *testing.T) {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("GET", "/", nil)

		// Create initial session
		userID1 := "user-1"
		store.CreateSession(w, r, userID1)

		cookies := w.Result().Cookies()
		r2 := httptest.NewRequest("GET", "/", nil)
		for _, cookie := range cookies {
			r2.AddCookie(cookie)
		}

		// Update session with new user ID
		w2 := httptest.NewRecorder()
		userID2 := "user-2"
		err := store.CreateSession(w2, r2, userID2)
		if err != nil {
			t.Fatalf("Failed to update session: %v", err)
		}

		// Verify new user ID
		cookies2 := w2.Result().Cookies()
		r3 := httptest.NewRequest("GET", "/", nil)
		for _, cookie := range cookies2 {
			r3.AddCookie(cookie)
		}

		retrievedUserID, err := store.GetUserID(r3)
		if err != nil {
			t.Fatalf("Failed to get user ID: %v", err)
		}

		if retrievedUserID != userID2 {
			t.Errorf("Expected user ID %s, got %s", userID2, retrievedUserID)
		}
	})

	t.Run("validate session with valid session", func(t *testing.T) {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("GET", "/", nil)

		userID := "test-user-789"
		err := store.CreateSession(w, r, userID)
		if err != nil {
			t.Fatalf("Failed to create session: %v", err)
		}

		// Create request with session cookie
		cookies := w.Result().Cookies()
		r2 := httptest.NewRequest("GET", "/", nil)
		for _, cookie := range cookies {
			r2.AddCookie(cookie)
		}

		// Validate session
		if !store.ValidateSession(r2) {
			t.Error("Expected ValidateSession to return true for valid session")
		}
	})

	t.Run("validate session without session", func(t *testing.T) {
		r := httptest.NewRequest("GET", "/", nil)

		// Validate session without any cookies
		if store.ValidateSession(r) {
			t.Error("Expected ValidateSession to return false for missing session")
		}
	})

	t.Run("validate session after logout", func(t *testing.T) {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("GET", "/", nil)

		userID := "test-user-logout"
		err := store.CreateSession(w, r, userID)
		if err != nil {
			t.Fatalf("Failed to create session: %v", err)
		}

		// Get the cookie
		cookies := w.Result().Cookies()
		r2 := httptest.NewRequest("GET", "/", nil)
		for _, cookie := range cookies {
			r2.AddCookie(cookie)
		}

		// Destroy the session
		w2 := httptest.NewRecorder()
		err = store.DestroySession(w2, r2)
		if err != nil {
			t.Fatalf("Failed to destroy session: %v", err)
		}

		// Create a new request without any cookies (simulating browser behavior after logout)
		r3 := httptest.NewRequest("GET", "/", nil)

		// Validate session should return false
		if store.ValidateSession(r3) {
			t.Error("Expected ValidateSession to return false after logout")
		}
	})

	t.Run("session cookie has security settings", func(t *testing.T) {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("GET", "/", nil)

		userID := "test-user-security"
		err := store.CreateSession(w, r, userID)
		if err != nil {
			t.Fatalf("Failed to create session: %v", err)
		}

		cookies := w.Result().Cookies()
		if len(cookies) == 0 {
			t.Fatal("Expected session cookie, got none")
		}

		cookie := cookies[0]

		// Verify HttpOnly flag is set
		if !cookie.HttpOnly {
			t.Error("Expected HttpOnly flag to be true")
		}

		// Verify SameSite is set
		if cookie.SameSite != http.SameSiteLaxMode {
			t.Errorf("Expected SameSite to be Lax, got %v", cookie.SameSite)
		}

		// Verify Path is set
		if cookie.Path != "/" {
			t.Errorf("Expected Path to be /, got %s", cookie.Path)
		}

		// Verify MaxAge is set
		if cookie.MaxAge != 3600 {
			t.Errorf("Expected MaxAge to be 3600, got %d", cookie.MaxAge)
		}
	})
}