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