From 42b5ef7a8b12cb277f6255b5ca520a8828cb03ad Mon Sep 17 00:00:00 2001
From: liuyi <forehalo@gmail.com>
Date: Wed, 14 Aug 2024 16:05:55 +0800
Subject: [PATCH] fix(server): redirect to setup page if not initialized
 (#7871)

---
 packages/backend/server/src/app.module.ts     | 27 ++++++++++----
 .../backend/server/src/core/setup/index.ts    | 35 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/packages/backend/server/src/app.module.ts b/packages/backend/server/src/app.module.ts
index ac80a578013ae..11be399735705 100644
--- a/packages/backend/server/src/app.module.ts
+++ b/packages/backend/server/src/app.module.ts
@@ -1,10 +1,13 @@
-import { join } from 'node:path';
+import { join, resolve } from 'node:path';
+import { fileURLToPath } from 'node:url';
 
 import {
   DynamicModule,
   ForwardReference,
   Logger,
+  MiddlewareConsumer,
   Module,
+  NestModule,
 } from '@nestjs/common';
 import { ScheduleModule } from '@nestjs/schedule';
 import { ServeStaticModule } from '@nestjs/serve-static';
@@ -16,7 +19,7 @@ import { ADD_ENABLED_FEATURES, ServerConfigModule } from './core/config';
 import { DocModule } from './core/doc';
 import { FeatureModule } from './core/features';
 import { QuotaModule } from './core/quota';
-import { CustomSetupModule } from './core/setup';
+import { CustomSetupModule, SetupMiddleware } from './core/setup';
 import { StorageModule } from './core/storage';
 import { SyncModule } from './core/sync';
 import { UserModule } from './core/user';
@@ -135,16 +138,26 @@ export class AppModuleBuilder {
   }
 
   compile() {
+    const configure = (consumer: MiddlewareConsumer) => {
+      if (this.config.isSelfhosted) {
+        consumer.apply(SetupMiddleware).forRoutes('*');
+      }
+    };
+
     @Module({
       imports: this.modules,
       controllers: this.config.isSelfhosted ? [] : [AppController],
     })
-    class AppModule {}
+    class AppModule implements NestModule {
+      configure = configure;
+    }
 
     return AppModule;
   }
 }
 
+const pwd = resolve(fileURLToPath(import.meta.url), '../../');
+
 function buildAppModule() {
   AFFiNE = mergeConfigOverride(AFFiNE);
   const factor = new AppModuleBuilder(AFFiNE);
@@ -179,12 +192,12 @@ function buildAppModule() {
       config => config.isSelfhosted,
       CustomSetupModule,
       ServeStaticModule.forRoot({
-        rootPath: join('/app', 'static'),
-        exclude: ['/admin*'],
+        rootPath: join(pwd, 'static', 'admin'),
+        renderPath: /^\/admin\/?/,
       }),
       ServeStaticModule.forRoot({
-        rootPath: join('/app', 'static', 'admin'),
-        serveRoot: '/admin',
+        rootPath: join(pwd, 'static'),
+        renderPath: '*',
       })
     );
 
diff --git a/packages/backend/server/src/core/setup/index.ts b/packages/backend/server/src/core/setup/index.ts
index 56fe7cfd790a5..aba4796d01f66 100644
--- a/packages/backend/server/src/core/setup/index.ts
+++ b/packages/backend/server/src/core/setup/index.ts
@@ -1,9 +1,42 @@
-import { Module } from '@nestjs/common';
+import { Injectable, Module, NestMiddleware } from '@nestjs/common';
+import { PrismaClient } from '@prisma/client';
+import type { Request, Response } from 'express';
 
 import { AuthModule } from '../auth';
 import { UserModule } from '../user';
 import { CustomSetupController } from './controller';
 
+@Injectable()
+export class SetupMiddleware implements NestMiddleware {
+  private initialized: boolean | null = null;
+  constructor(private readonly db: PrismaClient) {}
+
+  use(req: Request, res: Response, next: (error?: Error | any) => void) {
+    // never throw
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
+    this.allow().then(allowed => {
+      if (allowed) {
+        next();
+      } else if (!req.path.startsWith('/admin/setup')) {
+        res.redirect('/admin/setup');
+      }
+    });
+  }
+
+  async allow() {
+    try {
+      if (this.initialized === null) {
+        this.initialized = (await this.db.user.count()) > 0;
+      }
+    } catch (e) {
+      // avoid block the whole app
+      return true;
+    }
+
+    return this.initialized;
+  }
+}
+
 @Module({
   imports: [AuthModule, UserModule],
   controllers: [CustomSetupController],