Merge pull request #640 from Phala-Network/tee-plugin

fix: Add docs, update providers for TEE Plugin

docs/docs/packages/plugins.md

@@ -294,9 +294,7 @@ Supported networks:
    - `COINBASE_API_KEY`: API key for Coinbase SDK.
    - `COINBASE_PRIVATE_KEY`: Private key for secure transactions.
 
----
-
-### Wallet Management
+**Wallet Management:**
 
 The plugin automatically handles wallet creation or uses an existing wallet if the required details are provided during the first run.
 
@@ -313,9 +311,8 @@ The plugin automatically handles wallet creation or uses an existing wallet if t
    - Provide `COINBASE_GENERATED_WALLET_HEX_SEED` and `COINBASE_GENERATED_WALLET_ID` via `runtime.character.settings.secrets` or environment variables.
    - The plugin will **import the wallet** and use it for processing mass payouts.
 
----
 
-### Required Configurations
+**Required Configurations:**
 
 The following configurations must be provided for wallet management:
 
@@ -324,9 +321,8 @@ The following configurations must be provided for wallet management:
   - `COINBASE_GENERATED_WALLET_ID`: Unique wallet ID.
   - These variables must be securely stored in `runtime.character.settings.secrets` or as environment variables.
 
----
 
-### Wallet Creation Process
+**Wallet Creation Process:**
 
 1. **Automatic Wallet Creation**
    When no wallet details are available:
@@ -357,9 +353,7 @@ The following configurations must be provided for wallet management:
      elizaLogger.log("Imported existing wallet:", wallet.getDefaultAddress());
      ```
 
----
-
-### Example Configuration
+**Example Configuration:**
 
 #### Automatic Wallet Generation:
 
@@ -395,8 +389,6 @@ Output Log:
 [INFO] Imported existing wallet: 0x1234567890abcdef1234567890abcdef12345678
 ```
 
----
-
 3. **Example Call**
    An example of using the `SEND_MASS_PAYOUT` action:
 
@@ -438,6 +430,86 @@ When successful, a response similar to the following will be returned:
 
 ---
 
