feat: replace stateWalletClient with StateSigner

integration/common/identity.ts

@@ -2,10 +2,11 @@ import { Address, createWalletClient, Hex, http } from 'viem';
 import { privateKeyToAccount } from 'viem/accounts';
 import { chain } from './setup';
 import { createECDSAMessageSigner } from '@erc7824/nitrolite';
+import { SessionKeyStateSigner } from '@erc7824/nitrolite/dist/client/signer';
 
 export class Identity {
     public walletClient = null;
-    public stateWalletClient = null;
+    public stateSigner = null;
     public walletAddress: Address;
     public sessionAddress: Address;
     public messageSigner = null;
@@ -20,21 +21,8 @@ export class Identity {
             transport: http(),
         });
 
-        const sessionAccount = privateKeyToAccount(sessionPrivateKey);
-        this.sessionAddress = sessionAccount.address;
-
-        this.stateWalletClient = {
-            ...this.walletClient,
-            account: {
-                address: this.sessionAddress,
-            },
-            signMessage: async ({ message: { raw } }: { message: { raw: string } }) => {
-                const flatSignature = await sessionAccount.sign({ hash: raw as Hex });
-
-                return flatSignature as Hex;
-            },
-        };
-
+        this.stateSigner = new SessionKeyStateSigner(sessionPrivateKey);
         this.messageSigner = createECDSAMessageSigner(sessionPrivateKey);
+        this.sessionAddress = this.stateSigner.getAddress();
     }
 }

integration/common/nitroliteClient.ts

