From be8587694eb7616e6386281b261ea1215032896b Mon Sep 17 00:00:00 2001
From: static <kmc7468@naver.com>
Date: Sun, 12 Jan 2025 07:59:49 +0900
Subject: [PATCH] =?UTF-8?q?=EC=95=94=ED=98=B8=20=ED=82=A4=20=EB=93=B1?=
 =?UTF-8?q?=EB=A1=9D=20=EC=B1=8C=EB=A6=B0=EC=A7=80=20=EC=B2=98=EB=A6=AC=20?=
 =?UTF-8?q?=EB=B0=A9=EC=8B=9D=EC=9D=84=20=EC=84=B8=EC=85=98=20=EC=97=85?=
 =?UTF-8?q?=EA=B7=B8=EB=A0=88=EC=9D=B4=EB=93=9C=20=EC=B1=8C=EB=A6=B0?=
 =?UTF-8?q?=EC=A7=80=20=EC=B2=98=EB=A6=AC=20=EB=B0=A9=EC=8B=9D=EA=B3=BC=20?=
 =?UTF-8?q?=EB=8F=99=EC=9D=BC=ED=95=98=EA=B2=8C=20=EB=B3=80=EA=B2=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/lib/server/db/client.ts        | 13 +++------
 src/lib/server/db/schema/client.ts | 43 ++++++++++++++++++++----------
 src/lib/server/services/client.ts  | 11 +++-----
 3 files changed, 36 insertions(+), 31 deletions(-)

diff --git a/src/lib/server/db/client.ts b/src/lib/server/db/client.ts
index 7a7b58b..37a1054 100644
--- a/src/lib/server/db/client.ts
+++ b/src/lib/server/db/client.ts
@@ -119,26 +119,21 @@ export const registerUserClientChallenge = async (
   });
 };
 
-export const getUserClientChallenge = async (answer: string, ip: string) => {
+export const consumeUserClientChallenge = async (userId: number, answer: string, ip: string) => {
   const challenges = await db
-    .select()
-    .from(userClientChallenge)
+    .delete(userClientChallenge)
     .where(
       and(
+        eq(userClientChallenge.userId, userId),
         eq(userClientChallenge.answer, answer),
         eq(userClientChallenge.allowedIp, ip),
         gt(userClientChallenge.expiresAt, new Date()),
-        eq(userClientChallenge.isUsed, false),
       ),
     )
-    .limit(1);
+    .returning({ clientId: userClientChallenge.clientId });
   return challenges[0] ?? null;
 };
 
-export const markUserClientChallengeAsUsed = async (id: number) => {
-  await db.update(userClientChallenge).set({ isUsed: true }).where(eq(userClientChallenge.id, id));
-};
-
 export const cleanupExpiredUserClientChallenges = async () => {
   await db.delete(userClientChallenge).where(lte(userClientChallenge.expiresAt, new Date()));
 };
diff --git a/src/lib/server/db/schema/client.ts b/src/lib/server/db/schema/client.ts
index eacd9c9..1e9eb85 100644
--- a/src/lib/server/db/schema/client.ts
+++ b/src/lib/server/db/schema/client.ts
@@ -1,4 +1,11 @@
-import { sqliteTable, text, integer, primaryKey, unique } from "drizzle-orm/sqlite-core";
+import {
+  sqliteTable,
+  text,
+  integer,
+  primaryKey,
+  foreignKey,
+  unique,
+} from "drizzle-orm/sqlite-core";
 import { user } from "./user";
 
 export const client = sqliteTable(
@@ -31,16 +38,24 @@ export const userClient = sqliteTable(
   }),
 );
 
-export const userClientChallenge = sqliteTable("user_client_challenge", {
-  id: integer("id").primaryKey(),
-  userId: integer("user_id")
-    .notNull()
-    .references(() => user.id),
-  clientId: integer("client_id")
-    .notNull()
-    .references(() => client.id),
-  answer: text("answer").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),
-});
+export const userClientChallenge = sqliteTable(
+  "user_client_challenge",
+  {
+    id: integer("id").primaryKey(),
+    userId: integer("user_id")
+      .notNull()
+      .references(() => user.id),
+    clientId: integer("client_id")
+      .notNull()
+      .references(() => client.id),
+    answer: text("answer").notNull().unique(), // Base64
+    allowedIp: text("allowed_ip").notNull(),
+    expiresAt: integer("expires_at", { mode: "timestamp_ms" }).notNull(),
+  },
+  (t) => ({
+    ref: foreignKey({
+      columns: [t.userId, t.clientId],
+      foreignColumns: [userClient.userId, userClient.clientId],
+    }),
+  }),
+);
diff --git a/src/lib/server/services/client.ts b/src/lib/server/services/client.ts
index 73973bb..cc706e1 100644
--- a/src/lib/server/services/client.ts
+++ b/src/lib/server/services/client.ts
@@ -8,8 +8,7 @@ import {
   getUserClient,
   setUserClientStateToPending,
   registerUserClientChallenge,
-  getUserClientChallenge,
-  markUserClientChallengeAsUsed,
+  consumeUserClientChallenge,
 } from "$lib/server/db/client";
 import { IntegrityError } from "$lib/server/db/error";
 import { verifyPubKey, verifySignature, generateChallenge } from "$lib/server/modules/crypto";
@@ -81,15 +80,11 @@ export const verifyUserClient = async (
   answer: string,
   answerSig: string,
 ) => {
-  const challenge = await getUserClientChallenge(answer, ip);
+  const challenge = await consumeUserClientChallenge(userId, answer, ip);
   if (!challenge) {
     error(403, "Invalid challenge answer");
-  } else if (challenge.userId !== userId) {
-    error(403, "Forbidden");
   }
 
-  await markUserClientChallengeAsUsed(challenge.id);
-
   const client = await getClient(challenge.clientId);
   if (!client) {
     error(500, "Invalid challenge answer");
@@ -97,7 +92,7 @@ export const verifyUserClient = async (
     error(403, "Invalid challenge answer signature");
   }
 
-  await setUserClientStateToPending(userId, challenge.clientId);
+  await setUserClientStateToPending(userId, client.id);
 };
 
 export const getUserClientStatus = async (userId: number, clientId: number) => {