RenewCMS/test/stale_token_invalidation_http_test.go

package test

import (
	"RenewCMS/api/handlers/article"
	"RenewCMS/api/handlers/auth"
	"RenewCMS/api/middleware"
	"RenewCMS/api/pkg/apperr"
	validator "RenewCMS/api/pkg/validator"
	"RenewCMS/internal/domain/user"
	"RenewCMS/internal/infrastructure/useCases"
	"bytes"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/jwtauth/v5"
)

func newStaleTokenInvalidationRouter(userRepo *mockUserRepo, articleRepo *mockArticleRepo, imageRepo *mockImageRepo, mailRepo *mockMailRepo, hasher *stubHasher, token *jwtauth.JWTAuth) *chi.Mux {
	getUserUseCase := useCases.NewGetUserUseCase(userRepo)
	articleHandler := article.NewHandler(
		useCases.NewCreateArticleUseCase(articleRepo, &fakeBodyRenderer{}, discardRecorder, plainSites(), plainAuthors()),
		useCases.NewGetArticleUseCase(articleRepo, &fakeBodyRenderer{}),
		useCases.NewListArticlesUseCase(articleRepo, &fakeBodyRenderer{}, plainSites()),
		useCases.NewDeleteArticleUseCase(articleRepo, discardRecorder),
		useCases.NewUpdateArticleUseCase(articleRepo, &fakeBodyRenderer{}, discardRecorder, plainSites(), plainAuthors()),
		getUserUseCase,
	)
	loginHandler := auth.NewLoginHandler(
		token,
		useCases.NewCreateUserUseCase(userRepo, hasher, mailRepo, "https://cms.example.com", discardRecorder),
		useCases.NewAuthenticateUserUseCase(userRepo, hasher, discardRecorder),
		true,
	)
	verificationHandler := auth.NewVerificationHandler(
		useCases.NewVerifyUserUseCase(userRepo, hasher, discardRecorder),
		useCases.NewCancelRegistrationUseCase(userRepo),
		true,
	)
	passwordResetHandler := auth.NewPasswordResetHandler(
		useCases.NewRequestPasswordResetUseCase(userRepo, hasher, mailRepo, "https://cms.example.com", discardRecorder),
		useCases.NewResetPasswordUseCase(userRepo, hasher, newMockAPIKeyRepo(), time.Now, discardRecorder),
	)

	r := chi.NewRouter()
	r.Route("/v1", func(r chi.Router) {
		r.Route("/article", func(r chi.Router) {
			r.Group(func(r chi.Router) {
				r.Use(jwtauth.Verifier(token))
				r.Use(jwtauth.Authenticator(token))
				r.Use(middleware.RequireCurrentToken(getUserUseCase))
				r.Use(middleware.RequireVerifiedUser(getUserUseCase))
				r.Post("/", articleHandler.PostArticle)
			})
		})
		r.Route("/auth", func(r chi.Router) {
			r.Post("/register", loginHandler.Register)
			r.Post("/login", loginHandler.Login)
			r.Post("/password-reset", passwordResetHandler.RequestPasswordReset)
			r.Post("/password-reset/confirm", passwordResetHandler.ResetPassword)
			r.Group(func(r chi.Router) {
				r.Use(jwtauth.Verifier(token))
				r.Use(jwtauth.Authenticator(token))
				r.Use(middleware.RequireCurrentToken(getUserUseCase))
				r.Post("/verify", verificationHandler.Verify)
				r.Delete("/register", verificationHandler.CancelRegistration)
			})
		})
	})
	return r
}

func newStaleTokenAliceRepo(isVerified bool) *mockUserRepo {
	repo := newMockUserRepo()
	repo.users[7] = user.User{
		ID:                      7,
		Username:                "alice",
		Email:                   "alice@example.com",
		Password:                "hashed:oldpassword1",
		IsVerified:              isVerified,
		TokenVersion:            0,
		PasswordResetCode:       "hashed:reset-code-1",
		PasswordResetExpiration: time.Now().Add(time.Hour),
	}
	repo.nextID = 7
	return repo
}

func mintStaleTokenCookie(t *testing.T, token *jwtauth.JWTAuth, claims map[string]any) *http.Cookie {
	t.Helper()
	_, tokenString, err := token.Encode(claims)
	if err != nil {
		t.Fatalf("failed to encode jwt: %v", err)
	}
	return &http.Cookie{Name: "jwt", Value: tokenString}
}

