fix

2026-10-08 02:07:45 +07:00
parent 56bc009311
commit 0210a0ed15
6 changed files with 285 additions and 24 deletions
@@ -67,16 +67,25 @@ func (s *RefreshTokensRepo) ExistsRefreshToken(
 	return exists, nil
 }

-// DeleteRefreshTokenByJti removes a refresh token by its JTI.
-func (s *RefreshTokensRepo) DeleteRefreshTokenByJti(
+// ReplaceRefreshToken atomically consumes an unexpired token and stores its replacement.
+func (s *RefreshTokensRepo) ReplaceRefreshToken(
 	ctx context.Context,
-	jti string,
-) error {
-	_, err := s.pool.Exec(
-		ctx,
-		`DELETE FROM refresh_tokens 
-		WHERE jti = $1`,
-		jti,
-	)
-	return err
+	userID int,
+	oldJTI string,
+	now time.Time,
+	expiresAt time.Time,
+	newJTI string,
+) (bool, error) {
+	result, err := s.pool.Exec(ctx, `
+		WITH consumed AS (
+			DELETE FROM refresh_tokens
+			WHERE user_id = $1 AND jti = $2 AND expires_at > $3
+			RETURNING user_id
+		)
+		INSERT INTO refresh_tokens (user_id, expires_at, jti)
+		SELECT user_id, $4, $5 FROM consumed`, userID, oldJTI, now, expiresAt, newJTI)
+	if err != nil {
+		return false, err
+	}
+	return result.RowsAffected() == 1, nil
 }
@@ -0,0 +1,203 @@
+package refresh_tokens_repo_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"evening_detective_server/internal/modules/processor_jwt"
+	"evening_detective_server/internal/repos/refresh_tokens_repo"
+	"evening_detective_server/internal/repos/user_agreements_repo"
+	"evening_detective_server/internal/repos/users_repo"
+	"evening_detective_server/internal/services/users_service"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// newRefreshFixture prepares an isolated schema; scenarios use repository methods only.
+func newRefreshFixture(t *testing.T) (*refresh_tokens_repo.RefreshTokensRepo, int, *users_repo.UsersRepo) {
+	t.Helper()
+	dsn := os.Getenv("TEST_DATABASE_URL")
+	if dsn == "" {
+		t.Skip("set TEST_DATABASE_URL to run PostgreSQL integration checks")
+	}
+	ctx := context.Background()
+	admin, err := pgxpool.New(ctx, dsn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	schema := fmt.Sprintf("refresh_test_%d", time.Now().UnixNano())
+	quoted := pgx.Identifier{schema}.Sanitize()
+	if _, err := admin.Exec(ctx, "CREATE SCHEMA "+quoted); err != nil {
+		admin.Close()
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if _, err := admin.Exec(ctx, "DROP SCHEMA "+quoted+" CASCADE"); err != nil {
+			t.Error(err)
+		}
+		admin.Close()
+	})
+	config, err := pgxpool.ParseConfig(dsn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	config.ConnConfig.RuntimeParams["search_path"] = schema
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(pool.Close)
+	for _, migration := range []string{
+		"20260628175120_create_users_table.sql",
+		"20260704192907_refresh_tokens_table.sql",
+		"20260818231943_create_user_agreements_table.sql",
+	} {
+		data, err := os.ReadFile("../../../migrations/" + migration)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := pool.Exec(ctx, strings.Split(string(data), "-- +goose Down")[0]); err != nil {
+			t.Fatal(err)
+		}
+	}
+	users := users_repo.NewUserRepo(pool)
+	userID, err := users.AddUserWithAgreements(ctx, "test", "test@example.invalid", "hash", []string{"user"}, []user_agreements_repo.Agreement{{AgreementType: "terms", Version: "test"}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	return refresh_tokens_repo.NewRefreshTokensRepo(pool), userID, users
+}
+
+func TestReplaceRefreshTokenIsAtomic(t *testing.T) {
+	for _, tc := range []struct {
+		name          string
+		oldLifetime   time.Duration
+		wrongUser     bool
+		newJTI        string
+		attempts      int
+		wantSuccesses int
+		wantError     bool
+		wantOld       bool
+	}{
+		{"successful replacement", time.Hour, false, "new", 1, 1, false, false},
+		{"replay rejected", time.Hour, false, "new", 2, 1, false, false},
+		{"concurrent replacement", time.Hour, false, "new", 8, 1, false, false},
+		{"expired token", -time.Hour, false, "new", 1, 0, false, true},
+		{"wrong user", time.Hour, true, "new", 1, 0, false, true},
+		{"insert failure preserves old token", time.Hour, false, "collision", 1, 0, true, true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			repo, userID, _ := newRefreshFixture(t)
+			ctx := context.Background()
+			now := time.Now()
+			if err := repo.AddRefreshToken(ctx, userID, now.Add(tc.oldLifetime), "old"); err != nil {
+				t.Fatal(err)
+			}
+			if err := repo.AddRefreshToken(ctx, userID, now.Add(time.Hour), "collision"); err != nil {
+				t.Fatal(err)
+			}
+			replaceUser := userID
+			if tc.wrongUser {
+				replaceUser++
+			}
+			type result struct {
+				ok  bool
+				err error
+			}
+			results := make(chan result, tc.attempts)
+			for range tc.attempts {
+				go func() {
+					ok, err := repo.ReplaceRefreshToken(ctx, replaceUser, "old", now, now.Add(time.Hour), tc.newJTI)
+					results <- result{ok, err}
+				}()
+			}
+			successes := 0
+			for range tc.attempts {
+				got := <-results
+				if (got.err != nil) != tc.wantError {
+					t.Fatalf("unexpected replacement error: %v", got.err)
+				}
+				if got.ok {
+					successes++
+				}
+			}
+			if successes != tc.wantSuccesses {
+				t.Fatalf("successes=%d, want %d", successes, tc.wantSuccesses)
+			}
+			// Use a point before expiration to distinguish deletion from an expired record.
+			oldExists, err := repo.ExistsRefreshToken(ctx, userID, "old", now.Add(-2*time.Hour))
+			if err != nil || oldExists != tc.wantOld {
+				t.Fatalf("old token exists=%v, err=%v", oldExists, err)
+			}
+			if tc.wantSuccesses > 0 {
+				exists, err := repo.ExistsRefreshToken(ctx, userID, tc.newJTI, now)
+				if err != nil || !exists {
+					t.Fatalf("replacement exists=%v, err=%v", exists, err)
+				}
+			}
+		})
+	}
+}
+
+func TestRefreshExchangesPersistedTokens(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		persist bool
+	}{
+		{"valid refresh returns new access and refresh and rejects replay", true},
+		{"unregistered refresh is rejected", false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			repo, userID, users := newRefreshFixture(t)
+			ctx := context.Background()
+			processor := processor_jwt.NewProcessor("test-secret")
+			oldToken, err := processor.GenerateRefreshToken(userID)
+			if err != nil {
+				t.Fatal(err)
+			}
+			oldClaims, err := processor.GetClaims(oldToken)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if tc.persist {
+				if err := repo.AddRefreshToken(ctx, userID, oldClaims.ExpiresAt.Time, oldClaims.ID); err != nil {
+					t.Fatal(err)
+				}
+			}
+			service := users_service.NewUsersService(users, nil, nil, processor, repo, "")
+			access, refresh, err := service.Refresh(ctx, oldToken)
+			if !tc.persist {
+				if !errors.Is(err, users_service.ErrInvalidRefreshToken) {
+					t.Fatalf("err=%v", err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			accessClaims, err := processor.GetClaims(access)
+			if err != nil || accessClaims.UserID != userID || !accessClaims.HasRole("user") {
+				t.Fatalf("access claims=%#v, err=%v", accessClaims, err)
+			}
+			refreshClaims, err := processor.GetClaims(refresh)
+			if err != nil || refreshClaims.UserID != userID || refreshClaims.ID == oldClaims.ID {
+				t.Fatalf("refresh claims=%#v, err=%v", refreshClaims, err)
+			}
+			exists, err := repo.ExistsRefreshToken(ctx, userID, refreshClaims.ID, time.Now())
+			if err != nil || !exists {
+				t.Fatalf("replacement exists=%v, err=%v", exists, err)
+			}
+			if _, _, err := service.Refresh(ctx, oldToken); !errors.Is(err, users_service.ErrInvalidRefreshToken) {
+				t.Fatalf("replay err=%v", err)
+			}
+			if _, _, err := service.Refresh(ctx, refresh); err != nil {
+				t.Fatalf("new token cannot be used: %v", err)
+			}
+		})
+	}
+}