feat: Add JWT authentication and token-based access control

tercel · tercel · commit 5413122bf98e · 2025-02-08T15:57:19.000+08:00
diff --git a/.env.example b/.env.example
@@ -968,3 +968,9 @@ BUNDLE_EXECUTOR_ADDRESS=          # Address of the bundle executor contract
 DESK_EXCHANGE_PRIVATE_KEY=                  # Required for trading and cancelling orders
 DESK_EXCHANGE_NETWORK=                      # "mainnet" or "testnet
 
+# JWT
+JWT_ENABLED=
+JWT_SECRET_KEY=
+JWT_EXPIRED=
+JWT_USERNAME=
+JWT_PASSWORD=
diff --git a/packages/client-direct/package.json b/packages/client-direct/package.json
@@ -21,15 +21,16 @@
     "dependencies": {
         "@elizaos/core": "workspace:*",
         "@elizaos/plugin-image-generation": "workspace:*",
-        "@elizaos/plugin-tee-verifiable-log": "workspace:*",
         "@elizaos/plugin-tee-log": "workspace:*",
+        "@elizaos/plugin-tee-verifiable-log": "workspace:*",
         "@types/body-parser": "1.19.5",
         "@types/cors": "2.8.17",
         "@types/express": "5.0.0",
         "body-parser": "1.20.3",
         "cors": "2.8.5",
         "discord.js": "14.16.3",
         "express": "4.21.1",
+        "jsonwebtoken": "^9.0.2",
         "multer": "1.4.5-lts.1",
         "openai": "4.73.0"
     },
diff --git a/packages/client-direct/src/api.ts b/packages/client-direct/src/api.ts
@@ -8,47 +8,16 @@ import {
     type AgentRuntime,
     elizaLogger,
     getEnvVariable,
-    type UUID,
-    validateCharacterConfig,
     ServiceType,
     type Character,
+    settings,
 } from "@elizaos/core";
 
 import type { TeeLogQuery, TeeLogService } from "@elizaos/plugin-tee-log";
 import { REST, Routes } from "discord.js";
 import type { DirectClient } from ".";
-import { validateUuid } from "@elizaos/core";
-
-interface UUIDParams {
-    agentId: UUID;
-    roomId?: UUID;
-}
-
-function validateUUIDParams(
-    params: { agentId: string; roomId?: string },
-    res: express.Response
-): UUIDParams | null {
-    const agentId = validateUuid(params.agentId);
-    if (!agentId) {
-        res.status(400).json({
-            error: "Invalid AgentId format. Expected to be a UUID: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
-        });
-        return null;
-    }
-
-    if (params.roomId) {
-        const roomId = validateUuid(params.roomId);
-        if (!roomId) {
-            res.status(400).json({
-                error: "Invalid RoomId format. Expected to be a UUID: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
-            });
-            return null;
-        }
-        return { agentId, roomId };
-    }
-
-    return { agentId };
-}
+import { validateUUIDParams } from ".";
+import { md5, signToken } from "./auth";
 
 export function createApiRouter(
     agents: Map<string, AgentRuntime>,
@@ -65,14 +34,6 @@ export function createApiRouter(
         })
     );
 