func sendStaleTokenRequest(router *chi.Mux, method, target string, body string, cookie *http.Cookie) *httptest.ResponseRecorder {
	var req *http.Request
	if body == "" {
		req = httptest.NewRequest(method, target, nil)
	} else {
		req = httptest.NewRequest(method, target, bytes.NewBufferString(body))
		req.Header.Set("Content-Type", "application/json")
	}
	if cookie != nil {
		req.AddCookie(cookie)
	}
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)
	return rec
}

func decodeJwtPayload(t *testing.T, tokenString string) map[string]any {
	t.Helper()
	parts := bytes.Split([]byte(tokenString), []byte("."))
	if len(parts) != 3 {
		t.Fatalf("expected a 3-part jwt, got %q", tokenString)
	}
	payload, err := base64.RawURLEncoding.DecodeString(string(parts[1]))
	if err != nil {
		t.Fatalf("failed to base64-decode jwt payload: %v", err)
	}
	var claims map[string]any
	if err := json.Unmarshal(payload, &claims); err != nil {
		t.Fatalf("failed to unmarshal jwt payload: %v", err)
	}
	return claims
}

func jwtCookieFromResponse(t *testing.T, rec *httptest.ResponseRecorder) *http.Cookie {
	t.Helper()
	for _, value := range rec.Result().Header.Values("Set-Cookie") {
		cookie, err := http.ParseSetCookie(value)
		if err != nil {
			t.Fatalf("failed to parse Set-Cookie header %q: %v", value, err)
		}
		if cookie.Name == "jwt" {
			return cookie
		}
	}
	t.Fatalf("expected a jwt Set-Cookie header")
	return nil
}

func tokenVersionClaim(t *testing.T, claims map[string]any) float64 {
	t.Helper()
	raw, ok := claims["token_version"]
	if !ok {
		t.Fatalf("expected a token_version claim, got %+v", claims)
	}
	value, ok := raw.(float64)
	if !ok {
		t.Fatalf("expected token_version to be numeric, got %T", raw)
	}
	return value
}

