feat: self-hosted worker

packages/backend/server/package.json

@@ -67,6 +67,7 @@
     "graphql-scalars": "^1.23.0",
     "graphql-upload": "^17.0.0",
     "html-validate": "^8.20.1",
+    "htmlrewriter": "^0.0.12",
     "ioredis": "^5.3.2",
     "is-mobile": "^5.0.0",
     "keyv": "^5.0.0",
@@ -86,6 +87,7 @@
     "ses": "^1.4.1",
     "socket.io": "^4.7.5",
     "stripe": "^17.0.0",
+    "tldts": "^6.1.58",
     "ts-node": "^10.9.2",
     "typescript": "^5.6.3",
     "yjs": "patch:yjs@npm%3A13.6.18#~/.yarn/patches/yjs-npm-13.6.18-ad0d5f7c43.patch",

packages/backend/server/src/config/affine.self.ts

@@ -79,6 +79,8 @@ AFFiNE.use('captcha', {
   },
 });
 
+AFFiNE.use('worker');
+
 if (AFFiNE.deploy) {
   AFFiNE.mailer = {
     service: 'gmail',

packages/backend/server/src/fundamentals/error/def.ts

@@ -213,6 +213,10 @@ export const USER_FRIENDLY_ERRORS = {
     type: 'resource_not_found',
     message: 'Resource not found.',
   },
+  bad_request: {
+    type: 'bad_request',
+    message: 'Bad request.',
+  },
 
   // User Errors
   user_not_found: {

packages/backend/server/src/fundamentals/error/errors.gen.ts

@@ -22,6 +22,12 @@ export class NotFound extends UserFriendlyError {
   }
 }
 
+export class BadRequest extends UserFriendlyError {
+  constructor(message?: string) {
+    super('bad_request', 'bad_request', message);
+  }
+}
+
 export class UserNotFound extends UserFriendlyError {
   constructor(message?: string) {
     super('resource_not_found', 'user_not_found', message);
@@ -543,6 +549,7 @@ export enum ErrorNames {
   INTERNAL_SERVER_ERROR,
   TOO_MANY_REQUEST,
   NOT_FOUND,
+  BAD_REQUEST,
   USER_NOT_FOUND,
   USER_AVATAR_NOT_FOUND,
   EMAIL_ALREADY_USED,

packages/backend/server/src/plugins/index.ts

@@ -5,6 +5,7 @@ import './oauth';
 import './payment';
 import './redis';
 import './storage';
+import './worker';
 
 export {
   enablePlugin,

packages/backend/server/src/plugins/worker/config.ts

@@ -0,0 +1,15 @@
+import { defineStartupConfig, ModuleConfig } from '../../fundamentals/config';
+
+export interface WorkerStartupConfigurations {
+  allowedOrigin: string[];
+}
+
+declare module '../config' {
+  interface PluginsConfig {
+    worker: ModuleConfig<WorkerStartupConfigurations>;
+  }
+}
+
+defineStartupConfig('plugins.worker', {
+  allowedOrigin: [],
+});

packages/backend/server/src/plugins/worker/controller.ts

@@ -0,0 +1,277 @@
+import {
+  Controller,
+  Get,
+  Logger,
+  Options,
+  Post,
+  Req,
+  Res,
+} from '@nestjs/common';
+import type { Request, Response } from 'express';
+import { HTMLRewriter } from 'htmlrewriter';
+
+import { BadRequest, Cache, Config, URLHelper } from '../../fundamentals';
+import type { LinkPreviewRequest, LinkPreviewResponse } from './types';
+import {
+  appendUrl,
+  cloneHeader,
+  fixUrl,
+  getCorsHeaders,
+  isOriginAllowed,
+  isRefererAllowed,
+  OriginRules,
+  parseJson,
+  reduceUrls,
+} from './utils';
+
+@Controller('/api/worker')
+export class WorkerController {
+  private readonly logger = new Logger(WorkerController.name);
+  private readonly allowedOrigin: OriginRules;
+
+  constructor(
+    config: Config,
+    private readonly cache: Cache,
+    private readonly url: URLHelper
+  ) {
+    this.allowedOrigin = [
+      ...config.plugins.worker.allowedOrigin
+        .map(u => fixUrl(u)?.origin as string)
+        .filter(v => !!v),
+      url.origin,
+    ];
+  }
+
+  @Get('/image-proxy')
+  async imageProxy(@Req() req: Request, @Res() resp: Response) {
+    const origin = req.headers.origin ?? '';
+    const referer = req.headers.referer;
+    if (
+      (origin && !isOriginAllowed(origin, this.allowedOrigin)) ||
+      (referer && !isRefererAllowed(referer, this.allowedOrigin))
+    ) {
+      this.logger.error('Invalid Origin', 'ERROR', { origin, referer });
+      throw new BadRequest('Invalid header');
+    }
+    const url = new URL(req.url, this.url.baseUrl);
+    const imageURL = url.searchParams.get('url');
+    if (!imageURL) {
+      throw new BadRequest('Missing "url" parameter');
+    }
+
+    const targetURL = fixUrl(imageURL);
+    if (!targetURL) {
+      this.logger.error(`Invalid URL: ${url}`);
+      throw new BadRequest(`Invalid URL`);
+    }
+
+    const response = await fetch(
+      new Request(targetURL.toString(), {
+        method: 'GET',
+        headers: cloneHeader(req.headers),
+      })
+    );
+    if (response.ok) {
+      const contentType = response.headers.get('Content-Type');
+      const contentDisposition = response.headers.get('Content-Disposition');
+      if (contentType?.startsWith('image/')) {
+        return resp
+          .status(200)
+          .header({
+            'Access-Control-Allow-Origin': origin ?? 'null',
+            Vary: 'Origin',
+            'Access-Control-Allow-Methods': 'GET',
+            'Content-Type': contentType,
+            'Content-Disposition': contentDisposition,
+          })
+          .send(Buffer.from(await response.arrayBuffer()));
+      } else {
+        throw new BadRequest('Invalid content type');
+      }
+    } else {
+      this.logger.error('Failed to fetch image', {
+        origin,
+        url: imageURL,
+        status: resp.status,
+      });
+      throw new BadRequest('Failed to fetch image');
+    }
+  }
+
+  @Options('/link-preview')
+  linkPreviewOption(@Req() request: Request, @Res() resp: Response) {
+    const origin = request.headers.origin;
+    return resp
+      .status(200)
+      .header({
+        ...getCorsHeaders(origin),
+        'Access-Control-Allow-Methods': 'POST, OPTIONS',
+        'Access-Control-Allow-Headers': 'Content-Type',
+      })
+      .send();
+  }
+
+  @Post('/link-preview')
+  async linkPreview(
+    @Req() request: Request,
+    @Res() resp: Response
+  ): Promise<Response> {
+    const origin = request.headers.origin;
+    const referer = request.headers.referer;
+    if (
+      (origin && !isOriginAllowed(origin, this.allowedOrigin)) ||
+      (referer && !isRefererAllowed(referer, this.allowedOrigin))
+    ) {
+      this.logger.error('Invalid Origin', { origin, referer });
+      throw new BadRequest('Invalid header');
+    }
+
+    this.logger.debug('Received request', { origin, method: request.method });
+
+    const targetBody = parseJson<LinkPreviewRequest>(request.body);
+    const targetURL = fixUrl(targetBody?.url);
+    // not allow same site preview
+    if (!targetURL || isOriginAllowed(targetURL.origin, this.allowedOrigin)) {
+      this.logger.error('Invalid URL', { origin, url: targetBody?.url });
+      throw new BadRequest('Invalid URL');
+    }
+
+    this.logger.debug('Processing request', { origin, url: targetURL });
+
+    try {
+      const cachedResponse = await this.cache.get<string>(targetURL.toString());
+      if (cachedResponse) {
+        return resp
+          .status(200)
+          .header({
+            'content-type': 'application/json;charset=UTF-8',
+            ...getCorsHeaders(origin),
+          })
+          .send(cachedResponse);
+      }
+
+      const response = await fetch(targetURL, {
+        headers: cloneHeader(request.headers),
+      });
+      this.logger.error('Fetched URL', {
+        origin,
+        url: targetURL,
+        status: response.status,
+      });
+
+      const res: LinkPreviewResponse = {
+        url: response.url,
+        images: [],
+        videos: [],
+        favicons: [],
+      };
+      const baseUrl = new URL(request.url, this.url.baseUrl).toString();
+
+      if (response.body) {
+        const rewriter = new HTMLRewriter()
+          .on('meta', {
+            element(element) {
+              const property =
+                element.getAttribute('property') ??
+                element.getAttribute('name');
+              const content = element.getAttribute('content');
+              if (property && content) {
+                switch (property.toLowerCase()) {
+                  case 'og:title':
+                    res.title = content;
+                    break;
+                  case 'og:site_name':
+                    res.siteName = content;
+                    break;
+                  case 'og:description':
+                    res.description = content;
+                    break;
+                  case 'og:image':
+                    appendUrl(content, res.images);
+                    break;
+                  case 'og:video':
+                    appendUrl(content, res.videos);
+                    break;
+                  case 'og:type':
+                    res.mediaType = content;
+                    break;
+                  case 'description':
+                    if (!res.description) {
+                      res.description = content;
+                    }
+                }
+              }
+            },
+          })
+          .on('link', {
+            element(element) {
+              if (element.getAttribute('rel')?.toLowerCase().includes('icon')) {
+                appendUrl(element.getAttribute('href'), res.favicons);
+              }
+            },
+          })
+          .on('title', {
+            text(text) {
+              if (!res.title) {
+                res.title = text.text;
+              }
+            },
+          })
+          .on('img', {
+            element(element) {
+              appendUrl(element.getAttribute('src'), res.images);
+            },
+          })
+          .on('video', {
+            element(element) {
+              appendUrl(element.getAttribute('src'), res.videos);
+            },
+          });
+
+        await rewriter.transform(response).text();
+
+        res.images = await reduceUrls(baseUrl, res.images);
+
+        this.logger.error('Processed response with HTMLRewriter', {
+          origin,
+          url: response.url,
+        });
+      }
+
+      // fix favicon
+      {
+        // head default path of favicon
+        const faviconUrl = new URL('/favicon.ico', response.url);
+        const faviconResponse = await fetch(faviconUrl, { method: 'HEAD' });
+        if (faviconResponse.ok) {
+          appendUrl(faviconUrl.toString(), res.favicons);
+        }
+
+        res.favicons = await reduceUrls(baseUrl, res.favicons);
+      }
+
+      const json = JSON.stringify(res);
+      this.logger.debug('Sending response', {
+        origin,
+        url: res.url,
+        responseSize: json.length,
+      });
+
+      await this.cache.set(targetURL.toString(), res);
+      return resp
+        .status(200)
+        .header({
+          'content-type': 'application/json;charset=UTF-8',
+          ...getCorsHeaders(origin),
+        })
+        .send(json);
+    } catch (error) {
+      this.logger.error('Error fetching URL', {
+        origin,
+        url: targetURL,
+        error,
+      });
+      throw new BadRequest('Error fetching URL');
+    }
+  }
+}

packages/backend/server/src/plugins/worker/index.ts

@@ -0,0 +1,11 @@
+import './config';
+
+import { Plugin } from '../registry';
+import { WorkerController } from './controller';
+
+@Plugin({
+  name: 'worker',
+  controllers: [WorkerController],
+  if: config => config.isSelfhosted || config.node.dev || config.node.test,
+})
+export class WorkerModule {}

packages/backend/server/src/plugins/worker/types.ts

@@ -0,0 +1,16 @@
+export type LinkPreviewRequest = {
+  url: string;
+};
+
+export type LinkPreviewResponse = {
+  url: string;
+  title?: string;
+  siteName?: string;
+  description?: string;
+  images?: string[];
+  mediaType?: string;
+  contentType?: string;
+  charset?: string;
+  videos?: string[];
+  favicons?: string[];
+};