From e6b5bbe828e85624980bd4b9c3acd20e6b4835e3 Mon Sep 17 00:00:00 2001 From: Shakker Nerd <165377636+shakkernerd@users.noreply.github.com> Date: Thu, 9 Jan 2025 05:38:41 +0000 Subject: [PATCH] Revert "fix: Update Key Derive in TEE" --- docs/docs/advanced/eliza-in-tee.md | 10 +- packages/plugin-evm/src/providers/wallet.ts | 2 +- packages/plugin-solana/src/keypairUtils.ts | 2 +- packages/plugin-tee/package.json | 3 +- .../src/providers/deriveKeyProvider.ts | 24 +--- .../src/providers/walletProvider.ts | 2 +- .../plugin-tee/src/tests/deriveKey.test.ts | 135 ------------------ .../src/tests/remoteAttestation.test.ts | 81 ----------- .../src/tests/remoteAttestationAction.test.ts | 103 ------------- 9 files changed, 8 insertions(+), 354 deletions(-) delete mode 100644 packages/plugin-tee/src/tests/deriveKey.test.ts delete mode 100644 packages/plugin-tee/src/tests/remoteAttestation.test.ts delete mode 100644 packages/plugin-tee/src/tests/remoteAttestationAction.test.ts diff --git a/docs/docs/advanced/eliza-in-tee.md b/docs/docs/advanced/eliza-in-tee.md index c4ef9cd6392..76d2e9e2851 100644 --- a/docs/docs/advanced/eliza-in-tee.md +++ b/docs/docs/advanced/eliza-in-tee.md @@ -57,22 +57,16 @@ Example usage: const provider = new DeriveKeyProvider(teeMode); // For Solana const { keypair, attestation } = await provider.deriveEd25519Keypair( + "/", secretSalt, - "solana", agentId, ); // For EVM const { keypair, attestation } = await provider.deriveEcdsaKeypair( + "/", secretSalt, - "evm", agentId, ); - -// For raw key derivation -const rawKey = await provider.deriveRawKey( - secretSalt, - "raw", -); ``` --- diff --git a/packages/plugin-evm/src/providers/wallet.ts b/packages/plugin-evm/src/providers/wallet.ts index 509cd137cd9..ae05779dabf 100644 --- a/packages/plugin-evm/src/providers/wallet.ts +++ b/packages/plugin-evm/src/providers/wallet.ts @@ -292,8 +292,8 @@ export const initWalletProvider = async (runtime: IAgentRuntime) => { const deriveKeyProvider = new DeriveKeyProvider(teeMode); const deriveKeyResult = await deriveKeyProvider.deriveEcdsaKeypair( + "/", walletSecretSalt, - "evm", runtime.agentId ); return new WalletProvider( diff --git a/packages/plugin-solana/src/keypairUtils.ts b/packages/plugin-solana/src/keypairUtils.ts index f8405424bc1..6dbacdb8414 100644 --- a/packages/plugin-solana/src/keypairUtils.ts +++ b/packages/plugin-solana/src/keypairUtils.ts @@ -30,8 +30,8 @@ export async function getWalletKey( const deriveKeyProvider = new DeriveKeyProvider(teeMode); const deriveKeyResult = await deriveKeyProvider.deriveEd25519Keypair( + "/", walletSecretSalt, - "solana", runtime.agentId ); diff --git a/packages/plugin-tee/package.json b/packages/plugin-tee/package.json index 0e1297bee12..fbf5234d190 100644 --- a/packages/plugin-tee/package.json +++ b/packages/plugin-tee/package.json @@ -32,8 +32,7 @@ "scripts": { "build": "tsup --format esm --dts", "dev": "tsup --format esm --dts --watch", - "lint": "eslint --fix --cache .", - "test": "vitest run" + "lint": "eslint --fix --cache ." }, "peerDependencies": { "whatwg-url": "7.1.0" diff --git a/packages/plugin-tee/src/providers/deriveKeyProvider.ts b/packages/plugin-tee/src/providers/deriveKeyProvider.ts index 4284479326e..96430f23586 100644 --- a/packages/plugin-tee/src/providers/deriveKeyProvider.ts +++ b/packages/plugin-tee/src/providers/deriveKeyProvider.ts @@ -72,12 +72,6 @@ class DeriveKeyProvider { 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 @@ -100,13 +94,6 @@ class DeriveKeyProvider { } } - /** - * 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, @@ -143,13 +130,6 @@ class DeriveKeyProvider { } } - /** - * 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, @@ -204,13 +184,13 @@ const deriveKeyProvider: Provider = { 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({ diff --git a/packages/plugin-tee/src/providers/walletProvider.ts b/packages/plugin-tee/src/providers/walletProvider.ts index 0d08c0e07bf..a31cc65ceac 100644 --- a/packages/plugin-tee/src/providers/walletProvider.ts +++ b/packages/plugin-tee/src/providers/walletProvider.ts @@ -299,8 +299,8 @@ const walletProvider: Provider = { keypair: Keypair; attestation: RemoteAttestationQuote; } = await deriveKeyProvider.deriveEd25519Keypair( + "/", runtime.getSetting("WALLET_SECRET_SALT"), - "solana", agentId ); publicKey = derivedKeyPair.keypair.publicKey; diff --git a/packages/plugin-tee/src/tests/deriveKey.test.ts b/packages/plugin-tee/src/tests/deriveKey.test.ts deleted file mode 100644 index f12beab80ca..00000000000 --- a/packages/plugin-tee/src/tests/deriveKey.test.ts +++ /dev/null @@ -1,135 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { DeriveKeyProvider } from '../providers/deriveKeyProvider'; -import { TappdClient } from '@phala/dstack-sdk'; -import { TEEMode } from '../types/tee'; -import { Keypair } from '@solana/web3.js'; -import { privateKeyToAccount } from 'viem/accounts'; - -// 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 agentId = 'test-agent'; - - const result = await provider.deriveEd25519Keypair(path, subject, agentId); - - expect(result).toHaveProperty('keypair'); - expect(result).toHaveProperty('attestation'); - expect(result.keypair.publicKey.toBase58()).toBe('mock-solana-public-key'); - }); - }); - - describe('deriveEcdsaKeypair', () => { - let provider: DeriveKeyProvider; - - beforeEach(() => { - provider = new DeriveKeyProvider(TEEMode.LOCAL); - }); - - it('should derive ECDSA keypair successfully', async () => { - const path = 'test-path'; - const subject = 'test-subject'; - const agentId = 'test-agent'; - - const result = await provider.deriveEcdsaKeypair(path, subject, agentId); - - expect(result).toHaveProperty('keypair'); - expect(result).toHaveProperty('attestation'); - expect(result.keypair.address).toBe('mock-evm-address'); - }); - }); -}); \ 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 deleted file mode 100644 index d9c401f2568..00000000000 --- a/packages/plugin-tee/src/tests/remoteAttestation.test.ts +++ /dev/null @@ -1,81 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { RemoteAttestationProvider } from '../providers/remoteAttestationProvider'; -import { TappdClient } from '@phala/dstack-sdk'; -import { TEEMode } from '../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 deleted file mode 100644 index 11b91037aa8..00000000000 --- a/packages/plugin-tee/src/tests/remoteAttestationAction.test.ts +++ /dev/null @@ -1,103 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { remoteAttestationAction } from '../actions/remoteAttestation'; -import { RemoteAttestationProvider } from '../providers/remoteAttestationProvider'; - -// Mock dependencies -vi.mock('../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(), - getConversation: vi.fn(), - 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, - {}, - {}, - {}, - 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(agentMessage.user).toBe('{{user2}}'); - expect(agentMessage.content.text).toBe('Of course, one second...'); - expect(agentMessage.content.action).toBe('REMOTE_ATTESTATION'); - }); - }); -}); \ No newline at end of file