tokenMiddleware: pass token to streaming API

packages/ditto/app.ts

@@ -196,12 +196,19 @@ const ratelimit = every(
   rateLimitMiddleware(300, Time.minutes(5), false),
 );
 
+const socketTokenMiddleware = tokenMiddleware((c) => {
+  const token = c.req.header('sec-websocket-protocol');
+  if (token) {
+    return `Bearer ${token}`;
+  }
+});
+
 app.use('/api/*', metricsMiddleware, ratelimit, paginationMiddleware(), logiMiddleware);
 app.use('/.well-known/*', metricsMiddleware, ratelimit, logiMiddleware);
 app.use('/nodeinfo/*', metricsMiddleware, ratelimit, logiMiddleware);
 app.use('/oauth/*', metricsMiddleware, ratelimit, logiMiddleware);
 
-app.get('/api/v1/streaming', metricsMiddleware, ratelimit, streamingController);
+app.get('/api/v1/streaming', socketTokenMiddleware, metricsMiddleware, ratelimit, streamingController);
 app.get('/relay', metricsMiddleware, ratelimit, relayController);
 
 app.use(

packages/ditto/controllers/api/streaming.ts

@@ -12,13 +12,10 @@ import { z } from 'zod';
 import { type AppController } from '@/app.ts';
 import { getFeedPubkeys } from '@/queries.ts';
 import { hydrateEvents } from '@/storages/hydrate.ts';
-import { Storages } from '@/storages.ts';
-import { getTokenHash } from '@/utils/auth.ts';
 import { errorJson } from '@/utils/log.ts';
-import { bech32ToPubkey, Time } from '@/utils.ts';
+import { Time } from '@/utils.ts';
 import { renderReblog, renderStatus } from '@/views/mastodon/statuses.ts';
 import { renderNotification } from '@/views/mastodon/notifications.ts';
-import { HTTPException } from '@hono/hono/http-exception';
 
 /**
  * Streaming timelines/categories.
@@ -68,7 +65,7 @@ const limiter = new TTLCache<string, number>();
 const connections = new Set<WebSocket>();
 
 const streamingController: AppController = async (c) => {
-  const { conf, relay } = c.var;
+  const { conf, relay, user } = c.var;
   const upgrade = c.req.header('upgrade');
   const token = c.req.header('sec-websocket-protocol');
   const stream = streamSchema.optional().catch(undefined).parse(c.req.query('stream'));
@@ -78,11 +75,6 @@ const streamingController: AppController = async (c) => {
     return c.text('Please use websocket protocol', 400);
   }
 
-  const pubkey = token ? await getTokenPubkey(token) : undefined;
-  if (token && !pubkey) {
-    return c.json({ error: 'Invalid access token' }, 401);
-  }
-
   const ip = c.req.header('x-real-ip');
   if (ip) {
     const count = limiter.get(ip) ?? 0;
@@ -93,7 +85,8 @@ const streamingController: AppController = async (c) => {
 
   const { socket, response } = Deno.upgradeWebSocket(c.req.raw, { protocol: token, idleTimeout: 30 });
 
-  const policy = pubkey ? new MuteListPolicy(pubkey, await Storages.admin()) : undefined;
+  const pubkey = await user?.signer.getPublicKey();
+  const policy = pubkey ? new MuteListPolicy(pubkey, relay) : undefined;
 
   function send(e: StreamingEvent) {
     if (socket.readyState === WebSocket.OPEN) {
@@ -229,25 +222,4 @@ async function topicToFilter(
   }
 }
 
-async function getTokenPubkey(token: string): Promise<string | undefined> {
-  if (token.startsWith('token1')) {
-    const kysely = await Storages.kysely();
-    const tokenHash = await getTokenHash(token as `token1${string}`);
-
-    const row = await kysely
-      .selectFrom('auth_tokens')
-      .select('pubkey')
-      .where('token_hash', '=', tokenHash)
-      .executeTakeFirst();
-
-    if (!row) {
-      throw new HTTPException(401, { message: 'Invalid access token' });
-    }
-
-    return row.pubkey;
-  } else {
-    return bech32ToPubkey(token);
-  }
-}
-
 export { streamingController };

packages/mastoapi/middleware/tokenMiddleware.ts

@@ -13,9 +13,11 @@ import type { DittoEnv, DittoMiddleware } from '@ditto/router';
 import type { Context } from '@hono/hono';
 import type { User } from './User.ts';
 
-export function tokenMiddleware(): DittoMiddleware<{ user?: User }> {
+type CredentialsFn = (c: Context) => string | undefined;
+
+export function tokenMiddleware(fn?: CredentialsFn): DittoMiddleware<{ user?: User }> {
   return async (c, next) => {
-    const header = c.req.header('authorization');
+    const header = fn ? fn(c) : c.req.header('authorization');
 
     if (header) {
       const { relay, conf } = c.var;

packages/policies/MuteListPolicy.ts

@@ -15,7 +15,7 @@ export class MuteListPolicy implements NPolicy {
     }
 
     if (pubkeys.has(event.pubkey)) {
-      return ['OK', event.id, false, 'blocked: Your account has been deactivated.'];
+      return ['OK', event.id, false, 'blocked: account blocked'];
     }
 
     return ['OK', event.id, true, ''];