add signature verification for agent startup

0glabs · Jan 22, 2025 · a432fa7 · a432fa7
1 parent 4f6fc94
commit a432fa7
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 18 deletions.
diff --git a/agent/src/agent.ts b/agent/src/agent.ts
@@ -134,11 +134,22 @@ export function parseArguments(): ParsedArguments {
                 type: "string",
                 description: "Comma separated list of NFT token IDs to load configuration from",
             })
+            .option("proof", {
+                type: "string",
+                description: "Proof of ownership of the NFT",
+            })
             .option("dir", {
                 type: "string",
                 description: "Base directory storing the agent data",
                 default: ""
             })
+            .check((argv) => {
+                if (argv.token && !argv.proof) {
+                    elizaLogger.error("The 'proof' option is required when 'token' is provided.");
+                    throw new Error("The 'proof' option is required when 'token' is provided.");
+                }
+                return true;
+            })
             .parseSync();
     } catch (error) {
         elizaLogger.error("Error parsing arguments:", error);
@@ -1024,7 +1035,7 @@ export const checkPortAvailable = (port: number): Promise<boolean> => {
 };
 
 
-export async function loadFromNFT(tokenId: string, baseDir: string = ""): Promise<Character[]> {
+export async function loadFromNFT(tokenId: string, proof: string, baseDir: string = ""): Promise<Character[]> {
     const agentNFTClient = new AgentNFTClient(baseDir);
 
     const name = await agentNFTClient.getNFTName();
@@ -1041,7 +1052,7 @@ export async function loadFromNFT(tokenId: string, baseDir: string = ""): Promis
     elizaLogger.info("tokenData", tokenData);
 
     elizaLogger.info("Validating token ownership and permissions...");
-    const isValid = await agentNFTClient.validateToken(tokenData);
+    const isValid = await agentNFTClient.validateToken(tokenData, proof);
     if (!isValid) {
         elizaLogger.error("Token validation failed");
         throw new Error(`Token validation failed`);

diff --git a/agent/src/agentNFTClient.ts b/agent/src/agentNFTClient.ts
@@ -6,8 +6,7 @@ import { TokenData, AgentMetadata } from './types';
 import { AgentNFT } from './contracts/AgentNFT';
 import { AgentNFT__factory } from './contracts/factories/AgentNFT__factory';
 import { Indexer, ZgFile } from '@0glabs/0g-ts-sdk';
-
-const NUM_AGENT_HASHES = 2;
+import { createVerify } from 'crypto';
 
 export class AgentNFTClient {
     private provider: ethers.Provider;
@@ -123,17 +122,17 @@ export class AgentNFTClient {
         }
     }
 
-    async validateToken(tokenData: TokenData): Promise<boolean> {
+    async validateToken(tokenData: TokenData, proof: string): Promise<boolean> {
         try {
             const tokenOwner = tokenData.owner.toLowerCase();
-            const tokenOwnerPrivateKey = process.env.ZEROG_PRIVATE_KEY?.toLowerCase();
-            const claimedTokenOwner = new ethers.Wallet(tokenOwnerPrivateKey).address.toLowerCase();
-            if (tokenOwner === claimedTokenOwner) {
-                return true;
-            } else {
-                elizaLogger.error(`Token ${tokenData.tokenId} is not owned by ${claimedTokenOwner}, token owner is ${tokenOwner}`);
-                return false;
-            }
+            elizaLogger.info("proof", proof);
+            // parse proof and verify
+            const { signature, message } = JSON.parse(proof);
+            elizaLogger.info("signature", signature);
+            elizaLogger.info("message", message);
+            elizaLogger.info("signer", tokenOwner);
+            const isValid = createVerify('sha256').update(message).verify(tokenOwner, signature, 'base64');
+            return isValid;
         } catch (error) {
             elizaLogger.error(`Error when validating token ${tokenData.tokenId}:`, error);
             return false;
@@ -153,10 +152,6 @@ export class AgentNFTClient {
             fs.mkdirSync(this.baseDir, { recursive: true });
         }
 
-        if (dataHashes.length !== NUM_AGENT_HASHES || dataDescriptions.length !== NUM_AGENT_HASHES) {
-            throw new Error(`Expected ${NUM_AGENT_HASHES} data hashes and descriptions, got ${dataHashes.length} hashes and ${dataDescriptions.length} descriptions`);
-        }
-
         elizaLogger.info(`Downloading data for token ${tokenId}`);
 
         try {

diff --git a/agent/src/index.ts b/agent/src/index.ts
@@ -23,7 +23,7 @@ export const startAgents = async () => {
         } else if (args.token) {
             // load from nft
             elizaLogger.info("Starting in NFT mode...");
-            characters = await loadFromNFT(args.token, args.dir);
+            characters = await loadFromNFT(args.token, args.proof, args.dir);
         } else {
             // load from default character
             elizaLogger.info("Starting with default character...");

diff --git a/agent/src/types.ts b/agent/src/types.ts
@@ -2,6 +2,7 @@ export interface ParsedArguments {
     character?: string;
     characters?: string;
     token?: string;
+    proof?: string;
     dir?: string;
 }