looking good

Author: kentcdodds

exercises/99.final/01.solution/package.json

@@ -23,6 +23,7 @@
   "dependencies": {
     "@cloudflare/workers-oauth-provider": "^0.0.3",
     "@epic-web/invariant": "^1.0.0",
+    "@epic-web/totp": "^4.0.1",
     "@modelcontextprotocol/sdk": "^1.10.0",
     "agents": "^0.0.60",
     "zod": "^3.24.3"

exercises/99.final/01.solution/src/db/index.ts

@@ -17,6 +17,7 @@ import {
 	entryTagSchema,
 	newEntryTagSchema,
 	userSchema,
+	grantSchema,
 } from './schema.ts'
 import { sql, snakeToCamel } from './utils.ts'
 
@@ -34,18 +35,31 @@ export class DB {
 		return db
 	}
 
-	async getUserByToken(token: string) {
+	async createUnclaimedGrant(grantUserId: string) {
+		const insertResult = await this.db
+			.prepare(sql`INSERT INTO grants (grant_user_id) VALUES (?1)`)
+			.bind(grantUserId)
+			.run()
+
+		if (!insertResult.success || !insertResult.meta.last_row_id) {
+			throw new Error('Failed to create grant: ' + insertResult.error)
+		}
+
+		return insertResult.meta.last_row_id
+	}
+
+	async getUserByGrantId(grantId: string) {
 		// TODO: I don't have internet, so turn this into a single query when I do...
 		// also add TAKE 1 or whatever
-		const tokensResult = await this.db
-			.prepare(sql`SELECT user_id FROM access_tokens WHERE token_value = ?1`)
-			.bind(token)
+		const grantsResult = await this.db
+			.prepare(sql`SELECT user_id FROM grants WHERE id = ?1`)
+			.bind(grantId)
 			.first()
-		if (!tokensResult) return null
+		if (!grantsResult) return null
 
 		const userResult = await this.db
 			.prepare(sql`SELECT * FROM users WHERE id = ?1`)
-			.bind(tokensResult.user_id)
+			.bind(grantsResult.user_id)
 			.first()
 		if (!userResult) return null
 
@@ -62,35 +76,14 @@ export class DB {
 		return userSchema.parse(snakeToCamel(userResult))
 	}
 
-	async getAccessTokenIdByValue(tokenValue: string) {
-		const tokenResult = await this.db
-			.prepare(sql`SELECT id FROM access_tokens WHERE token_value = ?1`)
-			.bind(tokenValue)
+	async getGrant(grantId: string) {
+		const grantResult = await this.db
+			.prepare(sql`SELECT * FROM grants WHERE id = ?1`)
+			.bind(grantId)
 			.first()
-		if (!tokenResult) return null
+		if (!grantResult) return null
 
-		return tokenResult.id
-	}
-
-	async createAccessTokenIfNecessary(tokenValue: string) {
-		const existingAccessTokenId = await this.getAccessTokenIdByValue(tokenValue)
-		if (existingAccessTokenId) return existingAccessTokenId
-
-		const insertResult = await this.db
-			.prepare(
-				sql`
-					INSERT INTO access_tokens (token_value)
-					VALUES (?1)
-				`,
-			)
-			.bind(tokenValue)
-			.run()
-
-		if (!insertResult.success || !insertResult.meta.last_row_id) {
-			throw new Error('Failed to create access token: ' + insertResult.error)
-		}
-
-		return insertResult.meta.last_row_id
+		return grantSchema.parse(snakeToCamel(grantResult))
 	}
 
 	async createValidationToken(
@@ -117,15 +110,15 @@ export class DB {
 		return insertResult.meta.last_row_id
 	}
 
-	async validateAccessToken(accessTokenId: number, validationToken: string) {
+	async validateTokenToGrant(grantId: number, validationToken: string) {
 		const validationResult = await this.db
 			.prepare(
 				sql`
-					SELECT id, email, access_token_id FROM validation_tokens
-					WHERE access_token_id = ?1 AND token_value = ?2
+					SELECT id, email, grant_id FROM validation_tokens
+					WHERE grant_id = ?1 AND token_value = ?2
 				`,
 			)
-			.bind(accessTokenId, validationToken)
+			.bind(grantId, validationToken)
 			.first()
 
 		if (!validationResult) {
@@ -150,22 +143,19 @@ export class DB {
 		}
 
 		// set access token to user id
-		const claimAccessTokenResult = await this.db
+		const claimGrantResult = await this.db
 			.prepare(
 				sql`
-					UPDATE access_tokens
-					SET user_id = ?2, updated_at = CURRENT_TIMESTAMP 
-					WHERE id = ?1
+					UPDATE grants
+					SET user_id = ?1, updated_at = CURRENT_TIMESTAMP 
+					WHERE id = ?2
 				`,
 			)
-			.bind(validationResult.access_token_id, userId)
+			.bind(userId, validationResult.grant_id)
 			.run()
 
-		if (
-			!claimAccessTokenResult.success ||
-			!claimAccessTokenResult.meta.last_row_id
-		) {
-			throw new Error('Failed to create user: ' + claimAccessTokenResult.error)
+		if (!claimGrantResult.success) {
+			throw new Error('Failed to claim grant: ' + claimGrantResult.error)
 		}
 
 		// delete validation token (fire and forget)
@@ -181,16 +171,21 @@ export class DB {
 		}
 	}
 
-	async deleteAccessToken(userId: number, tokenValue: string) {
-		await this.db
+	async deleteGrant(userId: number, grantId: string) {
+		const deleteResult = await this.db
 			.prepare(
 				sql`
-					DELETE FROM access_tokens
-					WHERE user_id = ?1 AND token_value = ?2
+					DELETE FROM grants
+					WHERE user_id = ?1 AND grant_user_id = ?2
 				`,
 			)
-			.bind(userId, tokenValue)
+			.bind(userId, grantId)
 			.run()
+
+		if (!deleteResult.success) {
+			throw new Error('Failed to delete grant: ' + deleteResult.error)
+		}
+		// TODO: delete the grant from OAUTH_PROVIDER as well
 	}
 
 	async createUserByEmail(email: string) {
@@ -403,9 +398,10 @@ export class DB {
 		return tagSchema.parse(snakeToCamel(result))
 	}
 
-	async listTags() {
+	async listTags(userId: number) {
 		const results = await this.db
-			.prepare(sql`SELECT * FROM tags ORDER BY name`)
+			.prepare(sql`SELECT * FROM tags WHERE user_id = ?1 ORDER BY name`)
+			.bind(userId)
 			.all()
 
 		return z
@@ -521,7 +517,7 @@ export class DB {
 	async getEntryTag(userId: number, id: number) {
 		const result = await this.db
 			.prepare(sql`SELECT * FROM entry_tags WHERE id = ?1 AND user_id = ?2`)
-			.bind(id)
+			.bind(id, userId)
 			.first()
 
 		if (!result) return null

exercises/99.final/01.solution/src/db/migrations.ts

@@ -28,23 +28,29 @@ const migrations = [
 							updated_at integer DEFAULT (CURRENT_TIMESTAMP) NOT NULL
 						);
 					`),
-					// OAuth Access tokens. If user_id is null then it's not yet been claimed
+					// This is a mapping of a grant_user_id (accessible via props.grantId)
+					// to the user_id that can be used to claim the grant. If user_id is
+					// null then it's not yet been claimed. When a user validates their
+					// email, we can update the user_id for the grant and find the
+					// appropriate user in future requests. and then the client
+					// can use that user_id to list and revoke grants (find all grants for
+					// a user then list/revoke them).
 					db.prepare(sql`
-						CREATE TABLE IF NOT EXISTS access_tokens (
+						CREATE TABLE IF NOT EXISTS grants (
 							id integer PRIMARY KEY AUTOINCREMENT NOT NULL,
-							token_value text NOT NULL UNIQUE,
+							grant_user_id text NOT NULL,
 							user_id integer,
 							created_at integer DEFAULT (CURRENT_TIMESTAMP) NOT NULL,
 							updated_at integer DEFAULT (CURRENT_TIMESTAMP) NOT NULL
 						);
 					`),
-					// A OTP emailed to the user to allow them to claim an access_token
+					// An OTP emailed to the user to allow them to claim an access_token
 					db.prepare(sql`
 						CREATE TABLE IF NOT EXISTS validation_tokens (
 							id integer PRIMARY KEY AUTOINCREMENT NOT NULL,
 							token_value text NOT NULL,
 							email text NOT NULL,
-							access_token_id integer NOT NULL,
+							grant_id text NOT NULL,
 							created_at integer DEFAULT (CURRENT_TIMESTAMP) NOT NULL,
 							updated_at integer DEFAULT (CURRENT_TIMESTAMP) NOT NULL
 						);

exercises/99.final/01.solution/src/db/schema.ts

@@ -18,6 +18,23 @@ export const userSchema = z.object({
 	updatedAt: timestampSchema,
 })
 
+export const grantSchema = z.object({
+	id: z.coerce.number(),
+	grantUserId: z.string(),
+	userId: z.coerce.number().optional(),
+	createdAt: timestampSchema,
+	updatedAt: timestampSchema,
+})
+
+export const validationTokenSchema = z.object({
+	id: z.coerce.number(),
+	tokenValue: z.string(),
+	email: z.string().email(),
+	accessTokenId: z.coerce.number(),
+	createdAt: timestampSchema,
+	updatedAt: timestampSchema,
+})
+
 // Schema Validation
 export const entrySchema = z.object({
 	id: z.coerce.number(),

exercises/99.final/01.solution/src/index.ts

@@ -8,7 +8,7 @@ import { initializeTools } from './tools.ts'
 import { type Env } from './types'
 
 type State = {}
-type Props = { accessToken: string }
+type Props = { grantId: string; grantUserId: string }
 
 export class EpicMeMCP extends McpAgent<Env, State, Props> {
 	db!: DB
@@ -60,27 +60,23 @@ const defaultHandler = {
 					return new Response('Invalid client', { status: 400 })
 				}
 
-				const userId = 'EpicMeMCP'
-				// TODO: make this like crypto nice or whatever
-				const accessToken = Math.random().toString(16).slice(2)
+				const db = await DB.getInstance(env)
+				const grantUserId = crypto.randomUUID()
+				const grantId = await db.createUnclaimedGrant(grantUserId)
 
 				const result = await env.OAUTH_PROVIDER.completeAuthorization({
 					request: oauthReqInfo,
-					userId,
-					props: { accessToken },
+					// Here's one of the hacks. We don't know who the user is yet since the token at
+					// this point is unclaimed. But completeAuthorization expects a userId.
+					// So we'll generate a random UUID as a temporary userId
+					userId: grantUserId,
+					props: { grantId, grantUserId },
 					scope: ['full'],
-					metadata: {
-						grantDate: new Date().toISOString(),
-					},
+					metadata: { grantDate: new Date().toISOString() },
 				})
 
 				// Redirect to the client with the authorization code
-				return new Response(null, {
-					status: 302,
-					headers: {
-						Location: result.redirectTo,
-					},
-				})
+				return Response.redirect(result.redirectTo)
 			} catch (error) {
 				console.error('Authorization error:', error)
 				return new Response(
@@ -105,7 +101,6 @@ const oauthProvider = new OAuthProvider({
 	authorizeEndpoint: '/authorize',
 	tokenEndpoint: '/oauth/token',
 	clientRegistrationEndpoint: '/oauth/register',
-	scopesSupported: ['full'],
 })
 
 export default {