From a855004fd203414791c1233a9917ecb5baaffc3a Mon Sep 17 00:00:00 2001
From: Zoe Roux <zoe.roux@zoriya.dev>
Date: Sun, 14 Dec 2025 23:12:15 +0100
Subject: [PATCH] Implement watch websocket api

---
 api/src/auth.ts       |  23 +++++----
 api/src/base.ts       |   4 +-
 api/src/websockets.ts | 114 ++++++++++++++++++++++++++++++++++++------
 3 files changed, 114 insertions(+), 27 deletions(-)

diff --git a/api/src/auth.ts b/api/src/auth.ts
index d7f0d874..101a5dc0 100644
--- a/api/src/auth.ts
+++ b/api/src/auth.ts
@@ -33,6 +33,17 @@ const Jwt = t.Object({
 type Jwt = typeof Jwt.static;
 const validator = TypeCompiler.Compile(Jwt);
 
+export async function verifyJwt(bearer: string) {
+	// @ts-expect-error ts can't understand that there's two overload idk why
+	const { payload } = await jwtVerify(bearer, jwtSecret ?? jwks, {
+		issuer: process.env.JWT_ISSUER,
+	});
+	const raw = validator.Decode(payload);
+	const jwt = Value.Default(Jwt, raw) as Prettify<Jwt & { settings: Settings }>;
+
+	return { jwt };
+}
+
 export const auth = new Elysia({ name: "auth" })
 	.guard({
 		headers: t.Object(
@@ -50,18 +61,8 @@ export const auth = new Elysia({ name: "auth" })
 				message: "No authorization header was found.",
 			});
 		}
-
 		try {
-			// @ts-expect-error ts can't understand that there's two overload idk why
-			const { payload } = await jwtVerify(bearer, jwtSecret ?? jwks, {
-				issuer: process.env.JWT_ISSUER,
-			});
-			const raw = validator.Decode(payload);
-			const jwt = Value.Default(Jwt, raw) as Prettify<
-				Jwt & { settings: Settings }
-			>;
-
-			return { jwt };
+			return await verifyJwt(bearer);
 		} catch (err) {
 			return status(403, {
 				status: 403,
diff --git a/api/src/base.ts b/api/src/base.ts
index ad33b439..920d88ff 100644
--- a/api/src/base.ts
+++ b/api/src/base.ts
@@ -17,6 +17,7 @@ import { videosReadH, videosWriteH } from "./controllers/videos";
 import { db } from "./db";
 import type { KError } from "./models/error";
 import { otel } from "./otel";
+import { appWs } from "./websockets";
 
 export const base = new Elysia({ name: "base" })
 	.onError(({ code, error }) => {
@@ -91,8 +92,9 @@ export const base = new Elysia({ name: "base" })
 export const prefix = "/api";
 export const handlers = new Elysia({ prefix })
 	.use(base)
-	.use(auth)
 	.use(otel)
+	.use(appWs)
+	.use(auth)
 	.guard(
 		{
 			// Those are not applied for now. See https://github.com/elysiajs/elysia/issues/1139
diff --git a/api/src/websockets.ts b/api/src/websockets.ts
index c581c700..69bbf62b 100644
--- a/api/src/websockets.ts
+++ b/api/src/websockets.ts
@@ -1,20 +1,104 @@
-import Elysia, { t } from "elysia";
+import type { TObject, TString } from "@sinclair/typebox";
+import Elysia, { type TSchema, t } from "elysia";
+import { verifyJwt } from "./auth";
+import { updateHistory, updateWatchlist } from "./controllers/profiles/history";
+import { getOrCreateProfile } from "./controllers/profiles/profile";
+import { db } from "./db";
+import { SeedHistory } from "./models/history";
 
-const actionMap: Record<string, ][> = [
+const actionMap = {
+	ping: handler({
+		message(ws) {
+			ws.send({ response: "pong" });
+		},
+	}),
+	watch: handler({
+		body: t.Omit(SeedHistory, ["playedDate"]),
+		permissions: ["core.read"],
+		async message(ws, body) {
+			const profilePk = await getOrCreateProfile(ws.data.jwt.sub);
 
-]
+			await db.transaction(async (tx) => {
+				const hist = await updateHistory(tx, profilePk, [body]);
+				await updateWatchlist(tx, profilePk, hist);
+			});
+			ws.send({ response: "ok" });
+		},
+	}),
+};
 
-export const appWs = new Elysia().ws("/ws", {
-	body: t.Union([
-		t.Object({
-			action: t.Literal("ping"),
-		}),
-		t.Object({
-			action: t.Literal("watch"),
-			entry: t.String(),
-		}),
-	]),
-	message(ws, { message }) {
-		actionMap[message.action](message);
+const baseWs = new Elysia()
+	.guard({
+		headers: t.Object(
+			{
+				authorization: t.Optional(t.TemplateLiteral("Bearer ${string}")),
+				"Sec-WebSocket-Protocol": t.Optional(
+					t.Array(
+						t.Union([t.Literal("kyoo"), t.TemplateLiteral("Bearer ${string}")]),
+					),
+				),
+			},
+			{ additionalProperties: true },
+		),
+	})
+	.resolve(
+		async ({
+			headers: { authorization, "Sec-WebSocket-Protocol": wsProtocol },
+			status,
+		}) => {
+			const auth =
+				authorization ??
+				(wsProtocol?.length === 2 &&
+				wsProtocol[0] === "kyoo" &&
+				wsProtocol[1].startsWith("Bearer ")
+					? wsProtocol[1]
+					: null);
+			const bearer = auth?.slice(7);
+			if (!bearer) {
+				return status(403, {
+					status: 403,
+					message: "No authorization header was found.",
+				});
+			}
+			try {
+				return await verifyJwt(bearer);
+			} catch (err) {
+				return status(403, {
+					status: 403,
+					message: "Invalid jwt. Verification vailed",
+					details: err,
+				});
+			}
+		},
+	);
+
+export const appWs = baseWs.ws("/ws", {
+	body: t.Union(
+		Object.entries(actionMap).map(([k, v]) =>
+			t.Intersect([t.Object({ action: t.Literal(k) }), v.body ?? t.Object({})]),
+		),
+	) as unknown as TObject<{ action: TString }>,
+	async open(ws) {
+		if (!ws.data.jwt.sub) {
+			ws.close(3000, "Unauthorized");
+		}
+	},
+	async message(ws, { action, ...body }) {
+		const handler = actionMap[action as keyof typeof actionMap];
+		for (const perm of handler.permissions ?? []) {
+			if (!ws.data.jwt.permissions.includes(perm)) {
+				return ws.close(3000, `Missing permission: '${perm}'.`);
+			}
+		}
+		await handler.message(ws as any, body as any);
 	},
 });
+
+type Ws = Parameters<NonNullable<Parameters<typeof baseWs.ws>[1]["open"]>>[0];
+function handler<Schema extends TSchema = TObject<{}>>(ret: {
+	body?: Schema;
+	permissions?: string[];
+	message: (ws: Ws, body: Schema["static"]) => void | Promise<void>;
+}) {
+	return ret;
+}