Skip to main content

forge_script/
session.rs

1use alloy_primitives::{
2    Address,
3    map::{AddressHashSet, HashMap},
4};
5use eyre::Result;
6use foundry_cli::opts::TempoOpts;
7use foundry_common::tempo::ResolvedSessionSigner;
8use foundry_wallets::{TempoAccessKeyConfig, WalletSigner};
9use itertools::Itertools;
10
11/// A transaction sender scoped to one chain.
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
13pub(crate) struct SignerScope {
14    chain: u64,
15    sender: Address,
16}
17
18impl SignerScope {
19    pub(crate) const fn new(chain: u64, sender: Address) -> Self {
20        Self { chain, sender }
21    }
22}
23
24/// A remaining unsigned script transaction, represented only by the data needed for signer lookup.
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub(crate) struct RemainingScriptTransaction {
27    pub(crate) chain: u64,
28    pub(crate) from: Address,
29}
30
31impl RemainingScriptTransaction {
32    pub(crate) const fn scope(&self) -> SignerScope {
33        SignerScope::new(self.chain, self.from)
34    }
35}
36
37/// Returns the single sender a configured Tempo session is allowed to cover.
38///
39/// Session signing is intentionally fail-closed: a single session access key represents one root
40/// account, so scripts with multiple pending senders must not silently mix the session key with
41/// other wallets.
42pub(crate) fn script_session_expected_sender_if_configured(
43    tempo: &TempoOpts,
44    required_addresses: &AddressHashSet,
45) -> Result<Option<Address>> {
46    tempo.session_id()?.map_or(Ok(None), |_| single_session_sender(required_addresses))
47}
48
49fn single_session_sender(required_addresses: &AddressHashSet) -> Result<Option<Address>> {
50    required_addresses
51        .iter()
52        .copied()
53        .at_most_one()
54        .map_err(|_| eyre::eyre!("Tempo sessions require a single script sender"))
55}
56
57/// Inserts this session access key when it covers the remaining transaction set.
58///
59/// Transactions from the session root on any other chain are rejected up front, so callers do not
60/// accidentally fall back to a long-lived root signer for the same session account.
61pub(crate) fn insert_session_access_key_for_remaining_transactions(
62    access_keys: &mut HashMap<SignerScope, (WalletSigner, TempoAccessKeyConfig)>,
63    session: ResolvedSessionSigner,
64    remaining_transactions: &[RemainingScriptTransaction],
65) -> Result<()> {
66    let chain = session.session.chain_id;
67    let root = session.session.root_account;
68    if let Some(tx) = remaining_transactions.iter().find(|tx| tx.from == root && tx.chain != chain)
69    {
70        eyre::bail!(
71            "Tempo session is for chain {}, but a remaining transaction from session root {} is on chain {}",
72            chain,
73            root,
74            tx.chain,
75        );
76    }
77
78    if remaining_transactions.iter().any(|tx| tx.from == root) {
79        access_keys.insert(SignerScope::new(chain, root), (session.signer, session.access_key));
80    }
81
82    Ok(())
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use alloy_primitives::{B256, address};
89    use alloy_signer::Signer;
90    use foundry_common::tempo::{KeyType, SessionEntry, SessionKeyMaterial, SessionStatus};
91
92    const ROOT_PRIVATE_KEY: &str =
93        "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80";
94    const ACCESS_KEY_PRIVATE_KEY: &str =
95        "0x59c6995e998f97a5a004497e5da3b5d2b2b66a87f064d39c44da0b6d6e4f8ff0";
96
97    #[test]
98    fn session_sender_requires_single_root_account() {
99        let one = address!("0x1111111111111111111111111111111111111111");
100        let two = address!("0x2222222222222222222222222222222222222222");
101        let single_sender = [one].into_iter().collect();
102        let multiple_senders = [one, two].into_iter().collect();
103
104        assert_eq!(single_session_sender(&single_sender).unwrap(), Some(one));
105        assert!(single_session_sender(&multiple_senders).is_err());
106    }
107
108    #[test]
109    fn session_access_key_rejects_session_root_on_wrong_chain() {
110        let (session, root_address, _) = session_signer(4217);
111        let remaining = [RemainingScriptTransaction { chain: 1, from: root_address }];
112        let mut access_keys = HashMap::default();
113
114        let err = insert_session_access_key_for_remaining_transactions(
115            &mut access_keys,
116            session,
117            &remaining,
118        )
119        .unwrap_err();
120
121        assert!(access_keys.is_empty());
122        let message = err.to_string();
123        assert!(message.contains("Tempo session is for chain 4217"), "{message}");
124        assert!(message.contains("transaction from session root"), "{message}");
125        assert!(message.contains("chain 1"), "{message}");
126    }
127
128    #[test]
129    fn session_access_key_is_inserted_for_session_chain() {
130        let (session, root_address, access_key_address) = session_signer(4217);
131        let remaining = [RemainingScriptTransaction { chain: 4217, from: root_address }];
132        let mut access_keys = HashMap::default();
133
134        insert_session_access_key_for_remaining_transactions(&mut access_keys, session, &remaining)
135            .unwrap();
136
137        let (signer, config) =
138            access_keys.get(&SignerScope::new(4217, root_address)).expect("session access key");
139        assert_eq!(signer.address(), access_key_address);
140        assert_eq!(config.wallet_address, root_address);
141        assert_eq!(config.key_address, access_key_address);
142    }
143
144    fn session_signer(chain_id: u64) -> (ResolvedSessionSigner, Address, Address) {
145        let root = foundry_wallets::utils::create_private_key_signer(ROOT_PRIVATE_KEY).unwrap();
146        let root_address = root.address();
147        let signer =
148            foundry_wallets::utils::create_private_key_signer(ACCESS_KEY_PRIVATE_KEY).unwrap();
149        let key_address = signer.address();
150        let access_key = TempoAccessKeyConfig {
151            wallet_address: root_address,
152            key_address,
153            key_authorization: None,
154        };
155        let session = SessionEntry {
156            session_id: B256::ZERO,
157            root_account: root_address,
158            chain_id,
159            key_address,
160            expiry: u64::MAX,
161            scope: None,
162            limits: None,
163            status: SessionStatus::Active,
164            key: Some(SessionKeyMaterial {
165                key_type: KeyType::Secp256k1,
166                key: ACCESS_KEY_PRIVATE_KEY.to_string(),
167                key_authorization: None,
168            }),
169        };
170
171        (ResolvedSessionSigner { session, signer, access_key }, root_address, key_address)
172    }
173}