Skip to content

Commit

Permalink
feat: give state machine ability to sign and send transactions and to…
Browse files Browse the repository at this point in the history
… have several actions in one state
  • Loading branch information
FedericoAmura committed Dec 5, 2024
1 parent be62dbc commit 47cc68a
Show file tree
Hide file tree
Showing 4 changed files with 178 additions and 63 deletions.
9 changes: 9 additions & 0 deletions packages/automation/src/lib/litActions.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
export const signWithLitActionCode = `(async () => {
const signature = await Lit.Actions.signAndCombineEcdsa({
toSign,
publicKey,
sigName,
});
Lit.Actions.setResponse({ response: signature });
})();`;
195 changes: 144 additions & 51 deletions packages/automation/src/lib/state-machine.ts
Original file line number Diff line number Diff line change
@@ -1,17 +1,7 @@
import { ethers } from 'ethers';

import {
// createSiweMessageWithRecaps,
// generateAuthSig,
LitActionResource,
// LitPKPResource,
} from '@lit-protocol/auth-helpers';
import {
LIT_ABILITY,
LIT_EVM_CHAINS,
LIT_RPC,
LIT_NETWORK,
} from '@lit-protocol/constants';
import { LitActionResource } from '@lit-protocol/auth-helpers';
import { LIT_ABILITY, LIT_RPC } from '@lit-protocol/constants';
import { LitContracts } from '@lit-protocol/contracts-sdk';
import { EthWalletProvider } from '@lit-protocol/lit-auth-client';
import { LitNodeClient } from '@lit-protocol/lit-node-client';
Expand All @@ -23,25 +13,15 @@ import {
Listener,
TimerListener,
} from './listeners';
import { signWithLitActionCode } from './litActions';
import { State, StateParams } from './states';
import { Check, Transition } from './transitions';
import { getChain } from './utils/chain';
import { getBalanceTransitionCheck, getERC20Balance } from './utils/erc20';

import type {
Address,
BalanceTransitionDefinition,
BaseBalanceTransitionDefinition,
BaseStateMachineParams,
ERC20BalanceTransitionDefinition,
EvmContractEventTransitionDefinition,
IntervalTransitionDefinition,
NativeBalanceTransitionDefinition,
OnEvmChainEvent,
StateDefinition,
StateMachineDefinition,
TimerTransitionDefinition,
TransitionDefinition,
TransitionParams,
} from './types';

