kabu_types_entities/
signers.rs

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