+#### 7. TEE Plugin (`@ai16z/plugin-tee`)
+
+Integrates [Dstack SDK](https://github.com/Dstack-TEE/dstack) to enable TEE (Trusted Execution Environment) functionality and deploy secure & privacy-enhanced Eliza Agents:
+
+**Providers:**
+
+- `deriveKeyProvider` - Allows for secure key derivation within a TEE environment. It supports deriving keys for both Solana (Ed25519) and Ethereum (ECDSA) chains.
+- `remoteAttestationProvider` - Generate a Remote Attestation Quote based on `report_data`.
+
+**DeriveKeyProvider Usage**
+
+```typescript
+import { DeriveKeyProvider } from "@ai16z/plugin-tee";
+
+// Initialize the provider
+const provider = new DeriveKeyProvider();
+
+// Derive a raw key
+try {
+    const rawKey = await provider.rawDeriveKey(
+        "/path/to/derive",
+        "subject-identifier"
+    );
+    // rawKey is a DeriveKeyResponse that can be used for further processing
+    // to get the uint8Array do the following
+    const rawKeyArray = rawKey.asUint8Array()
+} catch (error) {
+    console.error("Raw key derivation failed:", error);
+}
+
+// Derive a Solana keypair (Ed25519)
+try {
+    const solanaKeypair = await provider.deriveEd25519Keypair(
+        "/path/to/derive",
+        "subject-identifier"
+    );
+    // solanaKeypair can now be used for Solana operations
+} catch (error) {
+    console.error("Solana key derivation failed:", error);
+}
+
+// Derive an Ethereum keypair (ECDSA)
+try {
+    const evmKeypair = await provider.deriveEcdsaKeypair(
+        "/path/to/derive",
+        "subject-identifier"
+    );
+    // evmKeypair can now be used for Ethereum operations
+} catch (error) {
+    console.error("EVM key derivation failed:", error);
+}
+```
+
+**RemoteAttestationProvider Usage**
+
+```typescript
+import { RemoteAttestationProvider } from "@ai16z/plugin-tee";
+// Initialize the provider
+const provider = new RemoteAttestationProvider();
+// Generate Remote Attestation
+try {
+    const attestation = await provider.generateAttestation("your-report-data");
+    console.log("Attestation:", attestation);
+} catch (error) {
+    console.error("Failed to generate attestation:", error);
+}
+```
+
+**Configuration**
+
+When using the provider through the runtime environment, ensure the following settings are configured:
+
+```env
+ # Optional, for simulator purposes if testing on mac or windows. Leave empty for Linux x86 machines.
+DSTACK_SIMULATOR_ENDPOINT="http://host.docker.internal:8090"
+WALLET_SECRET_SALT=your-secret-salt // Required to single agent deployments
+```
+
+___
+
 ### Writing Custom Plugins
 
 Create a new plugin by implementing the Plugin interface:

packages/plugin-tee/README.md

@@ -8,20 +8,49 @@ This plugin includes several providers for handling different TEE-related operat
 
 ### DeriveKeyProvider
 
-The `DeriveKeyProvider` allows for secure key derivation within a TEE environment.
+The `DeriveKeyProvider` allows for secure key derivation within a TEE environment. It supports deriving keys for both Solana (Ed25519) and Ethereum (ECDSA) chains.
 
 #### Usage
 
 ```typescript
 import { DeriveKeyProvider } from "@ai16z/plugin-tee";
+
 // Initialize the provider
 const provider = new DeriveKeyProvider();
-// Derive a key
+
+// Derive a raw key
+try {
+    const rawKey = await provider.rawDeriveKey(
+        "/path/to/derive",
+        "subject-identifier"
+    );
+    // rawKey is a DeriveKeyResponse that can be used for further processing
+    // to get the uint8Array do the following
+    const rawKeyArray = rawKey.asUint8Array()
+} catch (error) {
+    console.error("Raw key derivation failed:", error);
+}
+
+// Derive a Solana keypair (Ed25519)
+try {
+    const solanaKeypair = await provider.deriveEd25519Keypair(
+        "/path/to/derive",
+        "subject-identifier"
+    );
+    // solanaKeypair can now be used for Solana operations
+} catch (error) {
+    console.error("Solana key derivation failed:", error);
+}
+
+// Derive an Ethereum keypair (ECDSA)
 try {
-    const keypair = await provider.deriveKey("/path/to/derive", "subject-identifier");
-    // keypair can now be used for cryptographic operations
+    const evmKeypair = await provider.deriveEcdsaKeypair(
+        "/path/to/derive",
+        "subject-identifier"
+    );
+    // evmKeypair can now be used for Ethereum operations
 } catch (error) {
-    console.error("Key derivation failed:", error);
+    console.error("EVM key derivation failed:", error);
 }
 ```
 

packages/plugin-tee/src/index.ts

@@ -1,7 +1,6 @@
 import { Plugin } from "@ai16z/eliza";
 import { remoteAttestationProvider } from "./providers/remoteAttestationProvider";
 import { deriveKeyProvider } from "./providers/deriveKeyProvider";
-import { walletProvider } from "./providers/walletProvider";
 
 export const teePlugin: Plugin = {
     name: "tee",
@@ -17,7 +16,6 @@ export const teePlugin: Plugin = {
         /* custom providers */
         remoteAttestationProvider,
         deriveKeyProvider,
-        walletProvider,
     ],
     services: [
         /* custom services */

packages/plugin-tee/src/providers/deriveKeyProvider.ts

@@ -1,7 +1,9 @@
 import { IAgentRuntime, Memory, Provider, State } from "@ai16z/eliza";
 import { Keypair } from "@solana/web3.js";
 import crypto from "crypto";
-import { TappdClient } from "@phala/dstack-sdk";
+import { DeriveKeyResponse, TappdClient } from "@phala/dstack-sdk";
+import { privateKeyToAccount } from "viem/accounts"
+import { PrivateKeyAccount, keccak256 } from "viem";
 
 class DeriveKeyProvider {
     private client: TappdClient;
@@ -10,29 +12,70 @@ class DeriveKeyProvider {
         this.client = endpoint ? new TappdClient(endpoint) : new TappdClient();
     }
 
-    async deriveKey(path: string, subject: string): Promise<Keypair> {
+    async rawDeriveKey(path: string, subject: string): Promise<DeriveKeyResponse> {
         try {
             if (!path || !subject) {
-                console.error("Path and Subject are required for key derivation");
+                console.error(
+                    "Path and Subject are required for key derivation"
+                );
+            }
+
+            console.log("Deriving Raw Key in TEE...");
+            const derivedKey = await this.client.deriveKey(path, subject);
+
+            console.log("Raw Key Derived Successfully!");
+            return derivedKey;
+        } catch (error) {
+            console.error("Error deriving raw key:", error);
+            throw error;
+        }
+    }
+
+    async deriveEd25519Keypair(path: string, subject: string): Promise<Keypair> {
+        try {
+            if (!path || !subject) {
+                console.error(
+                    "Path and Subject are required for key derivation"
+                );
             }
 
             console.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));
-            
+
             console.log("Key Derived Successfully!");
             return keypair;
         } catch (error) {
             console.error("Error deriving key:", error);
             throw error;
         }
     }
+
+    async deriveEcdsaKeypair(path: string, subject: string): Promise<PrivateKeyAccount> {
+        try {
+            if (!path || !subject) {
+                console.error(
+                    "Path and Subject are required for key derivation"
+                );
+            }
+
+            console.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);
+            console.log("ECDSA Key Derived Successfully!");
+            return keypair;
+        } catch (error) {
+            console.error("Error deriving ecdsa key:", error);
+            throw error;
+        }
+    }
 }
 
 const deriveKeyProvider: Provider = {
@@ -52,7 +95,12 @@ const deriveKeyProvider: Provider = {
             try {
                 const secretSalt =
                     runtime.getSetting("WALLET_SECRET_SALT") || "secret_salt";
-                    return await provider.deriveKey("/", secretSalt);
+                const solanaKeypair = await provider.deriveEd25519Keypair("/", secretSalt);
+                const evmKeypair = await provider.deriveEcdsaKeypair("/", secretSalt);
+                return JSON.stringify({
+                    solana: solanaKeypair.publicKey,
+                    evm: evmKeypair.address
+                })
             } catch (error) {
                 console.error("Error creating PublicKey:", error);
                 return "";

packages/plugin-tee/src/providers/remoteAttestationProvider.ts

@@ -27,7 +27,7 @@ const remoteAttestationProvider: Provider = {
         const endpoint = runtime.getSetting("DSTACK_SIMULATOR_ENDPOINT");
         const provider = new RemoteAttestationProvider(endpoint);
         const agentId = runtime.agentId;
-        
+
         try {
             const attestation = await provider.generateAttestation(agentId);
             return `Your Agent's remote attestation is: ${attestation}`;

packages/plugin-tee/src/providers/walletProvider.ts

@@ -1,8 +1,9 @@
+/* This is an example of how WalletProvider can use DeriveKeyProvider to generate a Solana Keypair */
 import { IAgentRuntime, Memory, Provider, State } from "@ai16z/eliza";
 import { Connection, Keypair, PublicKey } from "@solana/web3.js";
 import BigNumber from "bignumber.js";
 import NodeCache from "node-cache";
-import { deriveKeyProvider } from "./deriveKeyProvider";
+import { DeriveKeyProvider } from "./deriveKeyProvider";
 // Provider configuration
 const PROVIDER_CONFIG = {
     BIRDEYE_API: "https://public-api.birdeye.so",
@@ -282,11 +283,8 @@ const walletProvider: Provider = {
 
             let publicKey: PublicKey;
             try {
-                const derivedKeyPair: Keypair = await deriveKeyProvider.get(
-                    runtime,
-                    _message,
-                    _state
-                );
+                const deriveKeyProvider = new DeriveKeyProvider();
+                const derivedKeyPair: Keypair = await deriveKeyProvider.deriveEd25519Keypair("/", runtime.getSetting("WALLET_SECRET_SALT"));
                 publicKey = derivedKeyPair.publicKey;
                 console.log("Wallet Public Key: ", publicKey.toBase58());
             } catch (error) {