func TestStaleTokenInvalidationHttp(t *testing.T) {
	validator.InitValidator()

	t.Run("A jwt minted before a password reset is refused afterward, but a fresh login's jwt is accepted", func(t *testing.T) {
		token := jwtauth.New("HS256", []byte("test-secret"), nil)
		userRepo := newStaleTokenAliceRepo(true)
		articleRepo := newMockArticleRepo()
		router := newStaleTokenInvalidationRouter(userRepo, articleRepo, newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
		preResetCookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(7), "token_version": uint32(0)})

		first := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, preResetCookie)
		if first.Code != http.StatusCreated {
			t.Fatalf("expected status 201, got %d: %s", first.Code, first.Body.String())
		}

		resetRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/password-reset/confirm", `{"email": "alice@example.com", "code": "reset-code-1", "password": "newpassword1"}`, nil)
		if resetRec.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d: %s", resetRec.Code, resetRec.Body.String())
		}
		if userRepo.users[7].TokenVersion != 1 {
			t.Fatalf("expected TokenVersion 1, got %d", userRepo.users[7].TokenVersion)
		}

		articlesBefore := len(articleRepo.articles)
		second := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "Another title", "body": "Some other body"}`, preResetCookie)
		if second.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", second.Code, second.Body.String())
		}
		if second.Body.String() != apperr.UnauthorizedErrorMessage+"\n" {
			t.Errorf("unexpected body: %q", second.Body.String())
		}
		if len(articleRepo.articles) != articlesBefore {
			t.Errorf("expected no article to be created by the stale request")
		}

		loginRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/login", `{"username": "alice", "password": "newpassword1"}`, nil)
		if loginRec.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d: %s", loginRec.Code, loginRec.Body.String())
		}
		newCookie := jwtCookieFromResponse(t, loginRec)
		if got := tokenVersionClaim(t, decodeJwtPayload(t, newCookie.Value)); got != 1 {
			t.Errorf("expected token_version 1, got %v", got)
		}

		third := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "Another title", "body": "Some other body"}`, newCookie)
		if third.Code != http.StatusCreated {
			t.Fatalf("expected status 201, got %d: %s", third.Code, third.Body.String())
		}
	})

	t.Run("A stale jwt is also refused on the auth group's own protected routes, not just article routes", func(t *testing.T) {
		token := jwtauth.New("HS256", []byte("test-secret"), nil)
		userRepo := newStaleTokenAliceRepo(false)
		router := newStaleTokenInvalidationRouter(userRepo, newMockArticleRepo(), newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
		preResetCookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(7), "token_version": uint32(0)})

		resetRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/password-reset/confirm", `{"email": "alice@example.com", "code": "reset-code-1", "password": "newpassword1"}`, nil)
		if resetRec.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d: %s", resetRec.Code, resetRec.Body.String())
		}

		verifyRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/verify", `{"code": "anything"}`, preResetCookie)
		if verifyRec.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", verifyRec.Code, verifyRec.Body.String())
		}
		if verifyRec.Body.String() != apperr.UnauthorizedErrorMessage+"\n" {
			t.Errorf("unexpected body: %q", verifyRec.Body.String())
		}

		cancelRec := sendStaleTokenRequest(router, http.MethodDelete, "/v1/auth/register", "", preResetCookie)
		if cancelRec.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", cancelRec.Code, cancelRec.Body.String())
		}
		if cancelRec.Body.String() != apperr.UnauthorizedErrorMessage+"\n" {
			t.Errorf("unexpected body: %q", cancelRec.Body.String())
		}
		if _, err := userRepo.Get(7); err != nil {
			t.Errorf("expected user 7 to still exist: %v", err)
		}
	})

	t.Run("Resetting the password twice invalidates both the original jwt and the jwt from the interim login", func(t *testing.T) {
		token := jwtauth.New("HS256", []byte("test-secret"), nil)
		userRepo := newStaleTokenAliceRepo(true)
		articleRepo := newMockArticleRepo()
		router := newStaleTokenInvalidationRouter(userRepo, articleRepo, newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
		originalCookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(7), "token_version": uint32(0)})

		firstResetRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/password-reset/confirm", `{"email": "alice@example.com", "code": "reset-code-1", "password": "newpassword1"}`, nil)
		if firstResetRec.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d: %s", firstResetRec.Code, firstResetRec.Body.String())
		}
		if userRepo.users[7].TokenVersion != 1 {
			t.Fatalf("expected TokenVersion 1, got %d", userRepo.users[7].TokenVersion)
		}

		interimLoginRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/login", `{"username": "alice", "password": "newpassword1"}`, nil)
		interimCookie := jwtCookieFromResponse(t, interimLoginRec)
		if got := tokenVersionClaim(t, decodeJwtPayload(t, interimCookie.Value)); got != 1 {
			t.Errorf("expected interim token_version 1, got %v", got)
		}

		updated := userRepo.users[7]
		updated.PasswordResetCode = "hashed:reset-code-2"
		updated.PasswordResetExpiration = time.Now().Add(time.Hour)
		userRepo.users[7] = updated

		secondResetRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/password-reset/confirm", `{"email": "alice@example.com", "code": "reset-code-2", "password": "newerpassword1"}`, nil)
		if secondResetRec.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d: %s", secondResetRec.Code, secondResetRec.Body.String())
		}
		if userRepo.users[7].TokenVersion != 2 {
			t.Fatalf("expected TokenVersion 2, got %d", userRepo.users[7].TokenVersion)
		}

		originalArticleRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, originalCookie)
		if originalArticleRec.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", originalArticleRec.Code, originalArticleRec.Body.String())
		}

		interimArticleRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, interimCookie)
		if interimArticleRec.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", interimArticleRec.Code, interimArticleRec.Body.String())
		}

		finalLoginRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/login", `{"username": "alice", "password": "newerpassword1"}`, nil)
		finalCookie := jwtCookieFromResponse(t, finalLoginRec)
		if got := tokenVersionClaim(t, decodeJwtPayload(t, finalCookie.Value)); got != 2 {
			t.Errorf("expected final token_version 2, got %v", got)
		}

		finalArticleRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, finalCookie)
		if finalArticleRec.Code != http.StatusCreated {
			t.Fatalf("expected status 201, got %d: %s", finalArticleRec.Code, finalArticleRec.Body.String())
		}
	})

	t.Run("Resetting one user's password never invalidates another user's jwt", func(t *testing.T) {
		token := jwtauth.New("HS256", []byte("test-secret"), nil)
		userRepo := newStaleTokenAliceRepo(true)
		userRepo.users[9] = user.User{ID: 9, Username: "bob", IsVerified: true, TokenVersion: 0}
		articleRepo := newMockArticleRepo()
		router := newStaleTokenInvalidationRouter(userRepo, articleRepo, newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
		bobCookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(9), "token_version": uint32(0)})

		resetRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/auth/password-reset/confirm", `{"email": "alice@example.com", "code": "reset-code-1", "password": "newpassword1"}`, nil)
		if resetRec.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d: %s", resetRec.Code, resetRec.Body.String())
		}
		if userRepo.users[9].TokenVersion != 0 {
			t.Fatalf("expected bob's TokenVersion to stay 0, got %d", userRepo.users[9].TokenVersion)
		}

		bobArticleRec := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, bobCookie)
		if bobArticleRec.Code != http.StatusCreated {
			t.Fatalf("expected status 201, got %d: %s", bobArticleRec.Code, bobArticleRec.Body.String())
		}
	})

	cancelledRegistrationCases := []struct {
		method string
		path   string
		body   string
	}{
		{http.MethodPost, "/v1/auth/verify", `{"code": "anything"}`},
		{http.MethodDelete, "/v1/auth/register", ""},
	}
	for _, tc := range cancelledRegistrationCases {
		t.Run("A cancelled registration still refuses a jwt minted before the cancellation, on every protected route: "+tc.method+" "+tc.path, func(t *testing.T) {
			token := jwtauth.New("HS256", []byte("test-secret"), nil)
			userRepo := newStaleTokenAliceRepo(false)
			router := newStaleTokenInvalidationRouter(userRepo, newMockArticleRepo(), newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
			preCancelCookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(7), "token_version": uint32(0)})

			cancelRec := sendStaleTokenRequest(router, http.MethodDelete, "/v1/auth/register", "", preCancelCookie)
			if cancelRec.Code != http.StatusOK {
				t.Fatalf("expected status 200, got %d: %s", cancelRec.Code, cancelRec.Body.String())
			}

			rec := sendStaleTokenRequest(router, tc.method, tc.path, tc.body, preCancelCookie)
			if rec.Code != http.StatusUnauthorized {
				t.Fatalf("expected status 401, got %d: %s", rec.Code, rec.Body.String())
			}
			if rec.Body.String() != apperr.UnauthorizedErrorMessage+"\n" {
				t.Errorf("unexpected body: %q", rec.Body.String())
			}
		})
	}

	t.Run("A jwt with no token_version claim at all is refused, treated the same as a mismatch", func(t *testing.T) {
		token := jwtauth.New("HS256", []byte("test-secret"), nil)
		userRepo := newStaleTokenAliceRepo(true)
		router := newStaleTokenInvalidationRouter(userRepo, newMockArticleRepo(), newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
		cookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(7)})

		rec := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, cookie)
		if rec.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", rec.Code, rec.Body.String())
		}
		if rec.Body.String() != apperr.UnauthorizedErrorMessage+"\n" {
			t.Errorf("unexpected body: %q", rec.Body.String())
		}
	})

	t.Run("A jwt whose token_version claim is a string, not a number, is refused instead of crashing", func(t *testing.T) {
		token := jwtauth.New("HS256", []byte("test-secret"), nil)
		userRepo := newStaleTokenAliceRepo(true)
		router := newStaleTokenInvalidationRouter(userRepo, newMockArticleRepo(), newMockImageRepo(), newMockMailRepo(), &stubHasher{}, token)
		cookie := mintStaleTokenCookie(t, token, map[string]any{"user_id": uint32(7), "token_version": "0"})

		rec := sendStaleTokenRequest(router, http.MethodPost, "/v1/article/", `{"title": "A valid title", "body": "Some article body"}`, cookie)
		if rec.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401, got %d: %s", rec.Code, rec.Body.String())
		}
		if rec.Body.String() != apperr.UnauthorizedErrorMessage+"\n" {
			t.Errorf("unexpected body: %q", rec.Body.String())
		}
	})
}