Add Auth for WebPage and APIs

src/app.ts

@@ -3,8 +3,10 @@ import express from 'express';
 import cors from 'cors';
 import http from 'http';
 import routes from './routes';
+import authRouter from './routes/auth';
 import { getPort } from './config';
 import logger from './utils/logger';
+import { authMiddleware } from './middleware/auth';
 
 const port = getPort();
 
@@ -18,11 +20,18 @@ const corsOptions = {
 app.use(cors(corsOptions));
 app.use(express.json());
 
-app.use('/api', routes);
+// 公开路由 - 无需认证
 app.get('/api', (_, res) => {
   res.status(200).json({ status: 'ok' });
 });
 
+// 认证路由 - 无需认证
+app.use('/api/auth', authRouter);
+
+// 受保护的路由 - 需要认证
+// 注意：这里必须在'/api/auth'之后定义，避免认证中间件拦截认证请求
+app.use('/api', authMiddleware, routes);
+
 server.listen(port, () => {
   logger.info(`Server is running on port ${port}`);
 });

src/config.ts

@@ -35,6 +35,9 @@ interface Config {
   API_ENDPOINTS: {
     SEARXNG: string;
   };
+  AUTH_PAGE: {
+    AUTH_SECRET: string;
+  };
 }
 
 type RecursivePartial<T> = {
@@ -64,6 +67,8 @@ export const getGeminiApiKey = () => loadConfig().MODELS.GEMINI.API_KEY;
 export const getSearxngApiEndpoint = () =>
   process.env.SEARXNG_API_URL || loadConfig().API_ENDPOINTS.SEARXNG;
 
+export const getAuthSecret = () => loadConfig().AUTH_PAGE.AUTH_SECRET;
+
 export const getOllamaApiEndpoint = () => loadConfig().MODELS.OLLAMA.API_URL;
 
 export const getCustomOpenaiApiKey = () =>

src/middleware/auth.ts

@@ -0,0 +1,54 @@
+import { Request, Response, NextFunction } from 'express';
+import jwt from 'jsonwebtoken';
+import { getAuthSecret } from '../config';
+import logger from '../utils/logger';
+
+// 扩展Express的Request类型，添加用户信息
+declare global {
+  namespace Express {
+    interface Request {
+      user?: {
+        authenticated: boolean;
+      };
+    }
+  }
+}
+
+/**
+ * 验证JWT令牌的中间件
+ */
+export const authMiddleware = (req: Request, res: Response, next: NextFunction) => {
+  // 从请求头获取令牌
+  const authHeader = req.headers.authorization;
+  const token = authHeader && authHeader.split(' ')[1]; // Bearer TOKEN格式
+  
+  // 如果没有令牌
+  if (!token) {
+    return res.status(401).json({ message: '未提供认证令牌' });
+  }
+  
+  try {
+    // 验证令牌
+    const secret = getAuthSecret();
+    const decoded = jwt.verify(token, secret);
+    
+    // 将用户信息添加到请求对象中
+    req.user = {
+      authenticated: true
+    };
+    
+    next();
+  } catch (err: any) {
+    logger.error(`令牌验证失败: ${err.message}`);
+    return res.status(403).json({ message: '无效的认证令牌' });
+  }
+};
+
+/**
+ * 生成JWT令牌
+ */
+export const generateToken = (): string => {
+  const secret = getAuthSecret();
+  // 令牌有效期为7天
+  return jwt.sign({ authenticated: true }, secret, { expiresIn: '7d' });
+}; 

src/routes/auth.ts

@@ -0,0 +1,46 @@
+import express from 'express';
+import logger from '../utils/logger';
+import { getAuthSecret } from '../config';
+import { generateToken } from '../middleware/auth';
+
+const router = express.Router();
+
+// 预设的密码，实际应用中应该从环境变量或配置文件中获取
+const AUTH_PASSWORD = getAuthSecret();
+
+router.post('/verify', (req, res) => {
+  try {
+    const { password } = req.body;
+    
+    if (!password) {
+      return res.status(400).json({ message: '密码不能为空' });
+    }
+    
+    if (password === AUTH_PASSWORD) {
+      // 生成JWT令牌
+      const token = generateToken();
+      return res.status(200).json({ 
+        message: '验证成功',
+        token: token 
+      });
+    } else {
+      return res.status(401).json({ message: '密码错误' });
+    }
+  } catch (err: any) {
+    logger.error(`认证错误: ${err.message}`);
+    return res.status(500).json({ message: '服务器错误' });
+  }
+});
+
+// 验证令牌是否有效的端点
+router.get('/verify-token', (req, res) => {
+  try {
+    // 令牌验证已经在中间件中完成，如果能到达这里，说明令牌有效
+    return res.status(200).json({ valid: true });
+  } catch (err: any) {
+    logger.error(`令牌验证错误: ${err.message}`);
+    return res.status(500).json({ message: '服务器错误' });
+  }
+});
+
+export default router;

src/websocket/websocketServer.ts

@@ -1,14 +1,51 @@
 import { WebSocketServer } from 'ws';
 import { handleConnection } from './connectionManager';
 import http from 'http';
-import { getPort } from '../config';
+import { getPort, getAuthSecret } from '../config';
 import logger from '../utils/logger';
+import jwt from 'jsonwebtoken';
+import { URL } from 'url';
 
 export const initServer = (
   server: http.Server<typeof http.IncomingMessage, typeof http.ServerResponse>,
 ) => {
   const port = getPort();
-  const wss = new WebSocketServer({ server });
+  const wss = new WebSocketServer({ noServer: true });
+
+  // 拦截升级请求
+  server.on('upgrade', (request, socket, head) => {
+    try {
+      // 从请求URL中获取令牌
+      const url = new URL(request.url!, `http://${request.headers.host}`);
+      const token = url.searchParams.get('token');
+
+      if (!token) {
+        logger.warn('WebSocket连接尝试未提供令牌');
+        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+        socket.destroy();
+        return;
+      }
+
+      // 验证令牌
+      const secret = getAuthSecret();
+      try {
+        jwt.verify(token, secret);
+        
+        // 令牌有效，允许升级连接
+        wss.handleUpgrade(request, socket, head, (ws) => {
+          wss.emit('connection', ws, request);
+        });
+      } catch (err) {
+        logger.warn(`无效的WebSocket认证令牌: ${err.message}`);
+        socket.write('HTTP/1.1 403 Forbidden\r\n\r\n');
+        socket.destroy();
+      }
+    } catch (err) {
+      logger.error(`WebSocket认证错误: ${err}`);
+      socket.write('HTTP/1.1 500 Internal Server Error\r\n\r\n');
+      socket.destroy();
+    }
+  });
 
   wss.on('connection', handleConnection);