Add WEBSOCKET_ONLY_MODE to disable REST routes in websocket deployment

.infra/index.ts

@@ -370,6 +370,7 @@ if (isAdhocEnv) {
       env: [
         nodeOptions(wsMemory),
         { name: 'ENABLE_SUBSCRIPTIONS', value: 'true' },
+        { name: 'WEBSOCKET_ONLY_MODE', value: 'true' },
         ...commonEnv,
         ...jwtEnv,
         {

__tests__/websocketOnlyMode.ts

@@ -0,0 +1,77 @@
+import { FastifyInstance } from 'fastify';
+import request from 'supertest';
+import appFunc from '../src/index';
+import createOrGetConnection from '../src/db';
+
+let app: FastifyInstance;
+
+beforeAll(async () => {
+  await createOrGetConnection();
+});
+
+afterEach(async () => {
+  if (app) {
+    await app.close();
+  }
+});
+
+describe('websocket only mode', () => {
+  it('should expose all routes when WEBSOCKET_ONLY_MODE is not set', async () => {
+    delete process.env.WEBSOCKET_ONLY_MODE;
+    app = await appFunc();
+    await app.listen({ port: 0, host: '0.0.0.0' });
+
+    // GraphQL should be available
+    const graphqlRes = await request(app.server)
+      .post('/graphql')
+      .send({ query: '{ __typename }' })
+      .expect(200);
+    expect(graphqlRes.body.data.__typename).toBe('Query');
+
+    // REST routes should be available
+    await request(app.server).get('/health').expect(200);
+
+    // /v1 compatibility routes should be available
+    await request(app.server).get('/v1/users/me').expect(401); // Expects auth
+  });
+
+  it('should only expose GraphQL and health endpoints when WEBSOCKET_ONLY_MODE is true', async () => {
+    process.env.WEBSOCKET_ONLY_MODE = 'true';
+    app = await appFunc();
+    await app.listen({ port: 0, host: '0.0.0.0' });
+
+    // GraphQL should still be available
+    const graphqlRes = await request(app.server)
+      .post('/graphql')
+      .send({ query: '{ __typename }' })
+      .expect(200);
+    expect(graphqlRes.body.data.__typename).toBe('Query');
+
+    // Health endpoints should still be available
+    await request(app.server).get('/health').expect(200);
+    await request(app.server).get('/liveness').expect(200);
+
+    // REST routes should NOT be available
+    await request(app.server).get('/v1/users/me').expect(404);
+
+    // Icon proxy should NOT be available
+    await request(app.server).get('/icon?url=example.com&size=64').expect(404);
+
+    // Routes should NOT be available
+    await request(app.server).get('/rss/f/popular').expect(404);
+  });
+
+  it('should support GraphQL subscriptions when WEBSOCKET_ONLY_MODE is true and ENABLE_SUBSCRIPTIONS is true', async () => {
+    process.env.WEBSOCKET_ONLY_MODE = 'true';
+    process.env.ENABLE_SUBSCRIPTIONS = 'true';
+    app = await appFunc();
+    await app.listen({ port: 0, host: '0.0.0.0' });
+
+    // GraphQL should be available with subscription support
+    const graphqlRes = await request(app.server)
+      .post('/graphql')
+      .send({ query: '{ __typename }' })
+      .expect(200);
+    expect(graphqlRes.body.data.__typename).toBe('Query');
+  });
+});

src/index.ts

@@ -64,6 +64,7 @@ export default async function app(
 ): Promise<FastifyInstance> {
   let isTerminating = false;
   const isProd = process.env.NODE_ENV === 'production';
+  const isWebsocketOnly = process.env.WEBSOCKET_ONLY_MODE === 'true';
   const connection = await runInRootSpan(
     'createOrGetConnection',
     createOrGetConnection,
@@ -337,86 +338,89 @@ export default async function app(
     });
   }
 
-  app.register(
-    async (instance) => {
-      await compatibility(instance, connection);
-    },
-    { prefix: '/v1' },
-  );
-  app.register(proxy, {
-    upstream: 'https://www.google.com/s2/favicons',
-    prefix: '/icon',
-    logLevel: 'warn',
-    replyOptions: {
-      queryString: (search, reqUrl, req) => {
-        const reqSearchParams = new URLSearchParams(
-          req.query as { url: string; size: string },
-        );
-        const proxySearchParams = new URLSearchParams();
-
-        proxySearchParams.set('domain', reqSearchParams.get('url') ?? '');
-        proxySearchParams.set('sz', reqSearchParams.get('size') ?? '');
-        return proxySearchParams.toString();
+  // Skip all REST routes when in websocket-only mode
+  if (!isWebsocketOnly) {
+    app.register(
+      async (instance) => {
+        await compatibility(instance, connection);
       },
-    },
-    preValidation: async (req: FastifyRequest, res) => {
-      const { url, size } = req.query as { url: string; size: string };
-      if (!url || !size) {
-        res.status(400).send({ error: 'url and size are required' });
-      }
-    },
-    preHandler: async (req, res) => {
-      res.helmet({
-        crossOriginResourcePolicy: {
-          policy: 'cross-origin',
+      { prefix: '/v1' },
+    );
+    app.register(proxy, {
+      upstream: 'https://www.google.com/s2/favicons',
+      prefix: '/icon',
+      logLevel: 'warn',
+      replyOptions: {
+        queryString: (search, reqUrl, req) => {
+          const reqSearchParams = new URLSearchParams(
+            req.query as { url: string; size: string },
+          );
+          const proxySearchParams = new URLSearchParams();
+
+          proxySearchParams.set('domain', reqSearchParams.get('url') ?? '');
+          proxySearchParams.set('sz', reqSearchParams.get('size') ?? '');
+          return proxySearchParams.toString();
         },
-      });
-    },
-  });
+      },
+      preValidation: async (req: FastifyRequest, res) => {
+        const { url, size } = req.query as { url: string; size: string };
+        if (!url || !size) {
+          res.status(400).send({ error: 'url and size are required' });
+        }
+      },
+      preHandler: async (req, res) => {
+        res.helmet({
+          crossOriginResourcePolicy: {
+            policy: 'cross-origin',
+          },
+        });
+      },
+    });
 
-  const letterProxy: FastifyRegisterOptions<FastifyHttpProxyOptions> = {
-    upstream:
-      'https://media.daily.dev/image/upload/s--zchx8x3n--/f_auto,q_auto/v1731056371/webapp/shortcut-placeholder',
-    preHandler: async (req: FastifyRequest, res: FastifyReply) => {
-      res.helmet({
-        crossOriginResourcePolicy: {
-          policy: 'cross-origin',
-        },
-      });
-    },
-    logLevel: 'warn',
-  };
+    const letterProxy: FastifyRegisterOptions<FastifyHttpProxyOptions> = {
+      upstream:
+        'https://media.daily.dev/image/upload/s--zchx8x3n--/f_auto,q_auto/v1731056371/webapp/shortcut-placeholder',
+      preHandler: async (req: FastifyRequest, res: FastifyReply) => {
+        res.helmet({
+          crossOriginResourcePolicy: {
+            policy: 'cross-origin',
+          },
+        });
+      },
+      logLevel: 'warn',
+    };
 
-  app.register(proxy, {
-    prefix: 'lettericons',
-    ...letterProxy,
-  });
-  app.register(proxy, {
-    prefix: '/lettericons/:word',
-    ...letterProxy,
-  });
+    app.register(proxy, {
+      prefix: 'lettericons',
+      ...letterProxy,
+    });
+    app.register(proxy, {
+      prefix: '/lettericons/:word',
+      ...letterProxy,
+    });
 
-  app.register(proxy, {
-    prefix: '/freyja',
-    httpMethods: ['POST'],
-    upstream: `${process.env.FREYJA_ORIGIN}/api`,
-    preHandler: async (req: FastifyRequest, res: FastifyReply) => {
-      res.helmet({
-        crossOriginResourcePolicy: {
-          policy: 'cross-origin',
-        },
-      });
+    app.register(proxy, {
+      prefix: '/freyja',
+      httpMethods: ['POST'],
+      upstream: `${process.env.FREYJA_ORIGIN}/api`,
+      preHandler: async (req: FastifyRequest, res: FastifyReply) => {
+        res.helmet({
+          crossOriginResourcePolicy: {
+            policy: 'cross-origin',
+          },
+        });
 
-      const regex = new RegExp('^/freyja/sessions/[^/]+/transition$');
-      if (!regex.test(req.url)) {
-        res.status(404).send();
-        return;
-      }
-    },
-    logLevel: 'warn',
-  });
+        const regex = new RegExp('^/freyja/sessions/[^/]+/transition$');
+        if (!regex.test(req.url)) {
+          res.status(404).send();
+          return;
+        }
+      },
+      logLevel: 'warn',
+    });
 
-  app.register(routes, { prefix: '/' });
+    app.register(routes, { prefix: '/' });
+  }
 
   return app;
 }

src/types.ts

@@ -19,6 +19,7 @@ declare global {
       TYPEORM_DATABASE: string;
       HEIMDALL_ORIGIN: string;
       ENABLE_PRIVATE_ROUTES: string;
+      WEBSOCKET_ONLY_MODE?: string;
       ACCESS_SECRET: string;
       ALLOCATION_QUEUE_CONCURRENCY: string;
       QUEUE_CONCURRENCY: string;