From a64e85848cab40e2ef487cba79c578ab519143ad Mon Sep 17 00:00:00 2001
From: static <kmc7468@naver.com>
Date: Tue, 31 Dec 2024 03:05:14 +0900
Subject: [PATCH] =?UTF-8?q?=EC=B1=8C=EB=A6=B0=EC=A7=80=20Reply=20Attack=20?=
 =?UTF-8?q?=EB=B0=A9=EC=96=B4=20=EA=B5=AC=ED=98=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/lib/server/db/client.ts        | 9 +++++++++
 src/lib/server/db/schema/client.ts | 1 +
 src/lib/server/db/schema/token.ts  | 1 +
 src/lib/server/db/token.ts         | 9 +++++++++
 src/lib/server/services/auth.ts    | 3 ++-
 src/lib/server/services/client.ts  | 4 ++--
 6 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/lib/server/db/client.ts b/src/lib/server/db/client.ts
index 5423949..29f4806 100644
--- a/src/lib/server/db/client.ts
+++ b/src/lib/server/db/client.ts
@@ -118,12 +118,21 @@ export const getUserClientChallenge = async (answer: string, ip: string) => {
         eq(userClientChallenge.answer, answer),
         eq(userClientChallenge.allowedIp, ip),
         gt(userClientChallenge.expiresAt, new Date()),
+        eq(userClientChallenge.isUsed, false),
       ),
     )
     .execute();
   return challenges[0] ?? null;
 };
 
+export const markUserClientChallengeAsUsed = async (id: number) => {
+  await db
+    .update(userClientChallenge)
+    .set({ isUsed: true })
+    .where(eq(userClientChallenge.id, id))
+    .execute();
+};
+
 export const cleanupExpiredUserClientChallenges = async () => {
   await db
     .delete(userClientChallenge)
diff --git a/src/lib/server/db/schema/client.ts b/src/lib/server/db/schema/client.ts
index bab3475..7d83435 100644
--- a/src/lib/server/db/schema/client.ts
+++ b/src/lib/server/db/schema/client.ts
@@ -42,4 +42,5 @@ export const userClientChallenge = sqliteTable("user_client_challenge", {
   answer: text("challenge").notNull().unique(), // Base64
   allowedIp: text("allowed_ip").notNull(),
   expiresAt: integer("expires_at", { mode: "timestamp_ms" }).notNull(),
+  isUsed: integer("is_used", { mode: "boolean" }).notNull().default(false),
 });
diff --git a/src/lib/server/db/schema/token.ts b/src/lib/server/db/schema/token.ts
index dcf995d..72106d7 100644
--- a/src/lib/server/db/schema/token.ts
+++ b/src/lib/server/db/schema/token.ts
@@ -28,4 +28,5 @@ export const tokenUpgradeChallenge = sqliteTable("token_upgrade_challenge", {
   answer: text("challenge").notNull().unique(), // Base64
   allowedIp: text("allowed_ip").notNull(),
   expiresAt: integer("expires_at", { mode: "timestamp_ms" }).notNull(),
+  isUsed: integer("is_used", { mode: "boolean" }).notNull().default(false),
 });
diff --git a/src/lib/server/db/token.ts b/src/lib/server/db/token.ts
index 51499cf..d93527f 100644
--- a/src/lib/server/db/token.ts
+++ b/src/lib/server/db/token.ts
@@ -102,12 +102,21 @@ export const getTokenUpgradeChallenge = async (answer: string, ip: string) => {
         eq(tokenUpgradeChallenge.answer, answer),
         eq(tokenUpgradeChallenge.allowedIp, ip),
         gt(tokenUpgradeChallenge.expiresAt, new Date()),
+        eq(tokenUpgradeChallenge.isUsed, false),
       ),
     )
     .execute();
   return challenges[0] ?? null;
 };
 
+export const markTokenUpgradeChallengeAsUsed = async (id: number) => {
+  await db
+    .update(tokenUpgradeChallenge)
+    .set({ isUsed: true })
+    .where(eq(tokenUpgradeChallenge.id, id))
+    .execute();
+};
+
 export const cleanupExpiredTokenUpgradeChallenges = async () => {
   await db
     .delete(tokenUpgradeChallenge)
diff --git a/src/lib/server/services/auth.ts b/src/lib/server/services/auth.ts
index d8b3fa1..aeaf858 100644
--- a/src/lib/server/services/auth.ts
+++ b/src/lib/server/services/auth.ts
@@ -13,6 +13,7 @@ import {
   revokeRefreshToken,
   registerTokenUpgradeChallenge,
   getTokenUpgradeChallenge,
+  markTokenUpgradeChallengeAsUsed,
 } from "$lib/server/db/token";
 import { issueToken, verifyToken, TokenError } from "$lib/server/modules/auth";
 import { verifySignature, generateChallenge } from "$lib/server/modules/crypto";
@@ -152,7 +153,7 @@ export const upgradeToken = async (
     error(401, "Invalid challenge answer signature");
   }
 
-  // TODO: Replay attack prevention
+  await markTokenUpgradeChallengeAsUsed(challenge.id);
 
   const newJti = uuidv4();
   if (!(await upgradeRefreshToken(oldJti, newJti, client.id))) {
diff --git a/src/lib/server/services/client.ts b/src/lib/server/services/client.ts
index c3ad525..07729c0 100644
--- a/src/lib/server/services/client.ts
+++ b/src/lib/server/services/client.ts
@@ -11,6 +11,7 @@ import {
   setUserClientStateToPending,
   registerUserClientChallenge,
   getUserClientChallenge,
+  markUserClientChallengeAsUsed,
 } from "$lib/server/db/client";
 import { verifyPubKey, verifySignature, generateChallenge } from "$lib/server/modules/crypto";
 import { isInitialMekNeeded } from "$lib/server/modules/mek";
@@ -107,7 +108,6 @@ export const verifyUserClient = async (
     error(401, "Invalid challenge answer signature");
   }
 
-  // TODO: Replay attack prevention
-
+  await markUserClientChallengeAsUsed(challenge.id);
   await setUserClientStateToPending(userId, challenge.clientId);
 };