Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add TEE support for plugin-env #1571

Merged
merged 4 commits into from
Dec 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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.

Loading