diff --git a/bun.lockb b/bun.lockb index 5dffd68ce80..9eda4ab2b65 100755 Binary files a/bun.lockb and b/bun.lockb differ diff --git a/packages/agent/src/api.ts b/packages/agent/src/api.ts index 9a244c034ed..84872a1e627 100644 --- a/packages/agent/src/api.ts +++ b/packages/agent/src/api.ts @@ -2,7 +2,10 @@ import { type Character, getEnvVariable, IAgentRuntime, + ITeeLogService, logger, + ServiceType, + TeeLogQuery, type UUID, validateCharacterConfig, validateUuid @@ -352,5 +355,101 @@ export function createApiRouter( } }); + router.get("/tee/agents", async (req, res) => { + try { + const allAgents = []; + + for (const agentRuntime of agents.values()) { + const teeLogService = agentRuntime + .getService(ServiceType.TEE) + .getInstance(); + + const agents = await teeLogService.getAllAgents(); + allAgents.push(...agents); + } + + const runtime: IAgentRuntime = agents.values().next().value; + const teeLogService = runtime + .getService(ServiceType.TEE) + .getInstance() as ITeeLogService; + const attestation = await teeLogService.generateAttestation( + JSON.stringify(allAgents) + ); + res.json({ agents: allAgents, attestation: attestation }); + } catch (error) { + logger.error("Failed to get TEE agents:", error); + res.status(500).json({ + error: "Failed to get TEE agents", + }); + } + }); + + router.get("/tee/agents/:agentId", async (req, res) => { + try { + const agentId = req.params.agentId; + const agentRuntime = agents.get(agentId); + if (!agentRuntime) { + res.status(404).json({ error: "Agent not found" }); + return; + } + + const teeLogService = agentRuntime + .getService(ServiceType.TEE) + .getInstance(); + + const teeAgent = await teeLogService.getAgent(agentId); + const attestation = await teeLogService.generateAttestation( + JSON.stringify(teeAgent) + ); + res.json({ agent: teeAgent, attestation: attestation }); + } catch (error) { + logger.error("Failed to get TEE agent:", error); + res.status(500).json({ + error: "Failed to get TEE agent", + }); + } + }); + + router.post( + "/tee/logs", + async (req: express.Request, res: express.Response) => { + try { + const query = req.body.query || {}; + const page = Number.parseInt(req.body.page) || 1; + const pageSize = Number.parseInt(req.body.pageSize) || 10; + + const teeLogQuery: TeeLogQuery = { + agentId: query.agentId || "", + roomId: query.roomId || "", + userId: query.userId || "", + type: query.type || "", + containsContent: query.containsContent || "", + startTimestamp: query.startTimestamp || undefined, + endTimestamp: query.endTimestamp || undefined, + }; + const agentRuntime: IAgentRuntime = agents.values().next().value; + const teeLogService = agentRuntime + .getService(ServiceType.TEE); + const pageQuery = await teeLogService.getLogs( + teeLogQuery, + page, + pageSize + ); + const attestation = await teeLogService.generateAttestation( + JSON.stringify(pageQuery) + ); + res.json({ + logs: pageQuery, + attestation: attestation, + }); + } catch (error) { + logger.error("Failed to get TEE logs:", error); + res.status(500).json({ + error: "Failed to get TEE logs", + }); + } + } + ); + return router; } diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 243f8908f34..4fd2bb6d84e 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -141,6 +141,7 @@ export enum ServiceType { REMOTE_FILES = "aws_s3", WEB_SEARCH = "web_search", EMAIL = "email", + TEE = "tee", } /** @@ -1105,6 +1106,23 @@ export interface IFileService extends Service { generateSignedUrl(fileName: string, expiresIn: number): Promise; } +export interface ITeeLogService extends Service { + getInstance(): ITeeLogService; + log( + agentId: string, + roomId: string, + userId: string, + type: string, + content: string, + ): Promise; + + generateAttestation(reportData: string, hashAlgorithm?: T | any): Promise; + getAllAgents(): Promise; + getAgent(agentId: string): Promise; + getLogs(query: TeeLogQuery, page: number, pageSize: number): Promise>; + +} + export interface TestCase { name: string; fn: (runtime: IAgentRuntime) => Promise | void; @@ -1113,4 +1131,104 @@ export interface TestCase { export interface TestSuite { name: string; tests: TestCase[]; +} + +// Represents a log entry in the TeeLog table, containing details about agent activities. +export interface TeeLog { + id: string; + agentId: string; + roomId: string; + userId: string; + type: string; + content: string; + timestamp: number; + signature: string; +} + +export interface TeeLogQuery { + agentId?: string; + roomId?: string; + userId?: string; + type?: string; + containsContent?: string; + startTimestamp?: number; + endTimestamp?: number; +} + +// Represents an agent in the TeeAgent table, containing details about the agent. +export interface TeeAgent { + id: string; // Primary key + // Allow duplicate agentId. + // This is to support the case where the same agentId is registered multiple times. + // Each time the agent restarts, we will generate a new keypair and attestation. + agentId: string; + agentName: string; + createdAt: number; + publicKey: string; + attestation: string; +} + +export interface TeePageQuery { + page: number; + pageSize: number; + total?: number; + data?: Result; +} + +export abstract class TeeLogDAO { + db: DB; + + abstract initialize(): Promise; + + abstract addLog(log: TeeLog): Promise; + + abstract getPagedLogs( + query: TeeLogQuery, + page: number, + pageSize: number + ): Promise>; + + abstract addAgent(agent: TeeAgent): Promise; + + abstract getAgent(agentId: string): Promise; + + abstract getAllAgents(): Promise; +} + +export enum TEEMode { + OFF = "OFF", + LOCAL = "LOCAL", // For local development with simulator + DOCKER = "DOCKER", // For docker development with simulator + PRODUCTION = "PRODUCTION", // For production without simulator +} + +export interface RemoteAttestationQuote { + quote: string; + timestamp: number; +} + +export interface DeriveKeyAttestationData { + agentId: string; + publicKey: string; + subject?: string; +} + +export interface RemoteAttestationMessage { + agentId: string; + timestamp: number; + message: { + userId: string; + roomId: string; + content: string; + } +} + +export interface SgxAttestation { + quote: string; + timestamp: number; +} + +export enum TeeType { + SGX_GRAMINE = "sgx_gramine", + TDX_DSTACK = "tdx_dstack", } \ No newline at end of file diff --git a/packages/plugin-tee/.npmignore b/packages/plugin-tee/.npmignore new file mode 100644 index 00000000000..078562eceab --- /dev/null +++ b/packages/plugin-tee/.npmignore @@ -0,0 +1,6 @@ +* + +!dist/** +!package.json +!readme.md +!tsup.config.ts \ No newline at end of file diff --git a/packages/plugin-tee/README.md b/packages/plugin-tee/README.md new file mode 100644 index 00000000000..cd8207b7c62 --- /dev/null +++ b/packages/plugin-tee/README.md @@ -0,0 +1,136 @@ +# TEE Log Plugin for Eliza + +The TEE Log Plugin for Eliza is designed to enhance the logging capabilities of the Eliza by providing a structured way to generate, store and verify TEE (Trusted Execution Environment) logs for agents. This plugin ensures that all sensitive interactions are securely logged, providing a transparent and tamper-resistant record of all sensitive activities. + +## Background + +As Eliza is a fully autonomous AI agent capable of running within a TEE, we need to demonstrate to the outside world that we are indeed operating within a TEE. This allows external parties to verify that our actions are protected by the TEE and that they are autonomously executed by Eliza, without any third-party interference. Therefore, it is necessary to leverage TEE's remote attestation and establish a TEE logging mechanism to prove that these operations are entirely and autonomously performed by Eliza within the TEE. + +## Requirements + +Since the TEE Logging is based on the TEE, it is necessary to have a TEE enabled environment. Currently, we support Intel SGX (Gramine) and Intel TDX (dstack). +- using Intel SGX (Gramine), you need to enable the plugin-sgx in the Eliza runtime, which is enabled in SGX env automatically. +- using Intel TDX (dstack), you need to enable the plugin-tee in the Eliza runtime. + +## TEE Logging Mechanism + +## TEE Logging Mechanism + +1. **Key Pair Generation and Attestation**: + - During startup, each agent generates a key pair and creates a remote attestation for the public key. The private key is securely stored in the TEE's encrypted memory. The agent's relevant information, along with the public key and attestation, is recorded in a local database. A new key pair is generated each time the agent is updated or restarted to ensure key security. + +2. **Log Recording**: + - For each log entry, basic information is recorded, including `agentId`, `roomId`, `userId`, `type`, `content`, and `timestamp`. This information is concatenated and signed using the agent's corresponding private key to ensure verifiability. The verification process follows this trust chain: + - Verify the attestation. + - Trust the public key contained in the attestation. + - Use the public key to verify the signature. + - Trust the complete log record. + +3. **Data Storage**: + - All log data must be stored in the TEE's encrypted file system in production environments. Storing data in plaintext is prohibited to prevent tampering. + +4. **Log Extraction for Verification**: + - Third parties can extract TEE logs for verification purposes. Two types of information can be extracted: + - **Agent Information**: This includes the agent's metadata, public key, and attestation, which can be used to verify the agent's public key. + - **Log Information**: Required logs can be extracted, with the agent's attestation and public key used to verify the signature, ensuring that each record remains untampered. + +5. **Integrity Protection**: + - When users extract TEE logs via the REST API, the results are hashed, and an attestation is generated. After extraction, users can verify the attestation by comparing the hash value contained within it to the extracted results, thereby ensuring the integrity of the data. + +## Services + +- **[TeeLogService]**: This service is responsible for generating and storing TEE logs for agents. + +### Class: TeeLogService + +The `TeeLogService` class implements the `ITeeLogService` interface and extends the `Service` class. It manages the logging of sensitive interactions within a Trusted Execution Environment (TEE). + +#### Methods + +- **getInstance()**: `TeeLogService` + - Returns the singleton instance of the `TeeLogService`. + +- **static get serviceType()**: `ServiceType` + - Returns the service type for TEE logging. + +- **async initialize(runtime: IAgentRuntime): Promise** + - Initializes the TEE log service. It checks the runtime settings to configure the TEE type and enables logging if configured. + +- **async log(agentId: string, roomId: string, userId: string, type: string, content: string): Promise** + - Logs an interaction with the specified parameters. Returns `false` if TEE logging is not enabled. + +- **async getAllAgents(): Promise** + - Retrieves all agents that have been logged. Returns an empty array if TEE logging is not enabled. + +- **async getAgent(agentId: string): Promise** + - Retrieves the details of a specific agent by their ID. Returns `undefined` if TEE logging is not enabled. + +- **async getLogs(query: TeeLogQuery, page: number, pageSize: number): Promise>** + - Retrieves logs based on the provided query parameters. Returns an empty result if TEE logging is not enabled. + +- **async generateAttestation(userReport: string): Promise** + - Generates an attestation based on the provided user report. + +### Storage + +The TEE logs are stored in a SQLite database, which is located at `./data/tee_log.sqlite`. The database is automatically created when the service is initialized. + +Important: You need to use the encrypted file system to store the database file in production, otherwise the database will be compromised. Since TEE only protects memory-in-use, the disk is not protected by the TEE. However, Many TEE development tools support the encrypted file system, for example, you can refer to the [Gramine Encrypted files](https://gramine.readthedocs.io/en/latest/manifest-syntax.html#encrypted-files) documentation for more information. + +### Usage + +To use the `TeeLogService`, ensure that the TEE environment is properly configured and initialized. + +Enable the TEE logging in the Eliza .env file: + +```env +TEE_LOG_ENABLED=true +``` + +The logging isn't integrated for actions by default, you need to integrate the logging for the actions you want to log. For example, if you want to log the `Continue` action of plugin-bootstrap, you can do the following: + +First, add plugin-tee-log to the dependencies of plugin-bootstrap: + +```json +"@elizaos/plugin-tee-log": "workspace:*", +``` + +Then, add the following code to the `Continue` action: + +```typescript +import { + ServiceType, + ITeeLogService, +} from "@elizaos/core"; + + +// In the handler of the action + handler: async ( + runtime: IAgentRuntime, + message: Memory, + state: State, + options: any, + callback: HandlerCallback + ) => { + // Continue the action + + // Log the action + const teeLogService = runtime + .getService(ServiceType.TEE_LOG) + .getInstance(); + if (teeLogService.log( + runtime.agentId, + message.roomId, + message.userId, + "The type of the log, for example, Action:CONTINUE", + "The content that you want to log" + ) + ) { + console.log("Logged TEE log successfully"); + } + + // Continue the action + } +``` + +After configuring the logging for the action, you can run the Eliza and see the logs through the client-direct REST API. See more details in the [Client-Direct REST API](../client-direct/src/README.md) documentation. \ No newline at end of file diff --git a/packages/plugin-tee/biome.json b/packages/plugin-tee/biome.json new file mode 100644 index 00000000000..818716a6219 --- /dev/null +++ b/packages/plugin-tee/biome.json @@ -0,0 +1,41 @@ +{ + "$schema": "https://biomejs.dev/schemas/1.5.3/schema.json", + "organizeImports": { + "enabled": false + }, + "linter": { + "enabled": true, + "rules": { + "recommended": true, + "correctness": { + "noUnusedVariables": "error" + }, + "suspicious": { + "noExplicitAny": "error" + }, + "style": { + "useConst": "error", + "useImportType": "off" + } + } + }, + "formatter": { + "enabled": true, + "indentStyle": "space", + "indentWidth": 4, + "lineWidth": 100 + }, + "javascript": { + "formatter": { + "quoteStyle": "single", + "trailingCommas": "es5" + } + }, + "files": { + "ignore": [ + "dist/**/*", + "extra/**/*", + "node_modules/**/*" + ] + } +} \ No newline at end of file diff --git a/packages/plugin-tee/package.json b/packages/plugin-tee/package.json new file mode 100644 index 00000000000..756e41297fd --- /dev/null +++ b/packages/plugin-tee/package.json @@ -0,0 +1,29 @@ +{ + "name": "@elizaos/plugin-tee", + "version": "0.25.6-alpha.1", + "main": "dist/index.js", + "type": "module", + "types": "dist/index.d.ts", + "dependencies": { + "@elizaos/core": "workspace:*", + "@phala/dstack-sdk": "^0.1.7", + "better-sqlite3": "11.8.1", + "elliptic": "6.6.1" + }, + "devDependencies": { + "@biomejs/biome": "1.5.3", + "@types/node": "^20.0.0", + "tsup": "^8.3.5" + }, + "scripts": { + "build": "tsup --format esm --dts", + "dev": "tsup --format esm --dts --watch", + "lint": "biome check src/", + "lint:fix": "biome check --apply src/", + "format": "biome format src/", + "format:fix": "biome format --write src/" + }, + "publishConfig": { + "access": "public" + } +} diff --git a/packages/plugin-tee/src/__tests__/deriveKey.test.ts b/packages/plugin-tee/src/__tests__/deriveKey.test.ts new file mode 100644 index 00000000000..deabae7bcaf --- /dev/null +++ b/packages/plugin-tee/src/__tests__/deriveKey.test.ts @@ -0,0 +1,111 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { DeriveKeyProvider } from '../src/providers/deriveKeyProvider'; +import { TappdClient } from '@phala/dstack-sdk'; +import { TEEMode } from '../src/types/tee'; + +// Mock dependencies +vi.mock('@phala/dstack-sdk', () => ({ + TappdClient: vi.fn().mockImplementation(() => ({ + deriveKey: vi.fn().mockResolvedValue({ + asUint8Array: () => new Uint8Array([1, 2, 3, 4, 5]) + }), + tdxQuote: vi.fn().mockResolvedValue({ + quote: 'mock-quote-data', + replayRtmrs: () => ['rtmr0', 'rtmr1', 'rtmr2', 'rtmr3'] + }), + rawDeriveKey: vi.fn() + })) +})); + +vi.mock('@solana/web3.js', () => ({ + Keypair: { + fromSeed: vi.fn().mockReturnValue({ + publicKey: { + toBase58: () => 'mock-solana-public-key' + } + }) + } +})); + +vi.mock('viem/accounts', () => ({ + privateKeyToAccount: vi.fn().mockReturnValue({ + address: 'mock-evm-address' + }) +})); + +describe('DeriveKeyProvider', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('constructor', () => { + it('should initialize with LOCAL mode', () => { + const _provider = new DeriveKeyProvider(TEEMode.LOCAL); + expect(TappdClient).toHaveBeenCalledWith('http://localhost:8090'); + }); + + it('should initialize with DOCKER mode', () => { + const _provider = new DeriveKeyProvider(TEEMode.DOCKER); + expect(TappdClient).toHaveBeenCalledWith('http://host.docker.internal:8090'); + }); + + it('should initialize with PRODUCTION mode', () => { + const _provider = new DeriveKeyProvider(TEEMode.PRODUCTION); + expect(TappdClient).toHaveBeenCalledWith(); + }); + + it('should throw error for invalid mode', () => { + expect(() => new DeriveKeyProvider('INVALID_MODE')).toThrow('Invalid TEE_MODE'); + }); + }); + + describe('rawDeriveKey', () => { + let _provider: DeriveKeyProvider; + + beforeEach(() => { + _provider = new DeriveKeyProvider(TEEMode.LOCAL); + }); + + it('should derive raw key successfully', async () => { + const path = 'test-path'; + const subject = 'test-subject'; + const result = await _provider.rawDeriveKey(path, subject); + + const client = TappdClient.mock.results[0].value; + expect(client.deriveKey).toHaveBeenCalledWith(path, subject); + expect(result.asUint8Array()).toEqual(new Uint8Array([1, 2, 3, 4, 5])); + }); + + it('should handle errors during raw key derivation', async () => { + const mockError = new Error('Key derivation failed'); + vi.mocked(TappdClient).mockImplementationOnce(() => { + const instance = new TappdClient(); + instance.deriveKey = vi.fn().mockRejectedValueOnce(mockError); + instance.tdxQuote = vi.fn(); + instance.rawDeriveKey = vi.fn(); + return instance; + }); + + const provider = new DeriveKeyProvider(TEEMode.LOCAL); + await expect(provider.rawDeriveKey('path', 'subject')).rejects.toThrow(mockError); + }); + }); + + describe('deriveEd25519Keypair', () => { + let provider: DeriveKeyProvider; + + beforeEach(() => { + provider = new DeriveKeyProvider(TEEMode.LOCAL); + }); + + it('should derive Ed25519 keypair successfully', async () => { + const path = 'test-path'; + const subject = 'test-subject'; + const result = await provider.deriveEd25519Keypair(path, subject); + + const client = TappdClient.mock.results[0].value; + expect(client.deriveKey).toHaveBeenCalledWith(path, subject); + expect(result.keypair.publicKey.toBase58()).toEqual('mock-solana-public-key'); + }); + }); +}); \ No newline at end of file diff --git a/packages/plugin-tee/src/__tests__/remoteAttestation.test.ts b/packages/plugin-tee/src/__tests__/remoteAttestation.test.ts new file mode 100644 index 00000000000..4890d6683e5 --- /dev/null +++ b/packages/plugin-tee/src/__tests__/remoteAttestation.test.ts @@ -0,0 +1,81 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { RemoteAttestationProvider } from '../src/providers/remoteAttestationProvider'; +import { TappdClient } from '@phala/dstack-sdk'; +import { TEEMode } from '../src/types/tee'; + +// Mock TappdClient +vi.mock('@phala/dstack-sdk', () => ({ + TappdClient: vi.fn().mockImplementation(() => ({ + tdxQuote: vi.fn().mockResolvedValue({ + quote: 'mock-quote-data', + replayRtmrs: () => ['rtmr0', 'rtmr1', 'rtmr2', 'rtmr3'] + }), + deriveKey: vi.fn() + })) +})); + +describe('RemoteAttestationProvider', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('constructor', () => { + it('should initialize with LOCAL mode', () => { + const _provider = new RemoteAttestationProvider(TEEMode.LOCAL); + expect(TappdClient).toHaveBeenCalledWith('http://localhost:8090'); + }); + + it('should initialize with DOCKER mode', () => { + const _provider = new RemoteAttestationProvider(TEEMode.DOCKER); + expect(TappdClient).toHaveBeenCalledWith('http://host.docker.internal:8090'); + }); + + it('should initialize with PRODUCTION mode', () => { + const _provider = new RemoteAttestationProvider(TEEMode.PRODUCTION); + expect(TappdClient).toHaveBeenCalledWith(); + }); + + it('should throw error for invalid mode', () => { + expect(() => new RemoteAttestationProvider('INVALID_MODE')).toThrow('Invalid TEE_MODE'); + }); + }); + + describe('generateAttestation', () => { + let provider: RemoteAttestationProvider; + + beforeEach(() => { + provider = new RemoteAttestationProvider(TEEMode.LOCAL); + }); + + it('should generate attestation successfully', async () => { + const reportData = 'test-report-data'; + const quote = await provider.generateAttestation(reportData); + + expect(quote).toEqual({ + quote: 'mock-quote-data', + timestamp: expect.any(Number) + }); + }); + + it('should handle errors during attestation generation', async () => { + const mockError = new Error('TDX Quote generation failed'); + const mockTdxQuote = vi.fn().mockRejectedValue(mockError); + vi.mocked(TappdClient).mockImplementationOnce(() => ({ + tdxQuote: mockTdxQuote, + deriveKey: vi.fn() + })); + + const provider = new RemoteAttestationProvider(TEEMode.LOCAL); + await expect(provider.generateAttestation('test-data')).rejects.toThrow('Failed to generate TDX Quote'); + }); + + it('should pass hash algorithm to tdxQuote when provided', async () => { + const reportData = 'test-report-data'; + const hashAlgorithm = 'raw'; + await provider.generateAttestation(reportData, hashAlgorithm); + + const client = TappdClient.mock.results[0].value; + expect(client.tdxQuote).toHaveBeenCalledWith(reportData, hashAlgorithm); + }); + }); +}); \ No newline at end of file diff --git a/packages/plugin-tee/src/__tests__/remoteAttestationAction.test.ts b/packages/plugin-tee/src/__tests__/remoteAttestationAction.test.ts new file mode 100644 index 00000000000..897de2c7d86 --- /dev/null +++ b/packages/plugin-tee/src/__tests__/remoteAttestationAction.test.ts @@ -0,0 +1,103 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { remoteAttestationAction } from '../src/actions/remoteAttestation'; +import { RemoteAttestationProvider } from '../src/providers/remoteAttestationProvider'; + +// Mock dependencies +vi.mock('../src/providers/remoteAttestationProvider'); +vi.mock('undici', () => ({ + fetch: vi.fn().mockResolvedValue({ + json: () => Promise.resolve({ checksum: 'mock-checksum' }) + }) +})); + +describe('remoteAttestationAction', () => { + const mockRuntime = { + agentId: 'test-agent-id', + getSetting: vi.fn().mockReturnValue('LOCAL'), + getState: vi.fn(), + setState: vi.fn(), + message:{ userId: 'user', roomId: 'room', content: { text: 'If you are running in a TEE, generate a remote attestation' } }, + setConversation: vi.fn() + }; + + const mockCallback = vi.fn(); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(RemoteAttestationProvider).mockImplementation(() => ({ + generateAttestation: vi.fn().mockResolvedValue({ + quote: '0x1234', + timestamp: Date.now() + }) + })); + }); + + it('should have correct name and similes', () => { + expect(remoteAttestationAction.name).toBe('REMOTE_ATTESTATION'); + expect(remoteAttestationAction.similes).toContain('REMOTE_ATTESTATION'); + expect(remoteAttestationAction.similes).toContain('TEE_REMOTE_ATTESTATION'); + expect(remoteAttestationAction.similes).toContain('TEE_ATTESTATION'); + }); + + describe('handler', () => { + it('should generate and upload attestation successfully', async () => { + const result = await remoteAttestationAction.handler( + mockRuntime, + mockRuntime.message, + {}, + {}, + mockCallback + ); + + expect(result).toBe(true); + expect(mockCallback).toHaveBeenCalledWith({ + text: expect.stringContaining('https://proof.t16z.com/reports/mock-checksum'), + action: 'NONE' + }); + }); + + it('should handle errors during attestation generation', async () => { + const mockError = new Error('Attestation generation failed'); + vi.mocked(RemoteAttestationProvider).mockImplementation(() => ({ + generateAttestation: vi.fn().mockRejectedValueOnce(mockError), + client: { + tdxQuote: vi.fn(), + deriveKey: vi.fn() + } + })); + + const result = await remoteAttestationAction.handler( + mockRuntime, + {}, + {}, + {}, + mockCallback + ); + + expect(result).toBe(false); + }); + }); + + describe('validate', () => { + it('should always return true', async () => { + const result = await remoteAttestationAction.validate(mockRuntime); + expect(result).toBe(true); + }); + }); + + describe('examples', () => { + it('should have valid example conversations', () => { + expect(remoteAttestationAction.examples).toBeInstanceOf(Array); + expect(remoteAttestationAction.examples[0]).toBeInstanceOf(Array); + + const [userMessage, agentMessage] = remoteAttestationAction.examples[0]; + expect(userMessage.user).toBe('{{user1}}'); + expect(userMessage.content.text).toBe('If you are running in a TEE, generate a remote attestation'); + expect(userMessage.content.action).toBe('REMOTE_ATTESTATION'); + + expect(agentMessage.user).toBe('{{user2}}'); + expect(agentMessage.content.text).toBe('Of course, one second...'); + expect(agentMessage.content.action).toBeUndefined(); + }); + }); +}); \ No newline at end of file diff --git a/packages/plugin-tee/src/__tests__/timeout.test.ts b/packages/plugin-tee/src/__tests__/timeout.test.ts new file mode 100644 index 00000000000..f16a54d8a49 --- /dev/null +++ b/packages/plugin-tee/src/__tests__/timeout.test.ts @@ -0,0 +1,117 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { RemoteAttestationProvider } from '../src/providers/remoteAttestationProvider'; +import { DeriveKeyProvider } from '../src/providers/deriveKeyProvider'; +import { TEEMode } from '../src/types/tee'; +import { TappdClient } from '@phala/dstack-sdk'; + +// Mock TappdClient +vi.mock('@phala/dstack-sdk', () => ({ + TappdClient: vi.fn().mockImplementation(() => ({ + tdxQuote: vi.fn(), + deriveKey: vi.fn() + })) +})); + +describe('TEE Provider Timeout Tests', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('RemoteAttestationProvider', () => { + it('should handle API timeout during attestation generation', async () => { + const mockTdxQuote = vi.fn() + .mockRejectedValueOnce(new Error('Request timed out')); + + vi.mocked(TappdClient).mockImplementation(() => ({ + tdxQuote: mockTdxQuote, + deriveKey: vi.fn() + })); + + const provider = new RemoteAttestationProvider(TEEMode.LOCAL); + await expect(() => provider.generateAttestation('test-data')) + .rejects + .toThrow('Failed to generate TDX Quote: Request timed out'); + + // Verify the call was made once + expect(mockTdxQuote).toHaveBeenCalledTimes(1); + expect(mockTdxQuote).toHaveBeenCalledWith('test-data', undefined); + }); + + it('should handle network errors during attestation generation', async () => { + const mockTdxQuote = vi.fn() + .mockRejectedValueOnce(new Error('Network error')); + + vi.mocked(TappdClient).mockImplementation(() => ({ + tdxQuote: mockTdxQuote, + deriveKey: vi.fn() + })); + + const provider = new RemoteAttestationProvider(TEEMode.LOCAL); + await expect(() => provider.generateAttestation('test-data')) + .rejects + .toThrow('Failed to generate TDX Quote: Network error'); + + expect(mockTdxQuote).toHaveBeenCalledTimes(1); + }); + + it('should handle successful attestation generation', async () => { + const mockQuote = { + quote: 'test-quote', + replayRtmrs: () => ['rtmr0', 'rtmr1', 'rtmr2', 'rtmr3'] + }; + + const mockTdxQuote = vi.fn().mockResolvedValueOnce(mockQuote); + + vi.mocked(TappdClient).mockImplementation(() => ({ + tdxQuote: mockTdxQuote, + deriveKey: vi.fn() + })); + + const provider = new RemoteAttestationProvider(TEEMode.LOCAL); + const result = await provider.generateAttestation('test-data'); + + expect(mockTdxQuote).toHaveBeenCalledTimes(1); + expect(result).toEqual({ + quote: 'test-quote', + timestamp: expect.any(Number) + }); + }); + }); + + describe('DeriveKeyProvider', () => { + it('should handle API timeout during key derivation', async () => { + const mockDeriveKey = vi.fn() + .mockRejectedValueOnce(new Error('Request timed out')); + + vi.mocked(TappdClient).mockImplementation(() => ({ + tdxQuote: vi.fn(), + deriveKey: mockDeriveKey + })); + + const provider = new DeriveKeyProvider(TEEMode.LOCAL); + await expect(() => provider.rawDeriveKey('test-path', 'test-subject')) + .rejects + .toThrow('Request timed out'); + + expect(mockDeriveKey).toHaveBeenCalledTimes(1); + expect(mockDeriveKey).toHaveBeenCalledWith('test-path', 'test-subject'); + }); + + it('should handle API timeout during Ed25519 key derivation', async () => { + const mockDeriveKey = vi.fn() + .mockRejectedValueOnce(new Error('Request timed out')); + + vi.mocked(TappdClient).mockImplementation(() => ({ + tdxQuote: vi.fn(), + deriveKey: mockDeriveKey + })); + + const provider = new DeriveKeyProvider(TEEMode.LOCAL); + await expect(() => provider.deriveEd25519Keypair('test-path', 'test-subject')) + .rejects + .toThrow('Request timed out'); + + expect(mockDeriveKey).toHaveBeenCalledTimes(1); + }); + }); +}); diff --git a/packages/plugin-tee/src/adapters/sqliteDAO.ts b/packages/plugin-tee/src/adapters/sqliteDAO.ts new file mode 100644 index 00000000000..9d985c23c8b --- /dev/null +++ b/packages/plugin-tee/src/adapters/sqliteDAO.ts @@ -0,0 +1,125 @@ +import type { Database } from "better-sqlite3"; +import { TeeLogDAO, type TeeAgent, type TeeLog, type TeeLogQuery, type TeePageQuery } from "../types.ts"; +import { sqliteTables } from "./sqliteTables.ts"; + +export class SqliteTeeLogDAO extends TeeLogDAO { + constructor(db: Database) { + super(); + this.db = db; + } + + async initialize(): Promise { + this.db.exec(sqliteTables); + } + + async addLog(log: TeeLog): Promise { + const stmt = this.db.prepare( + "INSERT INTO tee_logs (id, agentId, roomId, userId, type, content, timestamp, signature) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" + ); + try { + stmt.run(log.id, log.agentId, log.roomId, log.userId, log.type, log.content, log.timestamp, log.signature); + return true; + } catch (error) { + console.error("Error adding log to database", error); + return false; + } + } + + async getPagedLogs(query: TeeLogQuery, page: number, pageSize: number): Promise> { + if (page < 1) { + page = 1; + } + const offset = (page - 1) * pageSize; + const limit = pageSize; + + const whereConditions = []; + const params = []; + + if (query.agentId && query.agentId !== "") { + whereConditions.push("agentId = ?"); + params.push(query.agentId); + } + if (query.roomId && query.roomId !== "") { + whereConditions.push("roomId = ?"); + params.push(query.roomId); + } + if (query.userId && query.userId !== "") { + whereConditions.push("userId = ?"); + params.push(query.userId); + } + if (query.type && query.type !== "") { + whereConditions.push("type = ?"); + params.push(query.type); + } + if (query.containsContent && query.containsContent !== "") { + whereConditions.push("content LIKE ?"); + params.push(`%${query.containsContent}%`); + } + if (query.startTimestamp) { + whereConditions.push("timestamp >= ?"); + params.push(query.startTimestamp); + } + if (query.endTimestamp) { + whereConditions.push("timestamp <= ?"); + params.push(query.endTimestamp); + } + + const whereClause = + whereConditions.length > 0 ? `WHERE ${whereConditions.join(" AND ")}` : ""; + + try { + const total_stmt = this.db.prepare( + `SELECT COUNT(*) as total FROM tee_logs ${whereClause}` + ); + const total = total_stmt.get(params).total; + + const logs_stmt = this.db.prepare( + `SELECT * FROM tee_logs ${whereClause} ORDER BY timestamp ASC LIMIT ? OFFSET ?` + ); + const logs = logs_stmt.all(...params, limit, offset); + + return { + page, + pageSize, + total, + data: logs, + }; + } catch (error) { + console.error("Error getting paged logs from database", error); + throw error; + } + } + + async addAgent(agent: TeeAgent): Promise { + const stmt = this.db.prepare( + "INSERT INTO tee_agents (id, agentId, agentName, createdAt, publicKey, attestation) VALUES (?, ?, ?, ?, ?, ?)" + ); + try { + stmt.run(agent.id, agent.agentId, agent.agentName, agent.createdAt, agent.publicKey, agent.attestation); + return true; + } catch (error) { + console.error("Error adding agent to database", error); + return false; + } + } + + async getAgent(agentId: string): Promise { + const stmt = this.db.prepare("SELECT * FROM tee_agents WHERE agentId = ? ORDER BY createdAt DESC LIMIT 1"); + try { + return stmt.get(agentId); + } catch (error) { + console.error("Error getting agent from database", error); + throw error; + } + } + + async getAllAgents(): Promise { + const stmt = this.db.prepare("SELECT * FROM tee_agents"); + try { + return stmt.all(); + } catch (error) { + console.error("Error getting all agents from database", error); + throw error; + } + } +} diff --git a/packages/plugin-tee/src/adapters/sqliteTables.ts b/packages/plugin-tee/src/adapters/sqliteTables.ts new file mode 100644 index 00000000000..9d24b7820ad --- /dev/null +++ b/packages/plugin-tee/src/adapters/sqliteTables.ts @@ -0,0 +1,26 @@ +export const sqliteTables = ` +BEGIN TRANSACTION; + +-- Table: tee_logs +CREATE TABLE IF NOT EXISTS "tee_logs" ( + "id" TEXT PRIMARY KEY, + "agentId" TEXT NOT NULL, + "roomId" TEXT NOT NULL, + "userId" TEXT NOT NULL, + "type" TEXT NOT NULL, + "content" TEXT NOT NULL, + "timestamp" TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + "signature" TEXT NOT NULL +); + +-- Table: tee_agents +CREATE TABLE IF NOT EXISTS "tee_agents" ( + "id" TEXT PRIMARY KEY, + "agentId" TEXT NOT NULL, + "agentName" TEXT NOT NULL, + "createdAt" TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + "publicKey" TEXT NOT NULL, + "attestation" TEXT NOT NULL +); + +COMMIT;`; diff --git a/packages/plugin-tee/src/index.ts b/packages/plugin-tee/src/index.ts new file mode 100644 index 00000000000..7d2ba0a4af8 --- /dev/null +++ b/packages/plugin-tee/src/index.ts @@ -0,0 +1,34 @@ +import type { Plugin } from "@elizaos/core"; +import { + deriveKeyProvider, +} from "./providers/deriveKeyProvider"; +import { + remoteAttestationProvider +} from "./providers/remoteAttestationProvider"; + +export { + DeriveKeyProvider +} from "./providers/deriveKeyProvider"; + +import { sgxAttestationProvider } from "./providers/sgxAttestationProvider"; +import { TeeLogService } from "./services/teeLogService"; + +export { TeeLogService }; + +export const teePlugin: Plugin = { + name: "tee", + description: + "TEE plugin with actions to generate remote attestations and derive keys", + actions: [ + ], + evaluators: [ + ], + providers: [ + remoteAttestationProvider, + deriveKeyProvider, + sgxAttestationProvider, + ], + services: [ + new TeeLogService() + ], +}; \ No newline at end of file diff --git a/packages/plugin-tee/src/providers/deriveKeyProvider.ts b/packages/plugin-tee/src/providers/deriveKeyProvider.ts new file mode 100644 index 00000000000..2e41b14010a --- /dev/null +++ b/packages/plugin-tee/src/providers/deriveKeyProvider.ts @@ -0,0 +1,228 @@ +import { + type IAgentRuntime, + type Memory, + type Provider, + type State, + elizaLogger, +} from "@elizaos/core"; +import { Keypair } from "@solana/web3.js"; +import crypto from "crypto"; +import { type DeriveKeyResponse, TappdClient } from "@phala/dstack-sdk"; +import { privateKeyToAccount } from "viem/accounts"; +import { type PrivateKeyAccount, keccak256 } from "viem"; +import { RemoteAttestationProvider } from "./remoteAttestationProvider"; +import { TEEMode, type RemoteAttestationQuote, type DeriveKeyAttestationData } from "@elizaos/core"; + +class DeriveKeyProvider { + private client: TappdClient; + private raProvider: RemoteAttestationProvider; + + constructor(teeMode?: string) { + let endpoint: string | undefined; + + // Both LOCAL and DOCKER modes use the simulator, just with different endpoints + switch (teeMode) { + case TEEMode.LOCAL: + endpoint = "http://localhost:8090"; + elizaLogger.log( + "TEE: Connecting to local simulator at localhost:8090" + ); + break; + case TEEMode.DOCKER: + endpoint = "http://host.docker.internal:8090"; + elizaLogger.log( + "TEE: Connecting to simulator via Docker at host.docker.internal:8090" + ); + break; + case TEEMode.PRODUCTION: + endpoint = undefined; + elizaLogger.log( + "TEE: Running in production mode without simulator" + ); + break; + default: + throw new Error( + `Invalid TEE_MODE: ${teeMode}. Must be one of: LOCAL, DOCKER, PRODUCTION` + ); + } + + this.client = endpoint ? new TappdClient(endpoint) : new TappdClient(); + this.raProvider = new RemoteAttestationProvider(teeMode); + } + + private async generateDeriveKeyAttestation( + agentId: string, + publicKey: string, + subject?: string + ): Promise { + const deriveKeyData: DeriveKeyAttestationData = { + agentId, + publicKey, + subject, + }; + const reportdata = JSON.stringify(deriveKeyData); + elizaLogger.log( + "Generating Remote Attestation Quote for Derive Key..." + ); + const quote = await this.raProvider.generateAttestation(reportdata); + elizaLogger.log("Remote Attestation Quote generated successfully!"); + return quote; + } + + /** + * Derives a raw key from the given path and subject. + * @param path - The path to derive the key from. This is used to derive the key from the root of trust. + * @param subject - The subject to derive the key from. This is used for the certificate chain. + * @returns The derived key. + */ + async rawDeriveKey( + path: string, + subject: string + ): Promise { + try { + if (!path || !subject) { + elizaLogger.error( + "Path and Subject are required for key derivation" + ); + } + + elizaLogger.log("Deriving Raw Key in TEE..."); + const derivedKey = await this.client.deriveKey(path, subject); + + elizaLogger.log("Raw Key Derived Successfully!"); + return derivedKey; + } catch (error) { + elizaLogger.error("Error deriving raw key:", error); + throw error; + } + } + + /** + * Derives an Ed25519 keypair from the given path and subject. + * @param path - The path to derive the key from. This is used to derive the key from the root of trust. + * @param subject - The subject to derive the key from. This is used for the certificate chain. + * @param agentId - The agent ID to generate an attestation for. + * @returns An object containing the derived keypair and attestation. + */ + async deriveEd25519Keypair( + path: string, + subject: string, + agentId: string + ): Promise<{ keypair: Keypair; attestation: RemoteAttestationQuote }> { + try { + if (!path || !subject) { + elizaLogger.error( + "Path and Subject are required for key derivation" + ); + } + + elizaLogger.log("Deriving Key in TEE..."); + const derivedKey = await this.client.deriveKey(path, subject); + const uint8ArrayDerivedKey = derivedKey.asUint8Array(); + + const hash = crypto.createHash("sha256"); + hash.update(uint8ArrayDerivedKey); + const seed = hash.digest(); + const seedArray = new Uint8Array(seed); + const keypair = Keypair.fromSeed(seedArray.slice(0, 32)); + + // Generate an attestation for the derived key data for public to verify + const attestation = await this.generateDeriveKeyAttestation( + agentId, + keypair.publicKey.toBase58() + ); + elizaLogger.log("Key Derived Successfully!"); + + return { keypair, attestation }; + } catch (error) { + elizaLogger.error("Error deriving key:", error); + throw error; + } + } + + /** + * Derives an ECDSA keypair from the given path and subject. + * @param path - The path to derive the key from. This is used to derive the key from the root of trust. + * @param subject - The subject to derive the key from. This is used for the certificate chain. + * @param agentId - The agent ID to generate an attestation for. This is used for the certificate chain. + * @returns An object containing the derived keypair and attestation. + */ + async deriveEcdsaKeypair( + path: string, + subject: string, + agentId: string + ): Promise<{ + keypair: PrivateKeyAccount; + attestation: RemoteAttestationQuote; + }> { + try { + if (!path || !subject) { + elizaLogger.error( + "Path and Subject are required for key derivation" + ); + } + + elizaLogger.log("Deriving ECDSA Key in TEE..."); + const deriveKeyResponse: DeriveKeyResponse = + await this.client.deriveKey(path, subject); + const hex = keccak256(deriveKeyResponse.asUint8Array()); + const keypair: PrivateKeyAccount = privateKeyToAccount(hex); + + // Generate an attestation for the derived key data for public to verify + const attestation = await this.generateDeriveKeyAttestation( + agentId, + keypair.address + ); + elizaLogger.log("ECDSA Key Derived Successfully!"); + + return { keypair, attestation }; + } catch (error) { + elizaLogger.error("Error deriving ecdsa key:", error); + throw error; + } + } +} + +const deriveKeyProvider: Provider = { + get: async (runtime: IAgentRuntime, _message?: Memory, _state?: State) => { + const teeMode = runtime.getSetting("TEE_MODE"); + const provider = new DeriveKeyProvider(teeMode); + const agentId = runtime.agentId; + try { + // Validate wallet configuration + if (!runtime.getSetting("WALLET_SECRET_SALT")) { + elizaLogger.error( + "Wallet secret salt is not configured in settings" + ); + return ""; + } + + try { + const secretSalt = + runtime.getSetting("WALLET_SECRET_SALT") || "secret_salt"; + const solanaKeypair = await provider.deriveEd25519Keypair( + secretSalt, + "solana", + agentId + ); + const evmKeypair = await provider.deriveEcdsaKeypair( + secretSalt, + "evm", + agentId + ); + return JSON.stringify({ + solana: solanaKeypair.keypair.publicKey, + evm: evmKeypair.keypair.address, + }); + } catch (error) { + elizaLogger.error("Error creating PublicKey:", error); + return ""; + } + } catch (error) { + elizaLogger.error("Error in derive key provider:", error.message); + return `Failed to fetch derive key information: ${error instanceof Error ? error.message : "Unknown error"}`; + } + }, +}; + +export { deriveKeyProvider, DeriveKeyProvider }; diff --git a/packages/plugin-tee/src/providers/remoteAttestationProvider.ts b/packages/plugin-tee/src/providers/remoteAttestationProvider.ts new file mode 100644 index 00000000000..59bfd2bb2b9 --- /dev/null +++ b/packages/plugin-tee/src/providers/remoteAttestationProvider.ts @@ -0,0 +1,106 @@ +import { + type IAgentRuntime, + type Memory, + type Provider, + type State, + elizaLogger, +} from "@elizaos/core"; +import { type TdxQuoteResponse, TappdClient, type TdxQuoteHashAlgorithms } from "@phala/dstack-sdk"; +import { type RemoteAttestationQuote, TEEMode, type RemoteAttestationMessage } from "@elizaos/core"; + +class RemoteAttestationProvider { + private client: TappdClient; + + constructor(teeMode?: string) { + let endpoint: string | undefined; + + // Both LOCAL and DOCKER modes use the simulator, just with different endpoints + switch (teeMode) { + case TEEMode.LOCAL: + endpoint = "http://localhost:8090"; + elizaLogger.log( + "TEE: Connecting to local simulator at localhost:8090" + ); + break; + case TEEMode.DOCKER: + endpoint = "http://host.docker.internal:8090"; + elizaLogger.log( + "TEE: Connecting to simulator via Docker at host.docker.internal:8090" + ); + break; + case TEEMode.PRODUCTION: + endpoint = undefined; + elizaLogger.log( + "TEE: Running in production mode without simulator" + ); + break; + default: + throw new Error( + `Invalid TEE_MODE: ${teeMode}. Must be one of: LOCAL, DOCKER, PRODUCTION` + ); + } + + this.client = endpoint ? new TappdClient(endpoint) : new TappdClient(); + } + + async generateAttestation( + reportData: string, + hashAlgorithm?: TdxQuoteHashAlgorithms + ): Promise { + try { + elizaLogger.log("Generating attestation for: ", reportData); + const tdxQuote: TdxQuoteResponse = + await this.client.tdxQuote(reportData, hashAlgorithm); + const rtmrs = tdxQuote.replayRtmrs(); + elizaLogger.log( + `rtmr0: ${rtmrs[0]}\nrtmr1: ${rtmrs[1]}\nrtmr2: ${rtmrs[2]}\nrtmr3: ${rtmrs[3]}f` + ); + const quote: RemoteAttestationQuote = { + quote: tdxQuote.quote, + timestamp: Date.now(), + }; + elizaLogger.log("Remote attestation quote: ", quote); + return quote; + } catch (error) { + console.error("Error generating remote attestation:", error); + throw new Error( + `Failed to generate TDX Quote: ${ + error instanceof Error ? error.message : "Unknown error" + }` + ); + } + } +} + +// Keep the original provider for backwards compatibility +const remoteAttestationProvider: Provider = { + get: async (runtime: IAgentRuntime, message: Memory, _state?: State) => { + const teeMode = runtime.getSetting("TEE_MODE"); + const provider = new RemoteAttestationProvider(teeMode); + const agentId = runtime.agentId; + + try { + const attestationMessage: RemoteAttestationMessage = { + agentId: agentId, + timestamp: Date.now(), + message: { + userId: message.userId, + roomId: message.roomId, + content: message.content.text, + } + }; + elizaLogger.log("Generating attestation for: ", JSON.stringify(attestationMessage)); + const attestation = await provider.generateAttestation(JSON.stringify(attestationMessage)); + return `Your Agent's remote attestation is: ${JSON.stringify(attestation)}`; + } catch (error) { + console.error("Error in remote attestation provider:", error); + throw new Error( + `Failed to generate TDX Quote: ${ + error instanceof Error ? error.message : "Unknown error" + }` + ); + } + }, +}; + +export { remoteAttestationProvider, RemoteAttestationProvider }; diff --git a/packages/plugin-tee/src/providers/sgxAttestationProvider.ts b/packages/plugin-tee/src/providers/sgxAttestationProvider.ts new file mode 100644 index 00000000000..a50462f7d8b --- /dev/null +++ b/packages/plugin-tee/src/providers/sgxAttestationProvider.ts @@ -0,0 +1,105 @@ +import type { IAgentRuntime, Memory, Provider, State } from "@elizaos/core"; +import type { SgxAttestation } from "@elizaos/core"; +import { promises as fs } from 'node:fs'; // Fix: Use node: protocol +import { createHash } from 'node:crypto'; // Fix: Use node: protocol + +// Function to calculate SHA-256 and return a Buffer (32 bytes) +function calculateSHA256(input: string): Buffer { + const hash = createHash('sha256'); + hash.update(input); + return hash.digest(); +} + +class SgxAttestationProvider { + private readonly SGX_QUOTE_MAX_SIZE: number = 8192 * 4; + private readonly SGX_TARGET_INFO_SIZE: number = 512; + + private readonly MY_TARGET_INFO_PATH: string = "/dev/attestation/my_target_info"; + private readonly TARGET_INFO_PATH: string = "/dev/attestation/target_info"; + private readonly USER_REPORT_DATA_PATH: string = "/dev/attestation/user_report_data"; + private readonly QUOTE_PATH: string = "/dev/attestation/quote"; + + // Remove unnecessary constructor + // constructor() {} + + async generateAttestation( + reportData: string + ): Promise { + // Hash the report data to generate the raw user report. + // The resulting hash value is 32 bytes long. + // Ensure that the length of the raw user report does not exceed 64 bytes. + const rawUserReport = calculateSHA256(reportData); + + try { + // Check if the gramine attestation device file exists + await fs.access(this.MY_TARGET_INFO_PATH); + + const quote = await this.generateQuoteByGramine(rawUserReport); + const attestation: SgxAttestation = { + quote: quote, + timestamp: Date.now(), + }; + // console.log("SGX remote attestation: ", attestation); + return attestation; + } catch (error) { + console.error("Error generating SGX remote attestation:", error); + throw new Error( + `Failed to generate SGX Quote: ${ + error instanceof Error ? error.message : "Unknown error" + }` + ); + } + } + + async generateQuoteByGramine( + rawUserReport: Buffer + ): Promise { + if (rawUserReport.length > 64) { + throw new Error("the length of rawUserReport exceeds 64 bytes"); + } + + const myTargetInfo = await fs.readFile(this.MY_TARGET_INFO_PATH); + if (myTargetInfo.length !== this.SGX_TARGET_INFO_SIZE) { + throw new Error("Invalid my_target_info length"); + } + + await fs.writeFile(this.TARGET_INFO_PATH, myTargetInfo); + await fs.writeFile(this.USER_REPORT_DATA_PATH, rawUserReport); + + // Read quote + const quoteData = await fs.readFile(this.QUOTE_PATH); + if (quoteData.length > this.SGX_QUOTE_MAX_SIZE) { + throw new Error("Invalid quote length"); + } + + const realLen = quoteData.lastIndexOf(0); + if (realLen === -1) { + throw new Error("quote without EOF"); + } + + //return '0x' + quoteData.subarray(0, realLen + 1).toString('hex'); + return `0x${quoteData.subarray(0, realLen + 1).toString('hex')}`; // Fix: Use template literal + } +} + +const sgxAttestationProvider: Provider = { + get: async (runtime: IAgentRuntime, _message: Memory, _state?: State) => { + const provider = new SgxAttestationProvider(); + const agentId = runtime.agentId; + + try { + // console.log("Generating attestation for agent: ", agentId); + const attestation = await provider.generateAttestation(agentId); + return `Your Agent's remote attestation is: ${JSON.stringify(attestation)}`; + } catch (error) { + console.error("Error in remote attestation provider:", error); + throw new Error( + `Failed to generate SGX Quote: ${ + error instanceof Error ? error.message : "Unknown error" + }` + ); + } + }, +}; + +export { sgxAttestationProvider, SgxAttestationProvider }; diff --git a/packages/plugin-tee/src/providers/walletProvider.ts b/packages/plugin-tee/src/providers/walletProvider.ts new file mode 100644 index 00000000000..6f0f6dddb02 --- /dev/null +++ b/packages/plugin-tee/src/providers/walletProvider.ts @@ -0,0 +1,326 @@ +/* This is an example of how WalletProvider can use DeriveKeyProvider to generate a Solana Keypair */ +import { + type IAgentRuntime, + type Memory, + type Provider, + type State, + elizaLogger, +} from "@elizaos/core"; +import { Connection, type Keypair, type PublicKey } from "@solana/web3.js"; +import BigNumber from "bignumber.js"; +import NodeCache from "node-cache"; +import { DeriveKeyProvider } from "./deriveKeyProvider"; +import type { RemoteAttestationQuote } from "@elizaos/core"; +// Provider configuration +const PROVIDER_CONFIG = { + BIRDEYE_API: "https://public-api.birdeye.so", + MAX_RETRIES: 3, + RETRY_DELAY: 2000, + DEFAULT_RPC: "https://api.mainnet-beta.solana.com", + TOKEN_ADDRESSES: { + SOL: "So11111111111111111111111111111111111111112", + BTC: "3NZ9JMVBmGAqocybic2c7LQCJScmgsAZ6vQqTDzcqmJh", + ETH: "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs", + }, +}; + +export interface Item { + name: string; + address: string; + symbol: string; + decimals: number; + balance: string; + uiAmount: string; + priceUsd: string; + valueUsd: string; + valueSol?: string; +} + +interface WalletPortfolio { + totalUsd: string; + totalSol?: string; + items: Array; +} + +interface _BirdEyePriceData { + data: { + [key: string]: { + price: number; + priceChange24h: number; + }; + }; +} + +interface Prices { + solana: { usd: string }; + bitcoin: { usd: string }; + ethereum: { usd: string }; +} + +export class WalletProvider { + private cache: NodeCache; + + constructor( + private connection: Connection, + private walletPublicKey: PublicKey + ) { + this.cache = new NodeCache({ stdTTL: 300 }); // Cache TTL set to 5 minutes + } + + private async fetchWithRetry( + runtime, + url: string, + options: RequestInit = {} + ): Promise { + let lastError: Error; + + for (let i = 0; i < PROVIDER_CONFIG.MAX_RETRIES; i++) { + try { + const apiKey = runtime.getSetting("BIRDEYE_API_KEY"); + const response = await fetch(url, { + ...options, + headers: { + Accept: "application/json", + "x-chain": "solana", + "X-API-KEY": apiKey || "", + ...options.headers, + }, + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `HTTP error! status: ${response.status}, message: ${errorText}` + ); + } + + const data = await response.json(); + return data; + } catch (error) { + elizaLogger.error(`Attempt ${i + 1} failed:`, error); + lastError = error; + if (i < PROVIDER_CONFIG.MAX_RETRIES - 1) { + const delay = PROVIDER_CONFIG.RETRY_DELAY * Math.pow(2, i); + await new Promise((resolve) => setTimeout(resolve, delay)); + continue; + } + } + } + + elizaLogger.error( + "All attempts failed. Throwing the last error:", + lastError + ); + throw lastError; + } + + async fetchPortfolioValue(runtime): Promise { + try { + const cacheKey = `portfolio-${this.walletPublicKey.toBase58()}`; + const cachedValue = this.cache.get(cacheKey); + + if (cachedValue) { + elizaLogger.log("Cache hit for fetchPortfolioValue"); + return cachedValue; + } + elizaLogger.log("Cache miss for fetchPortfolioValue"); + + const walletData = await this.fetchWithRetry( + runtime, + `${PROVIDER_CONFIG.BIRDEYE_API}/v1/wallet/token_list?wallet=${this.walletPublicKey.toBase58()}` + ); + + if (!walletData?.success || !walletData?.data) { + elizaLogger.error("No portfolio data available", walletData); + throw new Error("No portfolio data available"); + } + + const data = walletData.data; + const totalUsd = new BigNumber(data.totalUsd.toString()); + const prices = await this.fetchPrices(runtime); + const solPriceInUSD = new BigNumber(prices.solana.usd.toString()); + + const items = data.items.map((item: any) => ({ + ...item, + valueSol: new BigNumber(item.valueUsd || 0) + .div(solPriceInUSD) + .toFixed(6), + name: item.name || "Unknown", + symbol: item.symbol || "Unknown", + priceUsd: item.priceUsd || "0", + valueUsd: item.valueUsd || "0", + })); + + const totalSol = totalUsd.div(solPriceInUSD); + const portfolio = { + totalUsd: totalUsd.toString(), + totalSol: totalSol.toFixed(6), + items: items.sort((a, b) => + new BigNumber(b.valueUsd) + .minus(new BigNumber(a.valueUsd)) + .toNumber() + ), + }; + this.cache.set(cacheKey, portfolio); + return portfolio; + } catch (error) { + elizaLogger.error("Error fetching portfolio:", error); + throw error; + } + } + + async fetchPrices(runtime): Promise { + try { + const cacheKey = "prices"; + const cachedValue = this.cache.get(cacheKey); + + if (cachedValue) { + elizaLogger.log("Cache hit for fetchPrices"); + return cachedValue; + } + elizaLogger.log("Cache miss for fetchPrices"); + + const { SOL, BTC, ETH } = PROVIDER_CONFIG.TOKEN_ADDRESSES; + const tokens = [SOL, BTC, ETH]; + const prices: Prices = { + solana: { usd: "0" }, + bitcoin: { usd: "0" }, + ethereum: { usd: "0" }, + }; + + for (const token of tokens) { + const response = await this.fetchWithRetry( + runtime, + `${PROVIDER_CONFIG.BIRDEYE_API}/defi/price?address=${token}`, + { + headers: { + "x-chain": "solana", + }, + } + ); + + if (response?.data?.value) { + const price = response.data.value.toString(); + prices[ + token === SOL + ? "solana" + : token === BTC + ? "bitcoin" + : "ethereum" + ].usd = price; + } else { + elizaLogger.warn( + `No price data available for token: ${token}` + ); + } + } + + this.cache.set(cacheKey, prices); + return prices; + } catch (error) { + elizaLogger.error("Error fetching prices:", error); + throw error; + } + } + + formatPortfolio( + runtime, + portfolio: WalletPortfolio, + prices: Prices + ): string { + let output = `${runtime.character.description}\n`; + output += `Wallet Address: ${this.walletPublicKey.toBase58()}\n\n`; + + const totalUsdFormatted = new BigNumber(portfolio.totalUsd).toFixed(2); + const totalSolFormatted = portfolio.totalSol; + + output += `Total Value: $${totalUsdFormatted} (${totalSolFormatted} SOL)\n\n`; + output += "Token Balances:\n"; + + const nonZeroItems = portfolio.items.filter((item) => + new BigNumber(item.uiAmount).isGreaterThan(0) + ); + + if (nonZeroItems.length === 0) { + output += "No tokens found with non-zero balance\n"; + } else { + for (const item of nonZeroItems) { + const valueUsd = new BigNumber(item.valueUsd).toFixed(2); + output += `${item.name} (${item.symbol}): ${new BigNumber( + item.uiAmount + ).toFixed(6)} ($${valueUsd} | ${item.valueSol} SOL)\n`; + } + } + + output += "\nMarket Prices:\n"; + output += `SOL: $${new BigNumber(prices.solana.usd).toFixed(2)}\n`; + output += `BTC: $${new BigNumber(prices.bitcoin.usd).toFixed(2)}\n`; + output += `ETH: $${new BigNumber(prices.ethereum.usd).toFixed(2)}\n`; + + return output; + } + + async getFormattedPortfolio(runtime): Promise { + try { + const [portfolio, prices] = await Promise.all([ + this.fetchPortfolioValue(runtime), + this.fetchPrices(runtime), + ]); + + return this.formatPortfolio(runtime, portfolio, prices); + } catch (error) { + elizaLogger.error("Error generating portfolio report:", error); + return "Unable to fetch wallet information. Please try again later."; + } + } +} + +const walletProvider: Provider = { + get: async ( + runtime: IAgentRuntime, + _message: Memory, + _state?: State + ): Promise => { + const agentId = runtime.agentId; + const teeMode = runtime.getSetting("TEE_MODE"); + const deriveKeyProvider = new DeriveKeyProvider(teeMode); + try { + // Validate wallet configuration + if (!runtime.getSetting("WALLET_SECRET_SALT")) { + elizaLogger.error( + "Wallet secret salt is not configured in settings" + ); + return ""; + } + + let publicKey: PublicKey; + try { + const derivedKeyPair: { + keypair: Keypair; + attestation: RemoteAttestationQuote; + } = await deriveKeyProvider.deriveEd25519Keypair( + runtime.getSetting("WALLET_SECRET_SALT"), + "solana", + agentId + ); + publicKey = derivedKeyPair.keypair.publicKey; + elizaLogger.log("Wallet Public Key: ", publicKey.toBase58()); + } catch (error) { + elizaLogger.error("Error creating PublicKey:", error); + return ""; + } + + const connection = new Connection(PROVIDER_CONFIG.DEFAULT_RPC); + const provider = new WalletProvider(connection, publicKey); + + const porfolio = await provider.getFormattedPortfolio(runtime); + return porfolio; + } catch (error) { + elizaLogger.error("Error in wallet provider:", error.message); + return `Failed to fetch wallet information: ${error instanceof Error ? error.message : "Unknown error"}`; + } + }, +}; + +// Module exports +export { walletProvider }; diff --git a/packages/plugin-tee/src/services/teeLogManager.ts b/packages/plugin-tee/src/services/teeLogManager.ts new file mode 100644 index 00000000000..fcba701a810 --- /dev/null +++ b/packages/plugin-tee/src/services/teeLogManager.ts @@ -0,0 +1,105 @@ +import { SgxAttestationProvider } from "../providers/sgxAttestationProvider"; +import { RemoteAttestationProvider as TdxAttestationProvider } from "../providers/remoteAttestationProvider"; +import { TEEMode, TeeType, type TeeLogDAO, type TeeAgent, type TeeLog, type TeeLogQuery, type TeePageQuery } from "@elizaos/core"; +import elliptic from "elliptic"; +import { v4 } from "uuid"; + +export class TeeLogManager { + private teeLogDAO: TeeLogDAO; + private teeType: TeeType; + private teeMode: TEEMode; // Only used for plugin-tee with TDX dstack + + // Map of agentId to its key pair + // These keypairs only store in memory. + // When the agent restarts, we will generate new keypair. + private keyPairs: Map = new Map(); + + constructor(teeLogDAO: TeeLogDAO, teeType: TeeType, teeMode: TEEMode) { + this.teeLogDAO = teeLogDAO; + this.teeType = teeType; + this.teeMode = teeMode; + } + + public async registerAgent(agentId: string, agentName: string): Promise { + if (!agentId) { + throw new Error("Agent ID is required"); + } + + const keyPair = this.generateKeyPair(); + this.keyPairs.set(agentId, keyPair); + + const publicKey = keyPair.getPublic().encode('hex', true); + const attestation = await this.generateAttestation(publicKey); + + const new_agent = { + id: v4(), + agentId, + agentName: agentName || "", + createdAt: new Date().getTime(), + publicKey, + attestation, + }; + + console.log("registerAgent new_agent", new_agent); + + return this.teeLogDAO.addAgent(new_agent); + } + + public async getAllAgents(): Promise { + return this.teeLogDAO.getAllAgents(); + } + + public async getAgent(agentId: string): Promise { + return this.teeLogDAO.getAgent(agentId); + } + + public async log(agentId: string, roomId: string, userId: string, type: string, content: string): Promise { + const keyPair = this.keyPairs.get(agentId); + if (!keyPair) { + throw new Error(`Agent ${agentId} not found`); + } + + const timestamp = new Date().getTime(); + + // Join the information into a single string + const messageToSign = `${agentId}|${roomId}|${userId}|${type}|${content}|${timestamp}`; + + // Sign the joined message + const signature = "0x" + keyPair.sign(messageToSign).toDER('hex'); + + return this.teeLogDAO.addLog({ + id: v4(), + agentId, + roomId, + userId, + type, + content, + timestamp, + signature, + }); + } + + public async getLogs(query: TeeLogQuery, page: number, pageSize: number): Promise> { + return this.teeLogDAO.getPagedLogs(query, page, pageSize); + } + + public generateKeyPair(): elliptic.ec.KeyPair { + const ec = new elliptic.ec('secp256k1'); + const key = ec.genKeyPair(); + return key; + } + + public async generateAttestation(userReport: string): Promise { + if (this.teeType === TeeType.SGX_GRAMINE) { + const sgxAttestationProvider = new SgxAttestationProvider(); + const sgxAttestation = await sgxAttestationProvider.generateAttestation(userReport); + return JSON.stringify(sgxAttestation); + } else if (this.teeType === TeeType.TDX_DSTACK) { + const tdxAttestationProvider = new TdxAttestationProvider(this.teeMode); + const tdxAttestation = await tdxAttestationProvider.generateAttestation(userReport); + return JSON.stringify(tdxAttestation); + } else { + throw new Error("Invalid TEE type"); + } + } +} diff --git a/packages/plugin-tee/src/services/teeLogService.ts b/packages/plugin-tee/src/services/teeLogService.ts new file mode 100644 index 00000000000..1178b4f92d4 --- /dev/null +++ b/packages/plugin-tee/src/services/teeLogService.ts @@ -0,0 +1,124 @@ +import { type IAgentRuntime, Service, ServiceType, type ITeeLogService, TeeType, type TeeLogDAO, type TeeAgent, type TeeLog, type TeeLogQuery, type TeePageQuery, TEEMode } from "@elizaos/core"; +import { SqliteTeeLogDAO } from "../adapters/sqliteDAO"; +import { TeeLogManager } from "./teeLogManager"; +import Database from "better-sqlite3"; +import path from "path"; + +export class TeeLogService extends Service implements ITeeLogService { + private dbPath: string; + + private initialized = false; + private enableTeeLog = false; + private teeType: TeeType; + private teeMode: TEEMode = TEEMode.OFF; // Only used for plugin-tee with TDX dstack + + private teeLogDAO: TeeLogDAO; + private teeLogManager: TeeLogManager; + + + getInstance(): ITeeLogService { + return this; + } + + static get serviceType(): ServiceType { + return ServiceType.TEE; + } + + async initialize(runtime: IAgentRuntime): Promise { + if (this.initialized) { + return; + } + + const enableValues = ["true", "1", "yes", "enable", "enabled", "on"]; + + const enableTeeLog = runtime.getSetting("ENABLE_TEE_LOG"); + if (enableTeeLog === null) { + throw new Error("ENABLE_TEE_LOG is not set."); + } + this.enableTeeLog = enableValues.includes(enableTeeLog.toLowerCase()); + if (!this.enableTeeLog) { + console.log("TEE log is not enabled."); + return; + } + + const runInSgx = runtime.getSetting("SGX"); + const teeMode = runtime.getSetting("TEE_MODE"); + const walletSecretSalt = runtime.getSetting("WALLET_SECRET_SALT"); + + this.teeMode = teeMode ? TEEMode[teeMode as keyof typeof TEEMode] : TEEMode.OFF; + + const useSgxGramine = runInSgx && enableValues.includes(runInSgx.toLowerCase()); + const useTdxDstack = teeMode && teeMode !== TEEMode.OFF && walletSecretSalt; + + if (useSgxGramine && useTdxDstack) { + throw new Error("Cannot configure both SGX and TDX at the same time."); + } else if (useSgxGramine) { + this.teeType = TeeType.SGX_GRAMINE; + } else if (useTdxDstack) { + this.teeType = TeeType.TDX_DSTACK; + } else { + throw new Error("Invalid TEE configuration."); + } + + const dbPathSetting = runtime.getSetting("TEE_LOG_DB_PATH"); + this.dbPath = dbPathSetting || path.resolve("data/tee_log.sqlite"); + + const db = new Database(this.dbPath); + this.teeLogDAO = new SqliteTeeLogDAO(db); + await this.teeLogDAO.initialize(); + this.teeLogManager = new TeeLogManager(this.teeLogDAO, this.teeType, this.teeMode); + + const isRegistered = await this.teeLogManager.registerAgent( + runtime?.agentId, + runtime?.character?.name, + ); + if (!isRegistered) { + throw new Error(`Failed to register agent ${runtime.agentId}`); + } + + this.initialized = true; + } + + async log(agentId: string, roomId: string, userId: string, type: string, content: string): Promise { + if (!this.enableTeeLog) { + return false; + } + + return this.teeLogManager.log(agentId, roomId, userId, type, content); + } + + async getAllAgents(): Promise { + if (!this.enableTeeLog) { + return []; + } + + return this.teeLogManager.getAllAgents(); + } + + async getAgent(agentId: string): Promise { + if (!this.enableTeeLog) { + return undefined; + } + + return this.teeLogManager.getAgent(agentId); + } + + async getLogs(query: TeeLogQuery, page: number, pageSize: number): Promise> { + if (!this.enableTeeLog) { + return { + data: [], + total: 0, + page: page, + pageSize: pageSize, + }; + } + + return this.teeLogManager.getLogs(query, page, pageSize); + } + + async generateAttestation(userReport: string): Promise { + return this.teeLogManager.generateAttestation(userReport); + } +} + +export default TeeLogService; diff --git a/packages/plugin-tee/src/types.ts b/packages/plugin-tee/src/types.ts new file mode 100644 index 00000000000..3f4bee57aac --- /dev/null +++ b/packages/plugin-tee/src/types.ts @@ -0,0 +1,66 @@ +export enum TeeType { + SGX_GRAMINE = "sgx_gramine", + TDX_DSTACK = "tdx_dstack", +} + +// Represents a log entry in the TeeLog table, containing details about agent activities. +export interface TeeLog { + id: string; + agentId: string; + roomId: string; + userId: string; + type: string; + content: string; + timestamp: number; + signature: string; +} + +export interface TeeLogQuery { + agentId?: string; + roomId?: string; + userId?: string; + type?: string; + containsContent?: string; + startTimestamp?: number; + endTimestamp?: number; +} + +// Represents an agent in the TeeAgent table, containing details about the agent. +export interface TeeAgent { + id: string; // Primary key + // Allow duplicate agentId. + // This is to support the case where the same agentId is registered multiple times. + // Each time the agent restarts, we will generate a new keypair and attestation. + agentId: string; + agentName: string; + createdAt: number; + publicKey: string; + attestation: string; +} + +export interface TeePageQuery { + page: number; + pageSize: number; + total?: number; + data?: Result; +} + +export abstract class TeeLogDAO { + db: DB; + + abstract initialize(): Promise; + + abstract addLog(log: TeeLog): Promise; + + abstract getPagedLogs( + query: TeeLogQuery, + page: number, + pageSize: number + ): Promise>; + + abstract addAgent(agent: TeeAgent): Promise; + + abstract getAgent(agentId: string): Promise; + + abstract getAllAgents(): Promise; +} \ No newline at end of file diff --git a/packages/plugin-tee/tsconfig.json b/packages/plugin-tee/tsconfig.json new file mode 100644 index 00000000000..18c600eec05 --- /dev/null +++ b/packages/plugin-tee/tsconfig.json @@ -0,0 +1,13 @@ +{ + "extends": "../core/tsconfig.json", + "compilerOptions": { + "outDir": "dist", + "rootDir": "src", + "types": [ + "node" + ] + }, + "include": [ + "src/**/*.ts", + ] +} \ No newline at end of file diff --git a/packages/plugin-tee/tsup.config.ts b/packages/plugin-tee/tsup.config.ts new file mode 100644 index 00000000000..1a96f24afa1 --- /dev/null +++ b/packages/plugin-tee/tsup.config.ts @@ -0,0 +1,21 @@ +import { defineConfig } from "tsup"; + +export default defineConfig({ + entry: ["src/index.ts"], + outDir: "dist", + sourcemap: true, + clean: true, + format: ["esm"], // Ensure you're targeting CommonJS + external: [ + "dotenv", // Externalize dotenv to prevent bundling + "fs", // Externalize fs to use Node.js built-in module + "path", // Externalize other built-ins if necessary + "@reflink/reflink", + "@node-llama-cpp", + "https", + "http", + "agentkeepalive", + "safe-buffer", + // Add other modules you want to externalize + ], +});