Rewrite streaming controller to use hono websocket helper

src/app.ts

@@ -1,4 +1,4 @@
-import { Context, Env as HonoEnv, Handler, Hono, Input as HonoInput, MiddlewareHandler } from '@hono/hono';
+import { Context, Env as HonoEnv, Handler, Hono, MiddlewareHandler } from '@hono/hono';
 import { cors } from '@hono/hono/cors';
 import { serveStatic } from '@hono/hono/deno';
 import { logger } from '@hono/hono/logger';
@@ -134,7 +134,7 @@ interface AppEnv extends HonoEnv {
 
 type AppContext = Context<AppEnv>;
 type AppMiddleware = MiddlewareHandler<AppEnv>;
-type AppController = Handler<AppEnv, any, HonoInput, Response | Promise<Response>>;
+type AppController = Handler<AppEnv, any>;
 
 const app = new Hono<AppEnv>({ strict: false });
 

src/controllers/api/streaming.ts

@@ -2,13 +2,14 @@ import { NostrEvent, NostrFilter } from '@nostrify/nostrify';
 import Debug from '@soapbox/stickynotes/debug';
 import { z } from 'zod';
 
-import { type AppController } from '@/app.ts';
+import { AppController } from '@/app.ts';
 import { Conf } from '@/config.ts';
 import { DittoDB } from '@/db/DittoDB.ts';
 import { MuteListPolicy } from '@/policies/MuteListPolicy.ts';
 import { getFeedPubkeys } from '@/queries.ts';
 import { hydrateEvents } from '@/storages/hydrate.ts';
 import { Storages } from '@/storages.ts';
+import { upgradeWebSocket } from '@/utils/websocket.ts';
 import { bech32ToPubkey } from '@/utils.ts';
 import { renderReblog, renderStatus } from '@/views/mastodon/statuses.ts';
 import { renderNotification } from '@/views/mastodon/notifications.ts';
@@ -51,14 +52,15 @@ const streamingController: AppController = async (c) => {
     return c.json({ error: 'Invalid access token' }, 401);
   }
 
-  const { socket, response } = Deno.upgradeWebSocket(c.req.raw, { protocol: token, idleTimeout: 30 });
-
   const store = await Storages.db();
   const pubsub = await Storages.pubsub();
 
   const policy = pubkey ? new MuteListPolicy(pubkey, await Storages.admin()) : undefined;
 
-  function send(name: string, payload: object) {
+  return upgradeWebSocket((c) => {
+    let socket: WebSocket;
+
+    function send(name: string, payload: object): void {
       if (socket.readyState === WebSocket.OPEN) {
         debug('send', name, JSON.stringify(payload));
         socket.send(JSON.stringify({
@@ -69,7 +71,11 @@ const streamingController: AppController = async (c) => {
       }
     }
 
-  async function sub(type: string, filters: NostrFilter[], render: (event: NostrEvent) => Promise<unknown>) {
+    async function sub(
+      type: string,
+      filters: NostrFilter[],
+      render: (event: NostrEvent) => Promise<unknown>,
+    ): Promise<void> {
       try {
         for await (const msg of pubsub.req(filters, { signal: controller.signal })) {
           if (msg[0] === 'EVENT') {
@@ -96,7 +102,10 @@ const streamingController: AppController = async (c) => {
       }
     }
 
-  socket.onopen = async () => {
+    return {
+      async onOpen(_event, ws) {
+        socket = ws.raw as WebSocket;
+
         if (!stream) return;
         const topicFilter = await topicToFilter(stream, c.req.query(), pubkey);
 
@@ -117,13 +126,17 @@ const streamingController: AppController = async (c) => {
           });
           return;
         }
-  };
-
-  socket.onclose = () => {
+      },
+      onMessage(event, ws) {
+        if (typeof event.data !== 'string') {
+          ws.close(1003, 'Unsupported data type');
+        }
+      },
+      onClose(_event) {
         controller.abort();
+      },
     };
-
-  return response;
+  }, { protocol: token, idleTimeout: 30 });
 };
 
 async function topicToFilter(

src/utils/websocket.ts

@@ -0,0 +1,49 @@
+import { Context, MiddlewareHandler } from '@hono/hono';
+import { WSContext, WSEvents, WSReadyState } from '@hono/hono/ws';
+
+type UpgradeWebSocket = (
+  createEvents: (c: Context) => WSEvents | Promise<WSEvents>,
+  options?: Deno.UpgradeWebSocketOptions,
+) => MiddlewareHandler<
+  any,
+  string,
+  {
+    outputFormat: 'ws';
+  }
+>;
+
+/**
+ * A modified version of Hono's WebSocket Helper, just to support Deno's `options` object.
+ *
+ * See: https://github.com/honojs/hono/issues/2997
+ *
+ * (If that issue gets fixed, we can remove this code.)
+ */
+export const upgradeWebSocket: UpgradeWebSocket = (createEvents, options) => async (c, next) => {
+  if (c.req.header('upgrade') !== 'websocket') {
+    return await next();
+  }
+
+  const events = await createEvents(c);
+  const { response, socket } = Deno.upgradeWebSocket(c.req.raw, options);
+
+  const wsContext: WSContext = {
+    binaryType: 'arraybuffer',
+    close: (code, reason) => socket.close(code, reason),
+    get protocol() {
+      return socket.protocol;
+    },
+    raw: socket,
+    get readyState() {
+      return socket.readyState as WSReadyState;
+    },
+    url: socket.url ? new URL(socket.url) : null,
+    send: (source) => socket.send(source),
+  };
+  socket.onopen = (evt) => events.onOpen?.(evt, wsContext);
+  socket.onmessage = (evt) => events.onMessage?.(evt, wsContext);
+  socket.onclose = (evt) => events.onClose?.(evt, wsContext);
+  socket.onerror = (evt) => events.onError?.(evt, wsContext);
+
+  return response;
+};