Skip to content

Commit 754ad22

Browse files
mantrakp04mantrakp04
committed
refactor: enhance CLI authentication flow with improved session management and error handling
- Integrated users CRUD handlers for better user validation during anonymous session handling. - Updated SQL queries to ensure proper fetching of refresh tokens and CLI auth attempts from the correct tenancy database. - Enhanced error handling for invalid refresh tokens and user not found scenarios. - Improved response schemas for CLI authentication completion to support various response types. - Updated demo page to utilize sessionStorage for refresh tokens, ensuring better security practices in the development environment.
1 parent 267d4ef commit 754ad22

File tree

6 files changed

+242
-125
lines changed

6 files changed

+242
-125
lines changed

‎apps/backend/src/app/api/latest/auth/cli/complete/route.tsx‎

Lines changed: 145 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,87 @@
11
import { Prisma } from "@/generated/prisma/client";
2+
import { usersCrudHandlers } from "@/app/api/latest/users/crud";
23
import { generateAccessTokenFromRefreshTokenIfValid } from "@/lib/tokens";
3-
import { globalPrismaClient } from "@/prisma-client";
4+
import { getPrismaClientForTenancy, getPrismaSchemaForTenancy, globalPrismaClient, sqlQuoteIdent } from "@/prisma-client";
45
import { createSmartRouteHandler } from "@/route-handlers/smart-route-handler";
5-
import { adaptSchema, clientOrHigherAuthTypeSchema, yupBoolean, yupNumber, yupObject, yupString } from "@stackframe/stack-shared/dist/schema-fields";
6+
import { Tenancy } from "@/lib/tenancies";
7+
import { KnownErrors } from "@stackframe/stack-shared";
8+
import { adaptSchema, clientOrHigherAuthTypeSchema, yupBoolean, yupNumber, yupObject, yupString, yupUnion } from "@stackframe/stack-shared/dist/schema-fields";
69
import { StatusError } from "@stackframe/stack-shared/dist/utils/errors";
10+
import type { InferType } from "yup";
711

812
type CliSessionState = "anonymous" | "none";
913

