kabu_types_entities/
signers.rs

1use alloy_consensus::SignableTransaction;
2use alloy_network::{TransactionBuilder, TxSigner as AlloyTxSigner, TxSignerSync};
3use alloy_primitives::{Address, B256, Bytes, hex};
4use alloy_rpc_types::TransactionRequest;
5use alloy_signer_local::PrivateKeySigner;
6use eyre::{OptionExt, Result, eyre};
7use indexmap::IndexMap;
8use rand::prelude::IteratorRandom;
9use reth_ethereum_primitives::{EthPrimitives, TransactionSigned};
10use reth_node_types::NodePrimitives;
11use std::fmt;
12use std::fmt::Debug;
13use std::future::Future;
14use std::pin::Pin;
15use std::sync::Arc;
16
17pub trait LoomTxSigner<N: NodePrimitives>: Send + Sync + Debug {
18    fn sign<'a>(&'a self, tx: TransactionRequest) -> Pin<Box<dyn std::future::Future<Output = Result<N::SignedTx>> + Send + 'a>>;
19    fn sign_sync(&self, tx: TransactionRequest) -> Result<N::SignedTx>;
20    fn address(&self) -> Address;
21}
22
23#[derive(Clone)]
24pub struct TxSignerEth {
25    address: Address,
26    wallet: PrivateKeySigner,
27}
28
29impl Default for TxSignerEth {
30    fn default() -> Self {
31        let wallet = PrivateKeySigner::random();
32        Self { address: wallet.address(), wallet }
33    }
34}
35
36impl fmt::Debug for TxSignerEth {
37    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
38        f.debug_struct("TxSigner").field("address", &self.address.to_string()).finish()
39    }
40}
41
42impl LoomTxSigner<EthPrimitives> for TxSignerEth {
43    fn address(&self) -> Address {
44        self.address
45    }
46    fn sign<'a>(&'a self, tx_req: TransactionRequest) -> Pin<Box<dyn Future<Output = Result<TransactionSigned>> + Send + 'a>> {
47        let fut = async move {
48            let mut typed_tx = tx_req
49                .build_typed_tx()
50                .map_err(|e| eyre!("TRANSACTION_TYPE_IS_MISSING"))?
51                .eip1559()
52                .ok_or_eyre("TRANSACTION_IS_NOT_EIP1559")?
53                .clone();
54            let signature = self.wallet.sign_transaction(&mut typed_tx).await?;
55            let signed_tx = typed_tx.clone().into_signed(signature);
56
57            Ok(signed_tx.into())
58        };
59        Box::pin(fut)
60    }
61
62    fn sign_sync(&self, tx_req: TransactionRequest) -> Result<TransactionSigned> {
63        let mut typed_tx = tx_req
64            .build_unsigned()
65            .map_err(|e| eyre!(format!("CANNOT_BUILD_UNSIGNED with error: {}", e)))?
66            .eip1559()
67            .ok_or_eyre("TRANSACTION_IS_NOT_EIP1559")?
68            .clone();
69
70        let signature = self.wallet.sign_transaction_sync(&mut typed_tx)?;
71        let signed_tx = typed_tx.clone().into_signed(signature);
72
73        Ok(signed_tx.into())
74    }
75}
76
77impl TxSignerEth {
78    pub fn new(wallet: PrivateKeySigner) -> TxSignerEth {
79        TxSignerEth { address: wallet.address(), wallet }
80    }
81}
82
83#[derive(Clone, Default)]
84pub struct TxSigners<N: NodePrimitives = EthPrimitives> {
85    signers: IndexMap<Address, Arc<dyn LoomTxSigner<N>>>,
86}
87
88impl TxSigners<EthPrimitives> {
89    pub fn add_privkey(&mut self, priv_key: Bytes) -> TxSignerEth {
90        let wallet = PrivateKeySigner::from_bytes(&B256::from_slice(priv_key.as_ref())).unwrap();
91        self.signers.insert(wallet.address(), Arc::new(TxSignerEth::new(wallet.clone())));
92        TxSignerEth::new(wallet)
93    }
94
95    pub fn add_testkey(&mut self) -> TxSignerEth {
96        self.add_privkey(Bytes::from(hex!("507485ea5bcf6864596cb51b2e727bb2d8ed5e64bb4f3d8c77a734d2fd610c6e")))
97    }
98}
99
100impl<N: NodePrimitives> TxSigners<N> {
101    pub fn new() -> TxSigners<N> {
102        TxSigners { signers: IndexMap::new() }
103    }
104
105    pub fn len(&self) -> usize {
106        self.signers.len()
107    }
108
109    pub fn is_empty(&self) -> bool {
110        self.signers.is_empty()
111    }
112
113    pub fn get_random_signer(&self) -> Option<Arc<dyn LoomTxSigner<N>>> {
114        if self.is_empty() {
115            None
116        } else {
117            let mut rng = rand::rng();
118            self.signers.values().choose(&mut rng).cloned()
119        }
120    }
121    pub fn get_signer_by_index(&self, index: usize) -> Result<Arc<dyn LoomTxSigner<N>>> {
122        match self.signers.get_index(index) {
123            Some((_, s)) => Ok(s.clone()),
124            None => Err(eyre!("SIGNER_NOT_FOUND")),
125        }
126    }
127
128    pub fn get_signer_by_address(&self, address: &Address) -> Result<Arc<dyn LoomTxSigner<N>>> {
129        match self.signers.get(address) {
130            Some(s) => Ok(s.clone()),
131            None => Err(eyre!("SIGNER_NOT_FOUND")),
132        }
133    }
134
135    pub fn get_address_vec(&self) -> Vec<Address> {
136        self.signers.keys().cloned().collect()
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use alloy_network::eip2718::Encodable2718;
144    use alloy_primitives::{TxHash, address};
145    use alloy_rpc_types::TransactionRequest;
146    use eyre::Result;
147
148    #[test]
149    fn test_new_signer() {
150        let wallet = PrivateKeySigner::random();
151        let signer = TxSignerEth::new(wallet.clone());
152        assert_eq!(signer.address(), wallet.address());
153    }
154
155    #[test]
156    fn test_address() {
157        let wallet = PrivateKeySigner::random();
158        let signer = TxSignerEth::new(wallet.clone());
159        assert_eq!(signer.address(), wallet.address());
160    }
161
162    #[tokio::test]
163    async fn test_sign() -> Result<()> {
164        let wallet = PrivateKeySigner::from_bytes(&B256::repeat_byte(1))?;
165        let signer = Box::new(TxSignerEth::new(wallet));
166        let tx_req = TransactionRequest::default()
167            .with_to(Address::ZERO)
168            .with_nonce(1)
169            .with_gas_limit(1)
170            .with_max_fee_per_gas(1)
171            .with_max_priority_fee_per_gas(1);
172        let tx = signer.sign(tx_req).await?;
173        let tx_hash = tx.tx_hash();
174        let tx_rlp = tx.encoded_2718();
175        assert_eq!(tx_hash, &TxHash::from(hex!("a43d09cb299eb6269f5a63fb10ea078c649cbf6a5f159cfd5b6f4be7ad0dfcfd")));
176        assert!(!tx_rlp.is_empty());
177        Ok(())
178    }
179
180    #[test]
181    fn test_sign_sync() -> Result<()> {
182        let wallet = PrivateKeySigner::from_bytes(&B256::repeat_byte(1))?;
183        let signer = TxSignerEth::new(wallet);
184        let tx_req = TransactionRequest::default()
185            .with_to(Address::ZERO)
186            .with_nonce(1)
187            .with_gas_limit(1)
188            .with_max_fee_per_gas(1)
189            .with_max_priority_fee_per_gas(1);
190        let tx = signer.sign_sync(tx_req)?;
191        let tx_hash = tx.tx_hash();
192        let tx_rlp = tx.encoded_2718();
193        assert_eq!(tx_hash, &TxHash::from(hex!("a43d09cb299eb6269f5a63fb10ea078c649cbf6a5f159cfd5b6f4be7ad0dfcfd")));
194        assert!(!tx_rlp.is_empty());
195        Ok(())
196    }
197
198    // TxSigners tests
199
200    #[test]
201    fn test_new_signers() {
202        let signers: TxSigners<EthPrimitives> = TxSigners::new();
203        assert!(signers.is_empty());
204    }
205
206    #[test]
207    fn test_len() {
208        let mut signers = TxSigners::new();
209        assert_eq!(signers.len(), 0);
210        signers.add_testkey();
211        assert_eq!(signers.len(), 1);
212    }
213
214    #[test]
215    fn test_is_empty() {
216        let mut signers = TxSigners::new();
217        assert!(signers.is_empty());
218        signers.add_testkey();
219        assert!(!signers.is_empty());
220    }
221
222    #[test]
223    fn test_add_privkey() {
224        let mut signers = TxSigners::new();
225        let priv_key = Bytes::from(hex!("507485ea5bcf6864596cb51b2e727bb2d8ed5e64bb4f3d8c77a734d2fd610c6e"));
226        let signer = signers.add_privkey(priv_key);
227        assert_eq!(signers.len(), 1);
228        assert_eq!(signer.address(), signers.get_address_vec()[0]);
229        assert_eq!(signer.address(), address!("16Df4b25e4E37A9116eb224799c1e0Fb17fd8d30"));
230    }
231
232    #[test]
233    fn test_add_testkey() {
234        let mut signers = TxSigners::new();
235        let signer = signers.add_testkey();
236        assert_eq!(signers.len(), 1);
237        assert_eq!(signer.address(), signers.get_address_vec()[0]);
238    }
239
240    #[test]
241    fn test_get_random_signer() {
242        let mut signers = TxSigners::new();
243        signers.add_testkey();
244        assert!(signers.get_random_signer().is_some());
245    }
246
247    #[test]
248    fn test_get_signer_by_index() {
249        let mut signers = TxSigners::new();
250        let signer = signers.add_testkey();
251        let address = signer.address();
252        assert!(signers.get_signer_by_index(0).is_ok());
253        // test negative case
254        let unknown_address = Address::random();
255        assert!(signers.get_signer_by_index(1).is_err());
256    }
257
258    #[test]
259    fn test_get_signer_by_address() {
260        let mut signers = TxSigners::new();
261        let signer = signers.add_testkey();
262        let address = signer.address();
263        assert!(signers.get_signer_by_address(&address).is_ok());
264        // test negative case
265        let unknown_address = Address::random();
266        assert!(signers.get_signer_by_address(&unknown_address).is_err());
267    }
268
269    #[test]
270    fn test_get_address_vec() {
271        let mut signers = TxSigners::new();
272        assert!(signers.get_address_vec().is_empty());
273        let signer = signers.add_testkey();
274        let addresses = signers.get_address_vec();
275        assert_eq!(addresses.len(), 1);
276        assert_eq!(addresses[0], signer.address());
277    }
278}
kabu_types_entities/signers.rs

kabu_types_entities/
signers.rs