-    router.get("/", (req, res) => {
-        res.send("Welcome, this is the REST API!");
-    });
-
-    router.get("/hello", (req, res) => {
-        res.json({ message: "Hello World!" });
-    });
-
     router.get("/agents", (req, res) => {
         const agentsList = Array.from(agents.values()).map((agent) => ({
             id: agent.agentId,
@@ -116,102 +77,6 @@ export function createApiRouter(
         });
     });
 
-    router.delete("/agents/:agentId", async (req, res) => {
-        const { agentId } = validateUUIDParams(req.params, res) ?? {
-            agentId: null,
-        };
-        if (!agentId) return;
-
-        const agent: AgentRuntime = agents.get(agentId);
-
-        if (agent) {
-            agent.stop();
-            directClient.unregisterAgent(agent);
-            res.status(204).json({ success: true });
-        } else {
-            res.status(404).json({ error: "Agent not found" });
-        }
-    });
-
-    router.post("/agents/:agentId/set", async (req, res) => {
-        const { agentId } = validateUUIDParams(req.params, res) ?? {
-            agentId: null,
-        };
-        if (!agentId) return;
-
-        let agent: AgentRuntime = agents.get(agentId);
-
-        // update character
-        if (agent) {
-            // stop agent
-            agent.stop();
-            directClient.unregisterAgent(agent);
-            // if it has a different name, the agentId will change
-        }
-
-        // stores the json data before it is modified with added data
-        const characterJson = { ...req.body };
-
-        // load character from body
-        const character = req.body;
-        try {
-            validateCharacterConfig(character);
-        } catch (e) {
-            elizaLogger.error(`Error parsing character: ${e}`);
-            res.status(400).json({
-                success: false,
-                message: e.message,
-            });
-            return;
-        }
-
-        // start it up (and register it)
-        try {
-            agent = await directClient.startAgent(character);
-            elizaLogger.log(`${character.name} started`);
-        } catch (e) {
-            elizaLogger.error(`Error starting agent: ${e}`);
-            res.status(500).json({
-                success: false,
-                message: e.message,
-            });
-            return;
-        }
-
-        if (process.env.USE_CHARACTER_STORAGE === "true") {
-            try {
-                const filename = `${agent.agentId}.json`;
-                const uploadDir = path.join(
-                    process.cwd(),
-                    "data",
-                    "characters"
-                );
-                const filepath = path.join(uploadDir, filename);
-                await fs.promises.mkdir(uploadDir, { recursive: true });
-                await fs.promises.writeFile(
-                    filepath,
-                    JSON.stringify(
-                        { ...characterJson, id: agent.agentId },
-                        null,
-                        2
-                    )
-                );
-                elizaLogger.info(
-                    `Character stored successfully at ${filepath}`
-                );
-            } catch (error) {
-                elizaLogger.error(
-                    `Failed to store character: ${error.message}`
-                );
-            }
-        }
-
-        res.json({
-            id: character.id,
-            character: character,
-        });
-    });
-
     router.get("/agents/:agentId/channels", async (req, res) => {
         const { agentId } = validateUUIDParams(req.params, res) ?? {
             agentId: null,
@@ -404,53 +269,14 @@ export function createApiRouter(
         }
     );
 
-    router.post("/agent/start", async (req, res) => {
-        const { characterPath, characterJson } = req.body;
-        console.log("characterPath:", characterPath);
-        console.log("characterJson:", characterJson);
-        try {
-            let character: Character;
-            if (characterJson) {
-                character = await directClient.jsonToCharacter(
-                    characterPath,
-                    characterJson
-                );
-            } else if (characterPath) {
-                character =
-                    await directClient.loadCharacterTryPath(characterPath);
-            } else {
-                throw new Error("No character path or JSON provided");
-            }
-            await directClient.startAgent(character);
-            elizaLogger.log(`${character.name} started`);
-
-            res.json({
-                id: character.id,
-                character: character,
-            });
-        } catch (e) {
-            elizaLogger.error(`Error parsing character: ${e}`);
-            res.status(400).json({
-                error: e.message,
-            });
-            return;
-        }
-    });
-
-    router.post("/agents/:agentId/stop", async (req, res) => {
-        const agentId = req.params.agentId;
-        console.log("agentId", agentId);
-        const agent: AgentRuntime = agents.get(agentId);
-
-        // update character
-        if (agent) {
-            // stop agent
-            agent.stop();
-            directClient.unregisterAgent(agent);
-            // if it has a different name, the agentId will change
-            res.json({ success: true });
+    router.post("/auth/login", async (req, res) => {
+        const { username, password } = req.body;
+        const valid = username === settings.JWT_USERNAME && password === md5(settings.JWT_PASSWORD);
+        if (valid) {
+            const token = signToken({ username });
+            res.json({ success: true, token: token });
         } else {
-            res.status(404).json({ error: "Agent not found" });
+            res.status(401).json({ error: "Invalid username or password" });
         }
     });
 
diff --git a/packages/client-direct/src/auth.ts b/packages/client-direct/src/auth.ts
@@ -0,0 +1,28 @@
+import jwt from 'jsonwebtoken';
+import { v4 as uuidv4 } from "uuid";
+import { settings } from "@elizaos/core";
+import crypto from 'crypto';
+
+export function md5(text: any) {
+    return crypto.createHash('md5').update(text).digest('hex');
+}
+
+export const signToken = (data: Record<string, any>, expiresIn: string | number = settings.JWT_EXPIRED): string => {
+    const _salt = uuidv4();
+    return jwt.sign({ ...data, _salt }, settings.JWT_SECRET_KEY, {
+        expiresIn: expiresIn
+    });
+};
+
+export const verifyToken = (authorization: string): Promise<any> => {
+    return new Promise((resolve, reject) => {
+        jwt.verify(authorization, settings.JWT_SECRET_KEY, async (err: any, decode: any) => {
+            if (err) {
+                reject(err);
+            } else {
+                resolve(decode);
+            }
+        });
+    });
+};
+
diff --git a/packages/client-direct/src/index.ts b/packages/client-direct/src/index.ts
@@ -22,12 +22,15 @@ import {
     settings,
     type IAgentRuntime,
     type TypeDatabaseAdapter,
+    type UUID,
+    validateUuid,
 } from "@elizaos/core";
 import { createApiRouter } from "./api.ts";
 import * as fs from "fs";
 import * as path from "path";
 import { createVerifiableLogApiRouter } from "./verifiable-log-api.ts";
 import { createManageApiRouter } from "./manage-api.ts";
+import { verifyToken } from "./auth.ts";
 import OpenAI from "openai";
 
 const storage = multer.diskStorage({
@@ -110,6 +113,64 @@ Response format should be formatted in a JSON block like this:
 \`\`\`
 `;
 
+async function verifyTokenMiddleware(req: any, res: any, next) {
+    // if JWT is not enabled, skip verification
+    if (!(settings.JWT_ENABLED && settings.JWT_ENABLED.toLowerCase() === 'true')) {
+        next();
+        return;
+    }
+
+    const url: string = req.url.split('?')[0];
+    if (url.indexOf('/manage/') !== 0) {
+    next();
+    } else {
+        try {
+            const { authorization } = req.headers;
+            if (!authorization) throw new Error('no token');
+            const token = await verifyToken(authorization.split(' ')[1]);
+            if (token) {
+                next();
+            } else {
+                throw new Error('fail to verify token');
+            }
+        } catch (err: any) {
+            res.status(401).json({ error: err.message });
+            return;
+        }
+    }
+};
+
+interface UUIDParams {
+    agentId: UUID;
+    roomId?: UUID;
+}
+
+export function validateUUIDParams(
+    params: { agentId: string; roomId?: string },
+    res: express.Response
+): UUIDParams | null {
+    const agentId = validateUuid(params.agentId);
+    if (!agentId) {
+        res.status(400).json({
+            error: "Invalid AgentId format. Expected to be a UUID: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
+        });
+        return null;
+    }
+
+    if (params.roomId) {
+        const roomId = validateUuid(params.roomId);
+        if (!roomId) {
+            res.status(400).json({
+                error: "Invalid RoomId format. Expected to be a UUID: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
+            });
+            return null;
+        }
+        return { agentId, roomId };
+    }
+
+    return { agentId };
+}
+
 export class DirectClient {
     public app: express.Application;
     private agents: Map<string, AgentRuntime>; // container management
@@ -128,6 +189,17 @@ export class DirectClient {
         this.app.use(bodyParser.json());
         this.app.use(bodyParser.urlencoded({ extended: true }));
 
+        this.app.use(verifyTokenMiddleware);
+
+
+        this.app.get("/", (req, res) => {
+            res.send("Welcome, this is the REST API!");
+        });
+
+        this.app.get("/hello", (req, res) => {
+            res.json({ message: "Hello World!" });
+        });
+
         // Serve both uploads and generated images
         this.app.use(
             "/media/uploads",
diff --git a/packages/client-direct/src/manage-api.ts b/packages/client-direct/src/manage-api.ts
