Auth handshake

src/app/api/run/route.ts

@@ -1,16 +1,31 @@
 import { WebSocketServer, WebSocket } from "ws";
 import { randomUUID } from "crypto";
-import { createClient } from "@/lib/supabase/server";
+import { createClientCookie } from "@/lib/supabase/server";
 import { getSession } from "@/modules/account/core";
 import * as Core from "@/modules/live-run/core";
 import { ClientMessage } from "./types";
+import { NextRequest } from "next/server";
+
+export function GET(req: NextRequest) {
+  const headers = new Headers();
+  headers.set("Connection", "Upgrade");
+  headers.set("Upgrade", "websocket");
+  return new Response("Upgrade Required", { status: 426, headers });
+}
 
 export function UPGRADE(client: WebSocket, server: WebSocketServer) {
-  client.on("message", async (raw) => {
-    try {
-      const msg: ClientMessage = JSON.parse(raw.toString());
+  client.once("message", async (raw) => {
+    const { cookies }: { cookies: Record<string, string> } = JSON.parse(
+      raw.toString(),
+    );
 
-      const supabase = await createClient();
+    if (!cookies) {
+      client.send(JSON.stringify({ ok: false, error: "Missing cookies" }));
+      client.close();
+      return;
+    }
+
+    const supabase = await createClientCookie(new Map(Object.entries(cookies)));
     const session = await getSession(supabase);
 
     if (!session) {
@@ -18,6 +33,10 @@ export function UPGRADE(client: WebSocket, server: WebSocketServer) {
       return;
     }
 
+    client.on("message", async (raw) => {
+      try {
+        const msg: ClientMessage = JSON.parse(raw.toString());
+
         switch (msg.type) {
           /* ---------------- CREATE ---------------- */
           case "create": {
@@ -74,16 +93,6 @@ export function UPGRADE(client: WebSocket, server: WebSocketServer) {
               return;
             }
 
-          const supabase = await createClient();
-          const session = await getSession(supabase);
-
-          if (!session) {
-            client.send(
-              JSON.stringify({ ok: false, error: "Invalid session" }),
-            );
-            return;
-          }
-
             const liveRun = await Core.getLiveRun(runId);
 
             if (!liveRun) {
@@ -119,7 +128,9 @@ export function UPGRADE(client: WebSocket, server: WebSocketServer) {
             const { runId } = msg;
 
             if (!runId) {
-            client.send(JSON.stringify({ ok: false, error: "runId required" }));
+              client.send(
+                JSON.stringify({ ok: false, error: "runId required" }),
+              );
               return;
             }
 
@@ -152,7 +163,9 @@ export function UPGRADE(client: WebSocket, server: WebSocketServer) {
             const { runId } = msg;
 
             if (!runId) {
-            client.send(JSON.stringify({ ok: false, error: "runId required" }));
+              client.send(
+                JSON.stringify({ ok: false, error: "runId required" }),
+              );
               return;
             }
 
@@ -186,6 +199,7 @@ export function UPGRADE(client: WebSocket, server: WebSocketServer) {
         );
       }
     });
+  });
 
   client.on("close", () => {
     // optional cleanup

src/lib/supabase/server.ts

@@ -27,22 +27,20 @@ export async function createClient() {
   );
 }
 
-export async function createClientApi() {
-  const cookieStore = await cookies();
-
+export async function createClientCookie(cookies: Map<string, string>) {
   return createServerClient(
     process.env.SUPABASE_URL!,
     process.env.SUPABASE_SECRET_KEY!,
     {
       cookies: {
         get(name: string) {
-          return cookieStore.get(name)?.value;
+          return cookies.get(name) || null;
         },
-        set(name: string, value: string, options: CookieOptions) {
-          cookieStore.set(name, value, options);
+        set(name: string, value: string) {
+          cookies.set(name, value);
         },
         remove(name: string, options: CookieOptions) {
-          cookieStore.set(name, "", options);
+          cookies.delete(name);
         },
       },
     },