Code-based authorization for Drop clients (#145)

* feat: code-based authorization

* fix: final touches

* fix: require session on code fetch endpoint

* feat: better error handling

* refactor: move auth send to client handler

* fix: lint

server/api/v1/client/auth/callback/index.post.ts

@@ -8,13 +8,19 @@ export default defineEventHandler(async (h3) => {
   const body = await readBody(h3);
   const clientId = await body.id;
 
-  const data = await clientHandler.fetchClientMetadata(clientId);
-  if (!data)
+  const client = await clientHandler.fetchClient(clientId);
+  if (!client)
     throw createError({
       statusCode: 400,
       statusMessage: "Invalid or expired client ID.",
     });
 
+  if (client.userId != user.userId)
+    throw createError({
+      statusCode: 403,
+      statusMessage: "Not allowed to authorize this client.",
+    });
+
   const token = await clientHandler.generateAuthToken(clientId);
 
   return {

server/api/v1/client/auth/code/index.get.ts

@@ -0,0 +1,21 @@
+import clientHandler from "~/server/internal/clients/handler";
+import sessionHandler from "~/server/internal/session";
+
+export default defineEventHandler(async (h3) => {
+  const user = await sessionHandler.getSession(h3);
+  if (!user) throw createError({ statusCode: 403 });
+
+  const query = getQuery(h3);
+  const code = query.code?.toString()?.toUpperCase();
+  if (!code)
+    throw createError({
+      statusCode: 400,
+      statusMessage: "Code required in query params.",
+    });
+
+  const clientId = await clientHandler.fetchClientIdByCode(code);
+  if (!clientId)
+    throw createError({ statusCode: 400, statusMessage: "Invalid code." });
+
+  return clientId;
+});

server/api/v1/client/auth/code/index.post.ts

@@ -0,0 +1,35 @@
+import clientHandler from "~/server/internal/clients/handler";
+import sessionHandler from "~/server/internal/session";
+
+export default defineEventHandler(async (h3) => {
+  const user = await sessionHandler.getSession(h3);
+  if (!user) throw createError({ statusCode: 403 });
+
+  const body = await readBody(h3);
+  const clientId = await body.id;
+
+  const client = await clientHandler.fetchClient(clientId);
+  if (!client)
+    throw createError({
+      statusCode: 400,
+      statusMessage: "Invalid or expired client ID.",
+    });
+
+  if (client.userId != user.userId)
+    throw createError({
+      statusCode: 403,
+      statusMessage: "Not allowed to authorize this client.",
+    });
+
+  if (!client.peer)
+    throw createError({
+      statusCode: 500,
+      statusMessage: "No client listening for authorization.",
+    });
+
+  const token = await clientHandler.generateAuthToken(clientId);
+
+  await clientHandler.sendAuthToken(clientId, token);
+
+  return;
+});

server/api/v1/client/auth/code/ws.get.ts

@@ -0,0 +1,25 @@
+import type { FetchError } from "ofetch";
+import clientHandler from "~/server/internal/clients/handler";
+
+export default defineWebSocketHandler({
+  async open(peer) {
+    try {
+      const h3 = { headers: peer.request?.headers ?? new Headers() };
+      const code = h3.headers.get("Authorization");
+      if (!code)
+        throw createError({
+          statusCode: 400,
+          statusMessage: "Code required in Authorization header.",
+        });
+      await clientHandler.connectCodeListener(code, peer);
+    } catch (e) {
+      peer.send(
+        JSON.stringify({
+          type: "error",
+          value: (e as FetchError)?.statusMessage,
+        }),
+      );
+      peer.close();
+    }
+  },
+});

server/api/v1/client/auth/index.get.ts

@@ -13,14 +13,20 @@ export default defineEventHandler(async (h3) => {
       statusMessage: "Provide client ID in request params as 'id'",
     });
 
-  const data = await clientHandler.fetchClientMetadata(providedClientId);
-  if (!data)
+  const client = await clientHandler.fetchClient(providedClientId);
+  if (!client)
     throw createError({
       statusCode: 404,
       statusMessage: "Request not found.",
     });
 
+  if (client.userId && user.userId !== client.userId)
+    throw createError({
+      statusCode: 400,
+      statusMessage: "Client already claimed.",
+    });
+
   await clientHandler.attachUserId(providedClientId, user.userId);
 
-  return data;
+  return client.data;
 });

server/api/v1/client/auth/initiate.post.ts

@@ -1,3 +1,5 @@
+import { type } from "arktype";
+import { readDropValidatedBody, throwingArktype } from "~/server/arktype";
 import type {
   CapabilityConfiguration,
   InternalClientCapability,
@@ -5,23 +7,23 @@ import type {
 import capabilityManager, {
   validCapabilities,
 } from "~/server/internal/clients/capabilities";
-import clientHandler from "~/server/internal/clients/handler";
+import clientHandler, { AuthMode } from "~/server/internal/clients/handler";
 import { parsePlatform } from "~/server/internal/utils/parseplatform";
 
-export default defineEventHandler(async (h3) => {
-  const body = await readBody(h3);
+const ClientAuthInitiate = type({
+  name: "string",
+  platform: "string",
+  capabilities: "object",
+  mode: type.valueOf(AuthMode).default(AuthMode.Callback),
+}).configure(throwingArktype);
+
+export default defineEventHandler(async (h3) => {
+  const body = await readDropValidatedBody(h3, ClientAuthInitiate);
 
-  const name = body.name;
   const platformRaw = body.platform;
   const capabilities: Partial<CapabilityConfiguration> =
     body.capabilities ?? {};
 
-  if (!name || !platformRaw)
-    throw createError({
-      statusCode: 400,
-      statusMessage: "Missing name or platform in body",
-    });
-
   const platform = parsePlatform(platformRaw);
   if (!platform)
     throw createError({
@@ -29,12 +31,6 @@ export default defineEventHandler(async (h3) => {
       statusMessage: "Invalid or unsupported platform",
     });
 
-  if (!capabilities || typeof capabilities !== "object")
-    throw createError({
-      statusCode: 400,
-      statusMessage: "Capabilities must be an array",
-    });
-
   const capabilityIterable = Object.entries(capabilities) as Array<
     [InternalClientCapability, object]
   >;
@@ -64,11 +60,12 @@ export default defineEventHandler(async (h3) => {
       statusMessage: "Invalid capability configuration.",
     });
 
-  const clientId = await clientHandler.initiate({
-    name,
+  const result = await clientHandler.initiate({
+    name: body.name,
     platform,
     capabilities,
+    mode: body.mode,
   });
 
-  return `/client/${clientId}/callback`;
+  return result;
 });

server/internal/clients/handler.ts

@@ -7,11 +7,18 @@ import type {
   InternalClientCapability,
 } from "./capabilities";
 import capabilityManager from "./capabilities";
+import type { PeerImpl } from "../tasks";
+
+export enum AuthMode {
+  Callback = "callback",
+  Code = "code",
+}
 
 export interface ClientMetadata {
   name: string;
   platform: Platform;
   capabilities: Partial<CapabilityConfiguration>;
+  mode: AuthMode;
 }
 
 export class ClientHandler {
@@ -22,8 +29,10 @@ export class ClientHandler {
       data: ClientMetadata;
       userId?: string;
       authToken?: string;
+      peer?: PeerImpl;
     }
   >();
+  private codeClientMap = new Map<string, string>();
 
   async initiate(metadata: ClientMetadata) {
     const clientId = randomUUID();
@@ -32,14 +41,61 @@ export class ClientHandler {
       data: metadata,
       timeout: setTimeout(
         () => {
-          if (this.temporaryClientTable.has(clientId))
+          const client = this.temporaryClientTable.get(clientId);
+          if (client) {
+            if (client.peer) {
+              client.peer.send(
+                JSON.stringify({ type: "error", value: "Request timed out." }),
+              );
+              client.peer.close();
+            }
             this.temporaryClientTable.delete(clientId);
+          }
+
+          const code = this.codeClientMap
+            .entries()
+            .find(([_, v]) => v === clientId);
+          if (code) this.codeClientMap.delete(code[0]);
         },
         1000 * 60 * 10,
       ), // 10 minutes
     });
 
-    return clientId;
+    switch (metadata.mode) {
+      case AuthMode.Callback:
+        return `/client/authorize/${clientId}`;
+      case AuthMode.Code: {
+        const code = randomUUID()
+          .replaceAll(/-/g, "")
+          .slice(0, 7)
+          .toUpperCase();
+        this.codeClientMap.set(code, clientId);
+        return code;
+      }
+    }
+  }
+
+  async connectCodeListener(code: string, peer: PeerImpl) {
+    const clientId = this.codeClientMap.get(code);
+    if (!clientId)
+      throw createError({
+        statusCode: 403,
+        statusMessage: "Invalid or unknown code.",
+      });
+    const metadata = this.temporaryClientTable.get(clientId);
+    if (!metadata)
+      throw createError({ statusCode: 500, statusMessage: "Broken code." });
+    if (metadata.peer)
+      throw createError({
+        statusCode: 400,
+        statusMessage: "Pre-existing listener for this code.",
+      });
+    metadata.peer = peer;
+    this.temporaryClientTable.set(clientId, metadata);
+  }
+
+  async fetchClientIdByCode(code: string) {
+    return this.codeClientMap.get(code);
   }
 
   async fetchClientMetadata(clientId: string) {
@@ -68,6 +124,23 @@ export class ClientHandler {
     return token;
   }
 
+  async sendAuthToken(clientId: string, token: string) {
+    const client = this.temporaryClientTable.get(clientId);
+    if (!client)
+      throw createError({
+        statusCode: 500,
+        statusMessage: "Corrupted code, please restart the process.",
+      });
+    if (!client.peer)
+      throw createError({
+        statusCode: 400,
+        statusMessage: "Client has not connected yet. Please try again later.",
+      });
+    await client.peer.send(
+      JSON.stringify({ type: "token", value: `${clientId}/${token}` }),
+    );
+  }
+
   async fetchClientMetadataByToken(token: string) {
     return this.temporaryClientTable
       .entries()

server/internal/tasks/index.ts

@@ -445,6 +445,7 @@ export type TaskMessage = {
 
 export type PeerImpl = {
   send: (message: string) => void;
+  close: () => void;
 };
 
 export interface BuildTask {