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#[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#[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
37pub(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
57pub(crate) fn insert_session_access_key_for_remaining_transactions(
62 access_keys: &mut HashMap<SignerScope, TempoAccountsWallet>,
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.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 wallet =
138 access_keys.get(&SignerScope::new(4217, root_address)).expect("session access key");
139 assert_eq!(wallet.account(), root_address);
140 assert_eq!(wallet.key_id().unwrap(), access_key_address);
141 }
142
143 fn session_signer(chain_id: u64) -> (ResolvedSessionSigner, Address, Address) {
144 let root = foundry_wallets::utils::create_private_key_signer(ROOT_PRIVATE_KEY).unwrap();
145 let root_address = root.address();
146 let signer = foundry_wallets::utils::create_local_signer(ACCESS_KEY_PRIVATE_KEY).unwrap();
147 let key_address = signer.address();
148 let access_key =
149 TempoAccountsWallet::from_secp256k1(root_address, signer, None).with_chain_id(chain_id);
150 let session = SessionEntry {
151 session_id: B256::ZERO,
152 root_account: root_address,
153 chain_id,
154 key_address,
155 expiry: u64::MAX,
156 scope: None,
157 limits: None,
158 status: SessionStatus::Active,
159 key: Some(SessionKeyMaterial {
160 key_type: KeyType::Secp256k1,
161 key: ACCESS_KEY_PRIVATE_KEY.to_string(),
162 key_authorization: None,
163 }),
164 };
165
166 (ResolvedSessionSigner { session, access_key }, root_address, key_address)
167 }
168}