lit-rust-sdk 1.0.0

Rust SDK for the Lit Protocol Naga network
Documentation
use crate::auth::{pkp_eth_address_from_pubkey, AuthContext};
use crate::client::LitClient;
use crate::error::LitSdkError;
use async_trait::async_trait;
use ethers::signers::{to_eip155_v, Signer};
use ethers::types::{
    transaction::{eip2718::TypedTransaction, eip712::Eip712},
    Address, Signature, U256,
};
use ethers::utils::hash_message;
use std::fmt;
use tokio::time::{sleep, Duration};

#[derive(Clone)]
pub struct PkpSigner {
    client: LitClient,
    pkp_public_key: String,
    address: Address,
    auth_context: AuthContext,
    chain_id: u64,
    user_max_price_wei: Option<U256>,
}

impl PkpSigner {
    pub fn new(
        client: LitClient,
        pkp_public_key: impl Into<String>,
        auth_context: AuthContext,
        chain_id: u64,
    ) -> Result<Self, LitSdkError> {
        let pkp_public_key = pkp_public_key.into();
        let checksum_address = pkp_eth_address_from_pubkey(&pkp_public_key)?;
        let address: Address = checksum_address
            .parse()
            .map_err(|e| LitSdkError::Config(format!("invalid PKP eth address: {e}")))?;

        Ok(Self {
            client,
            pkp_public_key,
            address,
            auth_context,
            chain_id,
            user_max_price_wei: None,
        })
    }

    pub fn with_user_max_price_wei(mut self, user_max_price_wei: U256) -> Self {
        self.user_max_price_wei = Some(user_max_price_wei);
        self
    }

    async fn sign_digest(&self, digest: [u8; 32]) -> Result<(U256, U256, u8), LitSdkError> {
        let mut last_err: Option<LitSdkError> = None;
        let mut res: Option<serde_json::Value> = None;
        for attempt in 0..3 {
            match self
                .client
                .pkp_sign_ethereum_with_options(
                    &self.pkp_public_key,
                    &digest,
                    &self.auth_context,
                    self.user_max_price_wei,
                    true,
                )
                .await
            {
                Ok(v) => {
                    res = Some(v);
                    break;
                }
                Err(err) => {
                    let msg = err.to_string();
                    let retryable = msg.contains("Rate Limit Exceeded")
                        || msg.contains("Pubkey share not found")
                        || msg.contains("unable to get signature share");
                    if retryable && attempt < 2 {
                        last_err = Some(err);
                        sleep(Duration::from_secs(10 * (attempt as u64 + 1))).await;
                        continue;
                    }
                    return Err(err);
                }
            }
        }
        let res = res.ok_or_else(|| {
            last_err.unwrap_or_else(|| LitSdkError::Network("pkpSign failed after retries".into()))
        })?;

        let sig_str = res
            .get("signature")
            .and_then(|v| v.as_str())
            .ok_or_else(|| {
                LitSdkError::Crypto("pkpSign response missing signature field".into())
            })?;

        let sig_str = sig_str.replace('"', "");
        let sig_hex = sig_str.trim_start_matches("0x");
        let sig = hex::decode(sig_hex).map_err(|e| LitSdkError::Crypto(e.to_string()))?;

        if sig.len() != 64 {
            return Err(LitSdkError::Crypto(format!(
                "expected 64-byte signature, got {} bytes",
                sig.len()
            )));
        }

        let r = U256::from_big_endian(&sig[..32]);
        let s = U256::from_big_endian(&sig[32..]);

        let mut recid = res
            .get("recovery_id")
            .or_else(|| res.get("recoveryId"))
            .and_then(|v| v.as_u64())
            .map(|v| v as u8);

        if recid.is_none() {
            for candidate in [0u8, 1u8] {
                let sig = Signature {
                    r,
                    s,
                    v: 27u64 + candidate as u64,
                };
                if sig.recover(digest).ok() == Some(self.address) {
                    recid = Some(candidate);
                    break;
                }
            }
        }

        let recid = recid
            .ok_or_else(|| LitSdkError::Crypto("failed to recover a valid recovery id".into()))?;

        Ok((r, s, recid))
    }
}

impl fmt::Debug for PkpSigner {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PkpSigner")
            .field("pkp_public_key", &self.pkp_public_key)
            .field("address", &self.address)
            .field("chain_id", &self.chain_id)
            .finish()
    }
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl Signer for PkpSigner {
    type Error = LitSdkError;

    async fn sign_message<S: Send + Sync + AsRef<[u8]>>(
        &self,
        message: S,
    ) -> Result<Signature, Self::Error> {
        let digest = hash_message(message.as_ref()).0;
        let (r, s, recid) = self.sign_digest(digest).await?;
        Ok(Signature {
            r,
            s,
            v: 27u64 + recid as u64,
        })
    }

    async fn sign_transaction(&self, tx: &TypedTransaction) -> Result<Signature, Self::Error> {
        let digest = tx.sighash().0;
        let chain_id = tx.chain_id().map(|id| id.as_u64()).unwrap_or(self.chain_id);
        let (r, s, recid) = self.sign_digest(digest).await?;

        Ok(Signature {
            r,
            s,
            v: to_eip155_v(recid, chain_id),
        })
    }

    async fn sign_typed_data<T: Eip712 + Send + Sync>(
        &self,
        payload: &T,
    ) -> Result<Signature, Self::Error> {
        let digest = payload.encode_eip712().map_err(|e| {
            LitSdkError::Crypto(format!("failed to encode EIP-712 typed data: {e}"))
        })?;
        let (r, s, recid) = self.sign_digest(digest).await?;
        Ok(Signature {
            r,
            s,
            v: 27u64 + recid as u64,
        })
    }

    fn address(&self) -> Address {
        self.address
    }

    fn chain_id(&self) -> u64 {
        self.chain_id
    }

    fn with_chain_id<T: Into<u64>>(mut self, chain_id: T) -> Self {
        self.chain_id = chain_id.into();
        self
    }
}