14+
const postCliAuthCompleteRequestSchema = yupObject({
15+
auth: yupObject({
16+
type: clientOrHigherAuthTypeSchema,
17+
tenancy: adaptSchema.defined(),
18+
}).defined(),
19+
body: yupObject({
20+
login_code: yupString().defined(),
21+
mode: yupString().oneOf(["check", "claim-anon-session", "complete"]).default("complete"),
22+
refresh_token: yupString().optional(),
23+
}).defined(),
24+
});
25+
26+
const postCliAuthCompleteResponseSchema = yupUnion(
27+
yupObject({
28+
statusCode: yupNumber().oneOf([200]).defined(),
29+
bodyType: yupString().oneOf(["json"]).defined(),
30+
body: yupObject({
31+
cli_session_state: yupString().oneOf(["anonymous", "none"]).defined(),
32+
}).defined(),
33+
}).defined(),
34+
yupObject({
35+
statusCode: yupNumber().oneOf([200]).defined(),
36+
bodyType: yupString().oneOf(["json"]).defined(),
37+
body: yupObject({
38+
access_token: yupString().defined(),
39+
refresh_token: yupString().defined(),
40+
}).defined(),
41+
}).defined(),
42+
yupObject({
43+
statusCode: yupNumber().oneOf([200]).defined(),
44+
bodyType: yupString().oneOf(["json"]).defined(),
45+
body: yupObject({
46+
success: yupBoolean().oneOf([true]).defined(),
47+
}).defined(),
48+
}).defined(),
49+
).defined();
50+
51+
type PostCliAuthCompleteRequest = InferType<typeof postCliAuthCompleteRequestSchema>;
52+
type PostCliAuthCompleteResponse = InferType<typeof postCliAuthCompleteResponseSchema>;
53+
54+
function cliAuthCompleteCheckResponse(cliSessionState: CliSessionState): PostCliAuthCompleteResponse {
55+
return {
56+
statusCode: 200,
57+
bodyType: "json",
58+
body: {
59+
cli_session_state: cliSessionState,
60+
},
61+
};
62+
}
63+
64+
function cliAuthCompleteClaimResponse(accessToken: string, refreshToken: string): PostCliAuthCompleteResponse {
65+
return {
66+
statusCode: 200,
67+
bodyType: "json",
68+
body: {
69+
access_token: accessToken,
70+
refresh_token: refreshToken,
71+
},
72+
};
73+
}
74+
75+
function cliAuthCompleteSuccessResponse(): PostCliAuthCompleteResponse {
76+
return {
77+
statusCode: 200,
78+
bodyType: "json",
79+
body: {
80+
success: true,
81+
},
82+
};
83+
}
84+
1085
type CliAuthAttemptRow = {
1186
id: string,
1287
tenancyId: string,
@@ -24,28 +99,27 @@ type RefreshTokenRow = {
2499
expiresAt: Date | null,
25100
};
26101

27-
async function getPendingCliAuthAttempt(tenancyId: string, loginCode: string) {
28-
const rows = await globalPrismaClient.$queryRaw<CliAuthAttemptRow[]>(Prisma.sql`
102+
async function getPendingCliAuthAttempt(tenancy: Tenancy, loginCode: string) {
103+
// CliAuthAttempt lives in the tenancy's source-of-truth DB, consistent with cli/poll/route.tsx.
104+
const prisma = await getPrismaClientForTenancy(tenancy);
105+
const schema = await getPrismaSchemaForTenancy(tenancy);
106+
const rows = await prisma.$queryRaw<CliAuthAttemptRow[]>(Prisma.sql`
29107
SELECT
30108
"id",
31109
"tenancyId",
32110
"refreshToken",
33111
"anonRefreshToken",
34112
"expiresAt",
35113
"usedAt"
36-
FROM "CliAuthAttempt"
37-
WHERE "loginCode" = ${loginCode}
114+
FROM ${sqlQuoteIdent(schema)}."CliAuthAttempt"
115+
WHERE "tenancyId" = ${tenancy.id}::UUID
116+
AND "loginCode" = ${loginCode}
38117
LIMIT 1
39118
`);
40-
const cliAuth = rows[0];
41-
42-
if (!cliAuth) {
119+
if (rows.length === 0) {
43120
throw new StatusError(400, "Invalid login code or the code has expired");
44121
}
45-
46-
if (cliAuth.tenancyId !== tenancyId) {
47-
throw new StatusError(400, "Project ID mismatch; please ensure that you are using the correct app url.");
48-
}
122+
const cliAuth = rows[0];
49123

50124
if (cliAuth.refreshToken !== null || cliAuth.usedAt !== null || cliAuth.expiresAt < new Date()) {
51125
throw new StatusError(400, "Invalid login code or the code has expired");
@@ -55,6 +129,7 @@ async function getPendingCliAuthAttempt(tenancyId: string, loginCode: string) {
55129
}
56130

57131
async function getRefreshTokenSession(tenancyId: string, refreshToken: string) {
132+
// ProjectUserRefreshToken lives in the global DB (see tokens.tsx and oauth/model.tsx).
58133
const rows = await globalPrismaClient.$queryRaw<RefreshTokenRow[]>(Prisma.sql`
59134
SELECT
60135
"id",
@@ -66,11 +141,10 @@ async function getRefreshTokenSession(tenancyId: string, refreshToken: string) {
66141
WHERE "refreshToken" = ${refreshToken}
67142
LIMIT 1
68143
`);
69-
const refreshTokenObj = rows[0];
70-
71-
if (!refreshTokenObj) {
144+
if (rows.length === 0) {
72145
return null;
73146
}
147+
const refreshTokenObj = rows[0];
74148

75149
if (refreshTokenObj.tenancyId !== tenancyId) {
76150
throw new StatusError(400, "Refresh token does not belong to this project");
@@ -83,84 +157,89 @@ async function getRefreshTokenSession(tenancyId: string, refreshToken: string) {
83157
return refreshTokenObj;
84158
}
85159

86-
async function getCliAnonymousSession(tenancyId: string, anonRefreshToken: string | null) {
160+
async function getCliAnonymousSession(tenancy: Tenancy, anonRefreshToken: string | null) {
87161
if (anonRefreshToken === null) {
88162
return null;
89163
}
90164

91-
const refreshTokenObj = await getRefreshTokenSession(tenancyId, anonRefreshToken);
165+
const refreshTokenObj = await getRefreshTokenSession(tenancy.id, anonRefreshToken);
92166
if (!refreshTokenObj) {
93167
return null;
94168
}
95169

96-
const userRows = await globalPrismaClient.$queryRaw<{ projectUserId: string, isAnonymous: boolean }[]>(Prisma.sql`
97-
SELECT "projectUserId", "isAnonymous"
98-
FROM "ProjectUser"
99-
WHERE "tenancyId" = ${tenancyId}::UUID
100-
AND "projectUserId" = ${refreshTokenObj.projectUserId}::UUID
101-
LIMIT 1
102-
`);
103-
const user = userRows[0];
170+
// ProjectUser lives in the tenancy's source-of-truth DB, not global.
171+
// Use the CRUD handler which is topology-aware (matches tokens.tsx:206).
172+
let user;
173+
try {
174+
user = await usersCrudHandlers.adminRead({
175+
tenancy,
176+
user_id: refreshTokenObj.projectUserId,
177+
allowedErrorTypes: [KnownErrors.UserNotFound],
178+
});
179+
} catch (error) {
180+
if (error instanceof KnownErrors.UserNotFound) {
181+
return null;
182+
}
183+
throw error;
184+
}
104185

105-
if (!user?.isAnonymous) {
186+
if (!user.is_anonymous) {
106187
return null;
107188
}
108189

109190
return {
110191
refreshTokenObj,
111-
userId: user.projectUserId,
192+
userId: user.id,
112193
};
113194
}
114195

115-
export const POST = createSmartRouteHandler({
196+
export const POST = createSmartRouteHandler<PostCliAuthCompleteRequest, PostCliAuthCompleteResponse>({
116197
metadata: {
117198
summary: "Complete CLI authentication",
118199
description: "Inspect, claim, or complete a CLI authentication session",
119200
tags: ["CLI Authentication"],
120201
},
121-
request: yupObject({
122-
auth: yupObject({
123-
type: clientOrHigherAuthTypeSchema,
124-
tenancy: adaptSchema.defined(),
125-
}).defined(),
126-
body: yupObject({
127-
login_code: yupString().defined(),
128-
mode: yupString().oneOf(["check", "claim-anon-session", "complete"]).default("complete"),
129-
refresh_token: yupString().optional(),
130-
}).defined(),
131-
}),
132-
response: yupObject({
133-
statusCode: yupNumber().oneOf([200]).defined(),
134-
bodyType: yupString().oneOf(["json"]).defined(),
135-
body: yupObject({
136-
success: yupBoolean().optional(),
137-
cli_session_state: yupString().oneOf(["anonymous", "none"]).optional(),
138-
access_token: yupString().optional(),
139-
refresh_token: yupString().optional(),
140-
}).defined(),
141-
}),
202+
request: postCliAuthCompleteRequestSchema,
203+
response: postCliAuthCompleteResponseSchema,
142204
async handler({ auth: { tenancy }, body: { login_code, mode, refresh_token } }) {
143-
const cliAuth = await getPendingCliAuthAttempt(tenancy.id, login_code);
205+
const cliAuth = await getPendingCliAuthAttempt(tenancy, login_code);
144206

145207
if (mode === "check") {
146-
const cliAnonymousSession = await getCliAnonymousSession(tenancy.id, cliAuth.anonRefreshToken);
208+
const cliAnonymousSession = await getCliAnonymousSession(tenancy, cliAuth.anonRefreshToken);
147209
const cliSessionState: CliSessionState = cliAnonymousSession ? "anonymous" : "none";
148210

149-
return {
150-
statusCode: 200,
151-
bodyType: "json" as const,
152-
body: {
153-
cli_session_state: cliSessionState,
154-
},
155-
};
211+
return cliAuthCompleteCheckResponse(cliSessionState);
156212
}
157213

158214
if (mode === "claim-anon-session") {
159-
const cliAnonymousSession = await getCliAnonymousSession(tenancy.id, cliAuth.anonRefreshToken);
215+
const cliAnonymousSession = await getCliAnonymousSession(tenancy, cliAuth.anonRefreshToken);
160216
if (!cliAnonymousSession) {
161217
throw new StatusError(400, "No anonymous session associated with this code");
162218
}
163219

220+
// Atomically consume the anon session (one-shot): null out anonRefreshToken
221+
// on the CliAuthAttempt row so subsequent claim-anon-session calls cannot
222+
// replay and re-retrieve the anon user's refresh token.
223+
const prisma = await getPrismaClientForTenancy(tenancy);
224+
const schema = await getPrismaSchemaForTenancy(tenancy);
225+
const consumed = await prisma.$queryRaw<{ id: string }[]>(Prisma.sql`
226+
UPDATE ${sqlQuoteIdent(schema)}."CliAuthAttempt"
227+
SET
228+
"anonRefreshToken" = NULL,
229+
"updatedAt" = NOW()
230+
WHERE "tenancyId" = ${tenancy.id}::UUID
231+
AND "id" = ${cliAuth.id}::UUID
232+
AND "anonRefreshToken" = ${cliAuth.anonRefreshToken}
233+
AND "refreshToken" IS NULL
234+
AND "usedAt" IS NULL
235+
AND "expiresAt" > NOW()
236+
RETURNING "id"
237+
`);
238+
239+
if (consumed.length === 0) {
240+
throw new StatusError(400, "No anonymous session associated with this code");
241+
}
242+
164243
const accessToken = await generateAccessTokenFromRefreshTokenIfValid({
165244
tenancy,
166245
refreshTokenObj: cliAnonymousSession.refreshTokenObj,
@@ -170,14 +249,7 @@ export const POST = createSmartRouteHandler({
170249
throw new StatusError(400, "Anonymous session is no longer valid");
171250
}
172251

173-
return {
174-
statusCode: 200,
175-
bodyType: "json" as const,
176-
body: {
177-
access_token: accessToken,
178-
refresh_token: cliAnonymousSession.refreshTokenObj.refreshToken,
179-
},
180-
};
252+
return cliAuthCompleteClaimResponse(accessToken, cliAnonymousSession.refreshTokenObj.refreshToken);
181253
}
182254

183255
if (!refresh_token) {
@@ -194,8 +266,10 @@ export const POST = createSmartRouteHandler({
194266
// the anonymous user into the authenticated user (that was a security risk).
195267
// The anonymous user is left untouched and will simply be orphaned from
196268
// this CLI flow.
197-
const claimed = await globalPrismaClient.$queryRaw<{ id: string }[]>(Prisma.sql`
198-
UPDATE "CliAuthAttempt"
269+
const prisma = await getPrismaClientForTenancy(tenancy);
270+
const schema = await getPrismaSchemaForTenancy(tenancy);
271+
const claimed = await prisma.$queryRaw<{ id: string }[]>(Prisma.sql`
272+
UPDATE ${sqlQuoteIdent(schema)}."CliAuthAttempt"
199273
SET
200274
"refreshToken" = ${refresh_token},
201275
"anonRefreshToken" = NULL,
@@ -212,12 +286,6 @@ export const POST = createSmartRouteHandler({
212286
throw new StatusError(400, "Invalid login code or the code has expired");
213287
}
214288

215-
return {
216-
statusCode: 200,
217-
bodyType: "json" as const,
218-
body: {
219-
success: true,
220-
},
221-
};
289+
return cliAuthCompleteSuccessResponse();
222290
},
223291
});

0 commit comments

Comments
 (0)