Skip to content

Commit

Permalink
Merge pull request elizaOS#1571 from samuveth/samuv/env-tee
Browse files Browse the repository at this point in the history
feat: add TEE support for plugin-env
  • Loading branch information
odilitime authored Dec 30, 2024
2 parents db8aa92 + 71d89c0 commit 16d23ea
Show file tree
Hide file tree
Showing 6 changed files with 139 additions and 40 deletions.
45 changes: 23 additions & 22 deletions packages/plugin-evm/package.json
Original file line number Diff line number Diff line change
@@ -1,24 +1,25 @@
{
"name": "@elizaos/plugin-evm",
"version": "0.1.7-alpha.2",
"main": "dist/index.js",
"type": "module",
"types": "dist/index.d.ts",
"dependencies": {
"@elizaos/core": "workspace:*",
"@lifi/data-types": "5.15.5",
"@lifi/sdk": "3.4.1",
"@lifi/types": "16.3.0",
"tsup": "8.3.5",
"viem": "2.21.53"
},
"scripts": {
"build": "tsup --format esm --dts",
"dev": "tsup --format esm --dts --watch",
"test": "vitest run",
"lint": "eslint --fix --cache ."
},
"peerDependencies": {
"whatwg-url": "7.1.0"
}
"name": "@elizaos/plugin-evm",
"version": "0.1.7-alpha.1",
"main": "dist/index.js",
"type": "module",
"types": "dist/index.d.ts",
"dependencies": {
"@elizaos/core": "workspace:*",
"@elizaos/plugin-tee": "workspace:*",
"@lifi/data-types": "5.15.5",
"@lifi/sdk": "3.4.1",
"@lifi/types": "16.3.0",
"tsup": "8.3.5",
"viem": "2.21.53"
},
"scripts": {
"build": "tsup --format esm --dts",
"dev": "tsup --format esm --dts --watch",
"test": "vitest run",
"lint": "eslint --fix --cache ."
},
"peerDependencies": {
"whatwg-url": "7.1.0"
}
}
2 changes: 1 addition & 1 deletion packages/plugin-evm/src/actions/bridge.ts
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ export const bridgeAction = {
callback?: any
) => {
console.log("Bridge action handler called");
const walletProvider = initWalletProvider(runtime);
const walletProvider = await initWalletProvider(runtime);
const action = new BridgeAction(walletProvider);

// Compose bridge context
Expand Down
2 changes: 1 addition & 1 deletion packages/plugin-evm/src/actions/swap.ts
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ export const swapAction = {
callback?: any
) => {
console.log("Swap action handler called");
const walletProvider = initWalletProvider(runtime);
const walletProvider = await initWalletProvider(runtime);
const action = new SwapAction(walletProvider);

// Compose swap context
Expand Down
2 changes: 1 addition & 1 deletion packages/plugin-evm/src/actions/transfer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ export const transferAction = {
callback?: HandlerCallback
) => {
console.log("Transfer action handler called");
const walletProvider = initWalletProvider(runtime);
const walletProvider = await initWalletProvider(runtime);
const action = new TransferAction(walletProvider);

// Compose transfer context
Expand Down
52 changes: 41 additions & 11 deletions packages/plugin-evm/src/providers/wallet.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import type {
PrivateKeyAccount,
} from "viem";
import * as viemChains from "viem/chains";
import { DeriveKeyProvider, TEEMode } from "@elizaos/plugin-tee";

import type { SupportedChain } from "../types";

Expand All @@ -24,8 +25,11 @@ export class WalletProvider {
chains: Record<string, Chain> = { mainnet: viemChains.mainnet };
account: PrivateKeyAccount;

constructor(privateKey: `0x${string}`, chains?: Record<string, Chain>) {
this.setAccount(privateKey);
constructor(
accountOrPrivateKey: PrivateKeyAccount | `0x${string}`,
chains?: Record<string, Chain>
) {
this.setAccount(accountOrPrivateKey);
this.setChains(chains);

if (chains && Object.keys(chains).length > 0) {
Expand Down Expand Up @@ -118,8 +122,14 @@ export class WalletProvider {
this.setCurrentChain(chainName);
}

private setAccount = (pk: `0x${string}`) => {
this.account = privateKeyToAccount(pk);
private setAccount = (
accountOrPrivateKey: PrivateKeyAccount | `0x${string}`
) => {
if (typeof accountOrPrivateKey === "string") {
this.account = privateKeyToAccount(accountOrPrivateKey);
} else {
this.account = accountOrPrivateKey;
}
};

private setChains = (chains?: Record<string, Chain>) => {
Expand Down Expand Up @@ -197,15 +207,35 @@ const genChainsFromRuntime = (
return chains;
};

export const initWalletProvider = (runtime: IAgentRuntime) => {
const privateKey = runtime.getSetting("EVM_PRIVATE_KEY");
if (!privateKey) {
throw new Error("EVM_PRIVATE_KEY is missing");
}
export const initWalletProvider = async (runtime: IAgentRuntime) => {
const teeMode = runtime.getSetting("TEE_MODE") || TEEMode.OFF;

const chains = genChainsFromRuntime(runtime);

return new WalletProvider(privateKey as `0x${string}`, chains);
if (teeMode !== TEEMode.OFF) {
const walletSecretSalt = runtime.getSetting("WALLET_SECRET_SALT");
if (!walletSecretSalt) {
throw new Error(
"WALLET_SECRET_SALT required when TEE_MODE is enabled"
);
}

const deriveKeyProvider = new DeriveKeyProvider(teeMode);
const deriveKeyResult = await deriveKeyProvider.deriveEcdsaKeypair(
"/",
walletSecretSalt,
runtime.agentId
);
return new WalletProvider(deriveKeyResult.keypair, chains);
} else {
const privateKey = runtime.getSetting(
"EVM_PRIVATE_KEY"
) as `0x${string}`;
if (!privateKey) {
throw new Error("EVM_PRIVATE_KEY is missing");
}
return new WalletProvider(privateKey, chains);
}
};

export const evmWalletProvider: Provider = {
Expand All @@ -215,7 +245,7 @@ export const evmWalletProvider: Provider = {
_state?: State
): Promise<string | null> {
try {
const walletProvider = initWalletProvider(runtime);
const walletProvider = await initWalletProvider(runtime);
const address = walletProvider.getAddress();
const balance = await walletProvider.getWalletBalance();
const chain = walletProvider.getCurrentChain();
Expand Down
76 changes: 72 additions & 4 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 16d23ea

Please sign in to comment.