feat(server): native safe fetch (#14931)

packages/backend/server/src/__tests__/app/ssrf.spec.ts

@@ -1,27 +0,0 @@
-import test from 'ava';
-
-import { readResponseBufferWithLimit } from '../../base';
-
-test('readResponseBufferWithLimit rejects timed out web streams without crashing', async t => {
-  const response = new Response(
-    new ReadableStream<Uint8Array>({
-      start(controller) {
-        controller.enqueue(new Uint8Array([1, 2, 3]));
-        queueMicrotask(() => {
-          controller.error(
-            new DOMException(
-              'The operation was aborted due to timeout',
-              'TimeoutError'
-            )
-          );
-        });
-      },
-    })
-  );
-
-  const error = await t.throwsAsync(
-    readResponseBufferWithLimit(response, 1024)
-  );
-
-  t.is(error?.name, 'TimeoutError');
-});

packages/backend/server/src/__tests__/e2e/create-app.ts

@@ -18,6 +18,7 @@ import {
   JobQueue,
   OneMB,
 } from '../../base';
+import { ThrottlerStorage } from '../../base/throttler';
 import { SocketIoAdapter } from '../../base/websocket';
 import { AuthGuard, AuthService } from '../../core/auth';
 import { Mailer } from '../../core/mail';
@@ -163,6 +164,10 @@ export class TestingApp extends NestApplication {
     return await this.create(MockUser, overrides);
   }
 
+  resetRateLimit() {
+    this.get(ThrottlerStorage, { strict: false }).storage.clear();
+  }
+
   async signup(overrides?: Partial<MockUserInput>) {
     const user = await this.create(MockUser, overrides);
     await this.login(user);
@@ -170,6 +175,7 @@ export class TestingApp extends NestApplication {
   }
 
   async login(user: MockedUser) {
+    this.resetRateLimit();
     return await this.POST('/api/auth/sign-in').send({
       email: user.email,
       password: user.password,
@@ -195,6 +201,7 @@ export class TestingApp extends NestApplication {
   }
 
   async logout(userId?: string) {
+    this.resetRateLimit();
     const res = await this.POST(
       '/api/auth/sign-out' + (userId ? `?user_id=${userId}` : '')
     ).expect(200);

packages/backend/server/src/__tests__/nestjs/throttler.spec.ts

@@ -92,7 +92,7 @@ test.before(async t => {
         throttle: {
           throttlers: {
             default: {
-              ttl: 60,
+              ttl: 60_000,
               limit: 120,
             },
           },
@@ -109,6 +109,7 @@ test.before(async t => {
 
 test.beforeEach(async t => {
   const { app } = t.context;
+  t.context.storage.storage.clear();
   await app.initTestingDB();
 });
 

packages/backend/server/src/__tests__/worker.e2e.ts

@@ -1,49 +1,79 @@
-import { LookupAddress } from 'node:dns';
-
+import serverNativeModule from '@affine/server-native';
 import type { ExecutionContext, TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';
 import type { Response } from 'supertest';
 
-import {
-  __resetDnsLookupForTests,
-  __setDnsLookupForTests,
-  type DnsLookup,
-} from '../base/utils/ssrf';
-import { createTestingApp, TestingApp } from './utils';
+import type { TestingApp } from './utils';
 
 type TestContext = {
   app: TestingApp;
 };
 const test = ava as TestFn<TestContext>;
 
-const LookupAddressStub = (async (_hostname, options) => {
-  const result = [{ address: '76.76.21.21', family: 4 }] as LookupAddress[];
-  const isOptions = options && typeof options === 'object';
-  if (isOptions && 'all' in options && options.all) {
-    return result;
+let safeFetchStub: Sinon.SinonStub | undefined;
+let safeFetchHandler:
+  | ((request: { url: string; method?: 'get' | 'head' }) => {
+      status?: number;
+      finalUrl?: string;
+      headers?: Record<string, string>;
+      body?: Buffer | string;
+    })
+  | undefined;
+
+const stubSafeFetch = (
+  handler: (request: { url: string; method?: 'get' | 'head' }) => {
+    status?: number;
+    finalUrl?: string;
+    headers?: Record<string, string>;
+    body?: Buffer | string;
   }
-  return result[0];
-}) as DnsLookup;
+) => {
+  safeFetchHandler = handler;
+  return {
+    restore() {
+      safeFetchHandler = undefined;
+    },
+  };
+};
 
 test.before(async t => {
   // @ts-expect-error test
   env.DEPLOYMENT_TYPE = 'selfhosted';
 
-  // Avoid relying on real DNS during tests. SSRF protection uses dns.lookup().
-  __setDnsLookupForTests(LookupAddressStub);
+  safeFetchStub = Sinon.stub(serverNativeModule, 'safeFetch').callsFake(
+    async request => {
+      if (!safeFetchHandler) {
+        throw new Error('Unexpected safeFetch call');
+      }
+      const nativeRequest = request as {
+        url: string;
+        method?: 'get' | 'head';
+      };
+      const response = safeFetchHandler(nativeRequest);
+      return {
+        status: response.status ?? 200,
+        finalUrl: response.finalUrl ?? nativeRequest.url,
+        headers: response.headers ?? {},
+        body: Buffer.isBuffer(response.body)
+          ? response.body
+          : Buffer.from(response.body ?? ''),
+      };
+    }
+  );
 
+  const { createTestingApp } = await import('./utils');
   const app = await createTestingApp();
 
   t.context.app = app;
 });
 
 test.afterEach.always(() => {
-  Sinon.restore();
+  safeFetchHandler = undefined;
 });
 
 test.after.always(async t => {
-  __resetDnsLookupForTests();
+  safeFetchStub?.restore();
   await t.context.app.close();
 });
 
@@ -128,15 +158,13 @@ test('should proxy image', async t => {
       'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jfJ8AAAAASUVORK5CYII=',
       'base64'
     );
-    const fakeResponse = new Response(fakeBuffer, {
-      status: 200,
+    const fetchSpy = stubSafeFetch(() => ({
+      body: fakeBuffer,
       headers: {
         'content-type': 'image/png',
         'content-disposition': 'inline',
       },
-    });
-
-    const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeResponse);
+    }));
     try {
       await assertAndSnapshot(
         `/api/worker/image-proxy?url=${imageUrl}`,
@@ -146,6 +174,42 @@ test('should proxy image', async t => {
       fetchSpy.restore();
     }
   }
+
+  {
+    const invalidImageUrl = `http://example.com/not-image-${Date.now()}.png`;
+    const invalidFetchSpy = stubSafeFetch(() => ({
+      body: 'not an image',
+      headers: { 'content-type': 'image/png' },
+    }));
+    try {
+      await t.context.app
+        .GET(`/api/worker/image-proxy?url=${invalidImageUrl}`)
+        .set('Origin', 'http://localhost:3010')
+        .send()
+        .expect(400);
+    } finally {
+      invalidFetchSpy.restore();
+    }
+
+    const validImageUrl = `http://example.com/valid-image-${Date.now()}.png`;
+    const fakeBuffer = Buffer.from(
+      'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jfJ8AAAAASUVORK5CYII=',
+      'base64'
+    );
+    const validFetchSpy = stubSafeFetch(() => ({
+      body: fakeBuffer,
+      headers: { 'content-type': 'image/png' },
+    }));
+    try {
+      await t.context.app
+        .GET(`/api/worker/image-proxy?url=${validImageUrl}`)
+        .set('Origin', 'http://localhost:3010')
+        .send()
+        .expect(200);
+    } finally {
+      validFetchSpy.restore();
+    }
+  }
 });
 
 test('should preview link', async t => {
@@ -190,7 +254,8 @@ test('should preview link', async t => {
   );
 
   {
-    const fakeHTML = new Response(`
+    const pageUrl = `http://external.com/page-${Date.now()}`;
+    const fakeHTML = `
         <html>
           <head>
             <meta property="og:title" content="Test Title" />
@@ -201,13 +266,18 @@ test('should preview link', async t => {
             <title>Fallback Title</title>
           </body>
         </html>
-      `);
+      `;
 
-    Object.defineProperty(fakeHTML, 'url', {
-      value: 'http://example.com/page',
+    const fetchSpy = stubSafeFetch(request => {
+      if (request.url.includes('/favicon.ico')) {
+        return { status: 204, finalUrl: request.url };
+      }
+      return {
+        body: fakeHTML,
+        finalUrl: 'http://example.com/page',
+        headers: { 'content-type': 'text/html;charset=UTF-8' },
+      };
     });
-
-    const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeHTML);
     try {
       await assertAndSnapshot(
         '/api/worker/link-preview',
@@ -215,7 +285,7 @@ test('should preview link', async t => {
         {
           status: 200,
           method: 'POST',
-          body: { url: 'http://external.com/page' },
+          body: { url: pageUrl },
         }
       );
     } finally {
@@ -244,6 +314,7 @@ test('should preview link', async t => {
     ];
 
     for (const { content, charset } of encoded) {
+      const pageUrl = `http://example.com/${charset}-${Date.now()}`;
       const before = Buffer.from(`<html>
           <head>
             <meta http-equiv="Content-Type" content="text/html; charset=${charset}" />
@@ -253,13 +324,18 @@ test('should preview link', async t => {
           </head>
         </html>
       `);
-      const fakeHTML = new Response(Buffer.concat([before, encoded, after]));
+      const fakeHTML = Buffer.concat([before, encoded, after]);
 
-      Object.defineProperty(fakeHTML, 'url', {
-        value: `http://example.com/${charset}`,
+      const fetchSpy = stubSafeFetch(request => {
+        if (request.url.includes('/favicon.ico')) {
+          return { status: 204, finalUrl: request.url };
+        }
+        return {
+          body: fakeHTML,
+          finalUrl: `http://example.com/${charset}`,
+          headers: { 'content-type': `text/html;charset=${charset}` },
+        };
       });
-
-      const fetchSpy = Sinon.stub(global, 'fetch').resolves(fakeHTML);
       try {
         await assertAndSnapshot(
           '/api/worker/link-preview',
@@ -267,7 +343,7 @@ test('should preview link', async t => {
           {
             status: 200,
             method: 'POST',
-            body: { url: `http://example.com/${charset}` },
+            body: { url: pageUrl },
           }
         );
       } finally {