Expand All @@ -51,11 +31,58 @@ const ethPrivateKey = process.env['ETHEREUM_PRIVATE_KEY'];
if (!ethPrivateKey) {
throw new Error('ethPrivateKey not defined');
}
const yellowstoneSigner = new ethers.Wallet(
const yellowstoneLitSigner = new ethers.Wallet(
ethPrivateKey,
new ethers.providers.JsonRpcProvider(LIT_RPC.CHRONICLE_YELLOWSTONE)
);

// TODO improve and move
interface ExecuteLitAction {
litNodeClient: LitNodeClient;
pkpPublicKey: string;
machineSigner: ethers.Wallet;
ipfsId?: string;
code: string;
jsParams: Record<string, any>;
}

async function executeLitAction({
litNodeClient,
pkpPublicKey,
machineSigner,
ipfsId,
code,
jsParams,
}: ExecuteLitAction) {
const pkpSessionSigs = await litNodeClient.getPkpSessionSigs({
pkpPublicKey,
capabilityAuthSigs: [],
authMethods: [
await EthWalletProvider.authenticate({
signer: machineSigner,
litNodeClient: litNodeClient,
expiration: new Date(Date.now() + 1000 * 60 * 10).toISOString(), // 10 minutes
}),
],
resourceAbilityRequests: [
{
resource: new LitActionResource('*'),
ability: LIT_ABILITY.LitActionExecution,
},
],
});

// Run a LitAction
const executeJsResponse = await litNodeClient.executeJs({
ipfsId,
code,
jsParams,
sessionSigs: pkpSessionSigs,
});

return executeJsResponse;
}

/**
* A StateMachine class that manages states and transitions between them.
*/
Expand Down Expand Up @@ -83,6 +110,7 @@ export class StateMachine {
static fromDefinition(machineConfig: StateMachineDefinition): StateMachine {
const { litNodeClient, litContracts = {} } = machineConfig;

// Create litNodeClient and litContracts instances
const litNodeClientInstance =
'connect' in litNodeClient
? litNodeClient
Expand All @@ -91,7 +119,7 @@ export class StateMachine {
'connect' in litContracts
? litContracts
: new LitContracts({
signer: yellowstoneSigner,
signer: yellowstoneLitSigner,
...litContracts,
});

Expand All @@ -110,17 +138,19 @@ export class StateMachine {
});

machineConfig.states.forEach((state) => {
const { litAction } = state;
const { litAction, transaction } = state;

const stateConfig: StateParams = {
key: state.key,
};

const onEnterFunctions = [] as (() => Promise<void>)[];

if (litAction) {
let pkpPublicKey: string = litAction.pkpPublicKey;

stateConfig.onEnter = async () => {
const yellowstoneSigner = new ethers.Wallet(
onEnterFunctions.push(async () => {
const yellowstoneMachineSigner = new ethers.Wallet(
litAction.pkpOwnerKey,
new ethers.providers.JsonRpcProvider(LIT_RPC.CHRONICLE_YELLOWSTONE)
);
Expand All @@ -134,37 +164,99 @@ export class StateMachine {
console.log(`Minted PKP: ${pkp}`);
}

const pkpSessionSigs = await litNodeClientInstance.getPkpSessionSigs({
const litActionResponse = await executeLitAction({
litNodeClient: litNodeClientInstance,
pkpPublicKey,
capabilityAuthSigs: [],
authMethods: [
await EthWalletProvider.authenticate({
signer: yellowstoneSigner,
litNodeClient: litNodeClientInstance,
expiration: new Date(Date.now() + 1000 * 60 * 10).toISOString(), // 10 minutes
}),
],
resourceAbilityRequests: [
{
resource: new LitActionResource('*'),
ability: LIT_ABILITY.LitActionExecution,
},
],
});

// Run a LitAction
const executeJsResponse = await litNodeClientInstance.executeJs({
sessionSigs: pkpSessionSigs,
machineSigner: yellowstoneMachineSigner,
ipfsId: litAction.ipfsId,
code: litAction.code,
jsParams: litAction.jsParams,
});

// TODO send user this result with a webhook maybe
console.log(`============ executeJsResponse:`, executeJsResponse);
};
// TODO send user this result with a webhook and log
console.log(`============ litActionResponse:`, litActionResponse);
});
}

if (transaction) {
onEnterFunctions.push(async () => {
const yellowstoneMachineSigner = new ethers.Wallet(
transaction.pkpOwnerKey,
new ethers.providers.JsonRpcProvider(LIT_RPC.CHRONICLE_YELLOWSTONE)
);

const chain = getChain(transaction);
const chainProvider = new ethers.providers.JsonRpcProvider(
chain.rpcUrls[0],
chain.chainId
);

const contract = new ethers.Contract(
transaction.contractAddress,
transaction.contractABI,
chainProvider
);

const txData = await contract.populateTransaction[transaction.method](
...(transaction.params || [])
);
const gasLimit = await chainProvider.estimateGas({
to: transaction.contractAddress,
data: txData.data,
from: transaction.pkpEthAddress,
});
const gasPrice = await chainProvider.getGasPrice();
const nonce = await chainProvider.getTransactionCount(
transaction.pkpEthAddress
);

const rawTx = {
chainId: chain.chainId,
data: txData.data,
gasLimit: gasLimit.toHexString(),
gasPrice: gasPrice.toHexString(),
nonce,
to: transaction.contractAddress,
};
const rawTxHash = ethers.utils.keccak256(
ethers.utils.serializeTransaction(rawTx)
);

// Sign with the PKP in a LitAction
const litActionResponse = await executeLitAction({
litNodeClient: litNodeClientInstance,
pkpPublicKey: transaction.pkpPublicKey,
machineSigner: yellowstoneMachineSigner,
code: signWithLitActionCode,
jsParams: {
toSign: ethers.utils.arrayify(rawTxHash),
publicKey: transaction.pkpPublicKey,
sigName: 'signedTransaction',
},
});

const signature = litActionResponse.response as string;
const jsonSignature = JSON.parse(signature);
jsonSignature.r = '0x' + jsonSignature.r.substring(2);
jsonSignature.s = '0x' + jsonSignature.s;
const hexSignature = ethers.utils.joinSignature(jsonSignature);

const signedTx = ethers.utils.serializeTransaction(
rawTx,
hexSignature
);

const receipt = await chainProvider.sendTransaction(signedTx);

// TODO send user this result with a webhook and log
console.log('Transaction Receipt:', receipt);
});
}

stateConfig.onEnter = async () => {
await Promise.all(onEnterFunctions.map((onEnter) => onEnter()));
};

stateMachine.addState(stateConfig);
});

Expand Down Expand Up @@ -430,6 +522,7 @@ export class StateMachine {
this.isRunning && (await this.enterState(stateKey));
} catch (e) {
this.currentState = undefined;
console.error(e);
throw new Error(`Could not enter state ${stateKey}`);
}
}
Expand Down
33 changes: 23 additions & 10 deletions packages/automation/src/lib/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,34 @@ import { BaseTransitionParams } from './transitions';

export type Address = `0x${string}`;

export interface LitActionStateDefinition {
export interface OnEvmChain {
evmChainId: number;
}

export interface UsesPkp {
pkpOwnerKey: string;
pkpPublicKey: string;
pkpEthAddress: Address;
}

export interface LitActionStateDefinition extends UsesPkp {
ipfsId?: string; // TODO separate into another without code
code: string;
jsParams: Record<string, any>;
}

export interface TransactionStateDefinition extends UsesPkp, OnEvmChain {
contractAddress: Address;
contractABI: ethers.ContractInterface;
method: string;
value?: string;
params?: any[];
}

export interface StateDefinition {
key: string;
litAction?: LitActionStateDefinition;
}

export interface OnEvmChainEvent {
evmChainId: number;
transaction?: TransactionStateDefinition;
}

export interface IntervalTransitionDefinition {
Expand All @@ -30,7 +43,7 @@ export interface IntervalTransitionDefinition {

export interface BaseBalanceTransitionDefinition
extends IntervalTransitionDefinition,
OnEvmChainEvent {
OnEvmChain {
address: Address;
comparator: '>' | '>=' | '=' | '!=' | '<=' | '<';
amount: string;
Expand All @@ -44,7 +57,7 @@ export interface NativeBalanceTransitionDefinition
export interface ERC20BalanceTransitionDefinition
extends BaseBalanceTransitionDefinition {
type: 'ERC20';
tokenAddress: string;
tokenAddress: Address;
tokenDecimals: number;
}

Expand All @@ -60,9 +73,9 @@ export interface TimerTransitionDefinition
until: number;
}

export interface EvmContractEventTransitionDefinition extends OnEvmChainEvent {
contractAddress: string;
abi: ethers.ContractInterface;
export interface EvmContractEventTransitionDefinition extends OnEvmChain {
contractAddress: Address;
abi: ethers.ContractInterface; // TODO rename a contractABI
eventName: string;
eventParams?: any;
}
Expand Down
4 changes: 2 additions & 2 deletions packages/automation/src/lib/utils/chain.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import { LIT_EVM_CHAINS } from '@lit-protocol/constants';

import { OnEvmChainEvent } from '../types';
import { OnEvmChain } from '../types';

export function getChain(event: OnEvmChainEvent) {
export function getChain(event: OnEvmChain) {
const chain = Object.values(LIT_EVM_CHAINS).find(
(chain) => chain.chainId === event.evmChainId
);
Expand Down

0 comments on commit 47cc68a

Please sign in to comment.