@@ -21,7 +21,7 @@ export class TestNitroliteClient extends NitroliteClient {
             // @ts-ignore
             publicClient,
             walletClient: identity.walletClient,
-            stateWalletClient: identity.stateWalletClient,
+            stateSigner: identity.stateSigner,
             account: identity.walletClient.account,
             chainId: chain.id,
             challengeDuration: BigInt(CONFIG.DEFAULT_CHALLENGE_TIMEOUT), // min

sdk/src/client/index.ts

@@ -21,6 +21,7 @@ import {
     ResizeChannelParams,
     State,
 } from './types';
+import { StateSigner } from './signer';
 
 const CUSTODY_MIN_CHALLENGE_DURATION = 3600n;
 
@@ -36,7 +37,7 @@ export class NitroliteClient {
     public readonly challengeDuration: bigint;
     public readonly txPreparer: NitroliteTransactionPreparer;
     public readonly chainId: number;
-    private readonly stateWalletClient: WalletClient<Transport, Chain, ParseAccount<Account>>;
+    private readonly stateSigner: StateSigner;
     private readonly nitroliteService: NitroliteService;
     private readonly erc20Service: Erc20Service;
     private readonly sharedDeps: PreparerDependencies;
@@ -57,8 +58,7 @@ export class NitroliteClient {
 
         this.publicClient = config.publicClient;
         this.walletClient = config.walletClient;
-        // Determine which wallet client to use for state signing
-        this.stateWalletClient = config.stateWalletClient ?? config.walletClient;
+        this.stateSigner = config.stateSigner;
         this.account = config.walletClient.account;
         this.addresses = config.addresses;
         this.challengeDuration = config.challengeDuration;
@@ -79,7 +79,7 @@ export class NitroliteClient {
             account: this.account,
             walletClient: this.walletClient,
             challengeDuration: this.challengeDuration,
-            stateWalletClient: this.stateWalletClient,
+            stateSigner: this.stateSigner,
             chainId: this.chainId,
         };
 

sdk/src/client/prepare.ts

@@ -24,6 +24,7 @@ import {
     CreateChannelParams,
     ResizeChannelParams,
 } from './types';
+import { StateSigner } from './signer';
 
 /**
  * Represents the data needed to construct a transaction or UserOperation call.
@@ -41,7 +42,7 @@ export interface PreparerDependencies {
     addresses: ContractAddresses;
     account: ParseAccount<Account>;
     walletClient: WalletClient<Transport, Chain, ParseAccount<Account>>;
-    stateWalletClient: WalletClient<Transport, Chain, ParseAccount<Account>>;
+    stateSigner: StateSigner;
     challengeDuration: bigint;
     chainId: number;
 }

sdk/src/client/signer.ts

@@ -0,0 +1,92 @@
+import { Account, Address, Chain, Hex, ParseAccount, toHex, Transport, WalletClient } from 'viem';
+import { State } from './types';
+import { getPackedState, getStateHash } from '../utils';
+import { signRawECDSAMessage } from '../utils/sign';
+import { privateKeyToAccount } from 'viem/accounts';
+
+// TODO: perhaps extend this interface with rpc signing methods and use it as universal signer interface
+
+/**
+ * Interface for signing protocol states.
+ * This interface is used to abstract the signing logic for state updates in the Nitrolite SDK.
+ * It allows for different implementations, such as using a wallet client or a session key.
+ * Also implementation could include data packing/encoding, which is crucial for some signatures (EIP-712, EIP-191)
+ */
+export interface StateSigner {
+    /**
+     * Get the address of the signer.
+     * @returns The address of the signer.
+     */
+    getAddress(): Address;
+    /**
+     * Sign a state for a given channel ID.
+     * @param channelId The ID of the channel.
+     * @param state The state to sign.
+     * @returns A Promise that resolves to the signature as a Hex string.
+     */
+    signState(channelId: Hex, state: State): Promise<Hex>;
+    /**
+     * Sign a raw message.
+     * @param message The message to sign as a Hex string.
+     * @returns A Promise that resolves to the signature as a Hex string.
+     * @dev use viem's `toHex` to convert the message to Hex if needed.
+     */
+    signRawMessage(message: Hex): Promise<Hex>;
+}
+
+/**
+ * Implementation of the StateSigner interface using a viem WalletClient.
+ * This class uses the wallet client to sign states and raw messages.
+ * It is suitable for use in scenarios where the wallet client is available and can sign messages,
+ * e.g. signing with MetaMask or other wallet providers.
+ */
+export class WalletStateSigner implements StateSigner {
+    private readonly walletClient: WalletClient<Transport, Chain, ParseAccount<Account>>;
+
+    constructor(walletClient: WalletClient<Transport, Chain, ParseAccount<Account>>) {
+        this.walletClient = walletClient;
+    }
+
+    getAddress(): Address {
+        return this.walletClient.account.address;
+    }
+
+    async signState(channelId: Hex, state: State): Promise<Hex> {
+        const packedState = getPackedState(channelId, state)
+
+        return this.walletClient.signMessage({ message: { raw: packedState } });
+    }
+
+    async signRawMessage(message: Hex): Promise<Hex> {
+        return this.walletClient.signMessage({ message: { raw: message } });
+    }
+}
+
+/**
+ * Implementation of the StateSigner interface using a session key.
+ * This class uses a session key to sign states and raw messages.
+ * It is suitable for scenarios where a session key is used for signing and private key could be exposed to application.
+ */
+export class SessionKeyStateSigner implements StateSigner {
+    private readonly sessionKey: Hex;
+    private readonly account: Account;
+
+    constructor(sessionKey: Hex) {
+        this.sessionKey = sessionKey;
+        this.account = privateKeyToAccount(sessionKey);
+    }
+
+    getAddress(): Address {
+        return this.account.address;
+    }
+
+    async signState(channelId: Hex, state: State): Promise<Hex> {
+        const packedState = getPackedState(channelId, state);
+
+        return signRawECDSAMessage(packedState, this.sessionKey);
+    }
+
+    async signRawMessage(message: Hex): Promise<Hex> {
+        return signRawECDSAMessage(message, this.sessionKey);
+    }
+}

sdk/src/client/state.ts

@@ -1,6 +1,6 @@
-import { Address, encodeAbiParameters, encodePacked, Hex, keccak256 } from 'viem';
+import { Address, Hex } from 'viem';
 import * as Errors from '../errors';
-import { generateChannelNonce, getChannelId, getPackedState, getStateHash, signChallengeState, signState } from '../utils';
+import { generateChannelNonce, getChannelId, getPackedChallengeState } from '../utils';
 import { PreparerDependencies } from './prepare';
 import {
     ChallengeChannelParams,
@@ -38,7 +38,7 @@ export async function _prepareAndSignInitialState(
     const channelNonce = generateChannelNonce(deps.account.address);
 
     const participants: [Hex, Hex] = [deps.account.address, deps.addresses.guestAddress];
-    const channelParticipants: [Hex, Hex] = [deps.stateWalletClient.account.address, deps.addresses.guestAddress];
+    const channelParticipants: [Hex, Hex] = [deps.stateSigner.getAddress(), deps.addresses.guestAddress];
     const adjudicatorAddress = deps.addresses.adjudicator;
     if (!adjudicatorAddress) {
         throw new Errors.MissingParameterError(
@@ -77,8 +77,7 @@ export async function _prepareAndSignInitialState(
         sigs: [],
     };
 
-
-    const accountSignature = await signState(channelId, stateToSign, deps.stateWalletClient.signMessage);
+    const accountSignature = await deps.stateSigner.signState(channelId, stateToSign);
     const initialState: State = {
         ...stateToSign,
         sigs: [accountSignature],
@@ -104,7 +103,8 @@ export async function _prepareAndSignChallengeState(
     challengerSig: Signature;
 }> {
     const { channelId, candidateState, proofStates = [] } = params;
-    const challengerSig = await signChallengeState(channelId, candidateState, deps.stateWalletClient.signMessage);
+    const challengeMsg = getPackedChallengeState(channelId, candidateState);
+    const challengerSig = await deps.stateSigner.signRawMessage(challengeMsg);
 
     return { channelId, candidateState, proofs: proofStates, challengerSig };
 }
@@ -136,7 +136,7 @@ export async function _prepareAndSignResizeState(
         sigs: [],
     };
 
-    const accountSignature = await signState(channelId, stateToSign, deps.stateWalletClient.signMessage);
+    const accountSignature = await deps.stateSigner.signState(channelId, stateToSign);
 
     // Create a new state with signatures in the requested style
     const resizeStateWithSigs: State = {
@@ -176,7 +176,7 @@ export async function _prepareAndSignFinalState(
         sigs: [],
     };
 
-    const accountSignature = await signState(channelId, stateToSign, deps.stateWalletClient.signMessage);
+    const accountSignature = await deps.stateSigner.signState(channelId, stateToSign);
 
     // Create a new state with signatures in the requested style
     const finalStateWithSigs: State = {

sdk/src/client/types.ts

@@ -1,5 +1,6 @@
 import { Account, Hex, PublicClient, WalletClient, Chain, Transport, ParseAccount, Address } from 'viem';
 import { ContractAddresses } from '../abis';
+import { StateSigner } from './signer';
 
 /**
  * Channel identifier
@@ -125,14 +126,9 @@ export interface NitroliteClientConfig {
     walletClient: WalletClient<Transport, Chain, ParseAccount<Account>>;
 
     /**
-     * Optional: A separate viem WalletClient used *only* for signing off-chain state updates (`signMessage`).
-     * Provide this if you want to use a different key (e.g., a "hot" key from localStorage)
-     * for state signing than the one used for on-chain transactions.
-     * If omitted, `walletClient` will be used for state signing.
-     * @dev Note that the client's `signMessage` function should NOT add an EIP-191 prefix to the message signed. See {@link SignMessageFn} for details.
-     * viem's `signMessage` can operate in `raw` mode, which suffice.
+     * Implementation of the StateSigner interface used for signing protocol states.
      */
-    stateWalletClient?: WalletClient<Transport, Chain, ParseAccount<Account>>;
+    stateSigner: StateSigner;
 
     /** Contract addresses required by the SDK. */
     addresses: ContractAddresses;

sdk/src/rpc/api.ts

@@ -1,4 +1,4 @@
-import { Address, Hex, keccak256, stringToBytes, WalletClient } from 'viem';
+import { Address, Hex, keccak256, stringToBytes, toHex, WalletClient } from 'viem';
 import { privateKeyToAccount } from 'viem/accounts';
 import {
     MessageSigner,
@@ -28,6 +28,7 @@ import {
     GetLedgerTransactionsRequestParams,
     TransferRequestParams,
 } from './types/request';
+import { signRawECDSAMessage } from '../utils/sign';
 
 /**
  * Creates the signed, stringified message body for an 'auth_request'.
@@ -642,12 +643,9 @@ export function createEIP712AuthMessageSigner(
 export function createECDSAMessageSigner(privateKey: Hex): MessageSigner {
     return async (payload: RequestData | ResponsePayload): Promise<Hex> => {
         try {
-            const messageBytes = keccak256(
-                stringToBytes(JSON.stringify(payload, (_, v) => (typeof v === 'bigint' ? v.toString() : v))),
-            );
-            const flatSignature = await privateKeyToAccount(privateKey).sign({ hash: messageBytes });
+            const message = toHex(JSON.stringify(payload, (_, v) => (typeof v === 'bigint' ? v.toString() : v)));
 
-            return flatSignature as Hex;
+            return signRawECDSAMessage(message, privateKey);
         } catch (error) {
             console.error('ECDSA signing failed:', error);
             throw new Error(`ECDSA signing failed: ${error}`);

sdk/src/utils/sign.ts

@@ -0,0 +1,9 @@
+import { Hex, keccak256 } from 'viem';
+import { privateKeyToAccount } from 'viem/accounts';
+
+export const signRawECDSAMessage = async (message: Hex, privateKey: Hex): Promise<Hex> => {
+    const hash = keccak256(message);
+    const flatSignature = await privateKeyToAccount(privateKey).sign({ hash });
+
+    return flatSignature;
+};

sdk/src/utils/state.ts

@@ -45,65 +45,28 @@ export function getStateHash(channelId: ChannelId, state: State): StateHash {
 }
 
 /**
- * Function type for signing messages, compatible with Viem's WalletClient or Account.
- * @dev Signing should NOT add an EIP-191 prefix to the message.
- * @param args An object containing the message to sign in the `{ message: { raw: Hex } }` format.
- * @returns A promise that resolves to the signature as a Hex string.
- * @throws If the signing fails.
+ * Get a packed challenge state for a channel.
+ * This function encodes the packed state and the challenge string.ß
+ * @param channelId The ID of the channel.
+ * @param state The state to calculate with.
+ * @returns The encoded and packed challenge state as a Hex string.
  */
-type SignMessageFn = (args: { message: { raw: Hex } }) => Promise<Hex>;
+export function getPackedChallengeState(channelId: ChannelId, state: State): Hex {
+    const packedState = getPackedState(channelId, state);
+    const encoded = encodePacked(['bytes', 'string'], [packedState, 'challenge']);
 
-// TODO: extract into an interface and provide on NitroliteClient creation
-/**
- * Create a raw ECDSA signature for a hash over a packed state using a Viem WalletClient or Account compatible signer.
- * Uses the locally defined parseSignature function.
- * @dev `signMessage` function should NOT add an EIP-191 prefix to the stateHash. See {@link SignMessageFn}.
- * @param stateHash The hash of the state to sign.
- * @param signer An object with a `signMessage` method compatible with Viem's interface (e.g., WalletClient, Account).
- * @returns The signature over the state hash.
- * @throws If the signer cannot sign messages or signing/parsing fails.
- */
-export async function signState(
-    channelId: ChannelId,
-    state: State,
-    signMessage: SignMessageFn,
-): Promise<Signature> {
-    const stateHash = getStateHash(channelId, state);
-    try {
-        return await signMessage({ message: { raw: stateHash } });
-    } catch (error) {
-        console.error('Error signing state hash:', error);
-        throw new Error(`Failed to sign state hash: ${error instanceof Error ? error.message : String(error)}`);
-    }
+    return encoded;
 }
 
 /**
- * Signs a challenge state for a channel.
- * This function encodes the packed state and the challenge string, hashes it, and signs it.
+ * Calculate a challenge state for a channel.
+ * This function encodes the packed state and the challenge string and hashes it
  * @param channelId The ID of the channel.
- * @param state The state to sign.
- * @param signMessage The signing function compatible with Viem's WalletClient or Account.
- * @returns The signature as a Hex string.
- * @throws If signing fails.
+ * @param state The state to calculate with.
+ * @returns The challenge hash as a Hex string.
  */
-export async function signChallengeState(
-    channelId: ChannelId,
-    state: State,
-    signMessage: SignMessageFn,
-): Promise<Signature> {
-    const packedState = getPackedState(channelId, state);
-    const encoded = encodePacked(
-        [ 'bytes', 'string' ],
-        [packedState, 'challenge'],
-    );
-    const challengeHash = keccak256(encoded) as Hex;
-
-    try {
-        return await signMessage({ message: { raw: challengeHash } });
-    } catch (error) {
-        console.error('Error signing challenge state:', error);
-        throw new Error(`Failed to sign challenge state: ${error instanceof Error ? error.message : String(error)}`);
-    }
+export function getChallengeHash(channelId: ChannelId, state: State): Hex {
+    return keccak256(getPackedChallengeState(channelId, state));
 }
 
 // TODO: extract into an interface and provide on NitroliteClient creation