From 7a8872cc415cf99f0122dcc92e87b6e09932d465 Mon Sep 17 00:00:00 2001 From: Diana Carvalho Date: Thu, 30 Jan 2025 13:17:23 +0000 Subject: [PATCH] fix: Don't leak evm specific code to interfaces(PrivateKeySigner, Chain) Pass the signer's private key as a string around and only create a PrivateKeySigner inside Permit2 (where it's needed) Use tycho-core Chain instead of Alloy ChainID. Create a temp method to convert between them Create a EVMSwapEncoder that implements the executor_selector (that is evm specific) --- don't change below this line --- ENG-4081 Took 1 hour 7 minutes --- src/encoding/evm/approvals/permit2.rs | 35 +++++++++++-------- src/encoding/evm/router_encoder.rs | 16 ++++----- src/encoding/evm/strategy_encoder/encoder.rs | 26 +++++++------- src/encoding/evm/strategy_encoder/selector.rs | 11 +++--- src/encoding/evm/swap_encoder/encoders.rs | 14 +++++++- src/encoding/evm/utils.rs | 9 ++++- src/encoding/strategy_encoder.rs | 8 ++--- src/encoding/swap_encoder.rs | 8 ----- 8 files changed, 69 insertions(+), 58 deletions(-) diff --git a/src/encoding/evm/approvals/permit2.rs b/src/encoding/evm/approvals/permit2.rs index c9b687c..900d4cb 100644 --- a/src/encoding/evm/approvals/permit2.rs +++ b/src/encoding/evm/approvals/permit2.rs @@ -9,17 +9,18 @@ use alloy::{ }; #[allow(deprecated)] use alloy_primitives::Signature; +use alloy_primitives::B256; use alloy_sol_types::{eip712_domain, sol, SolStruct, SolValue}; use chrono::Utc; use num_bigint::BigUint; use tokio::runtime::Runtime; -use tycho_core::Bytes; +use tycho_core::{models::Chain, Bytes}; use crate::encoding::{ errors::EncodingError, evm::{ approvals::protocol_approvals_manager::get_client, - utils::{biguint_to_u256, bytes_to_address, encode_input}, + utils::{biguint_to_u256, bytes_to_address, encode_input, to_chain_id}, }, }; @@ -60,10 +61,17 @@ sol! { #[allow(dead_code)] impl Permit2 { - pub fn new(signer: PrivateKeySigner, chain_id: ChainId) -> Result { + pub fn new(signer_pk: String, chain: Chain) -> Result { + let chain_id = to_chain_id(chain)?; let runtime = Runtime::new() .map_err(|_| EncodingError::FatalError("Failed to create runtime".to_string()))?; let client = runtime.block_on(get_client())?; + let pk = B256::from_str(&signer_pk).map_err(|_| { + EncodingError::FatalError("Failed to convert signer private key to B256".to_string()) + })?; + let signer = PrivateKeySigner::from_bytes(&pk).map_err(|_| { + EncodingError::FatalError("Failed to create signer from private key".to_string()) + })?; Ok(Self { address: Address::from_str("0x000000000022D473030F116dDEE9F6B43aC78BA3") .map_err(|_| EncodingError::FatalError("Permit2 address not valid".to_string()))?, @@ -158,7 +166,7 @@ impl Permit2 { mod tests { use std::str::FromStr; - use alloy_primitives::{Uint, B256}; + use alloy_primitives::Uint; use num_bigint::BigUint; use super::*; @@ -196,8 +204,9 @@ mod tests { #[test] fn test_get_existing_allowance() { - let signer = PrivateKeySigner::random(); - let manager = Permit2::new(signer, 1).unwrap(); + let signer_pk = + "4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033".to_string(); + let manager = Permit2::new(signer_pk, Chain::Ethereum).unwrap(); let token = Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(); let owner = Bytes::from_str("0x2c6a3cd97c6283b95ac8c5a4459ebb0d5fd404f4").unwrap(); @@ -216,10 +225,8 @@ mod tests { fn test_get_permit() { // Set up a mock private key for signing let private_key = - B256::from_str("4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033") - .expect("Invalid private key"); - let signer = PrivateKeySigner::from_bytes(&private_key).expect("Failed to create signer"); - let permit2 = Permit2::new(signer, 1).expect("Failed to create Permit2"); + "4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033".to_string(); + let permit2 = Permit2::new(private_key, Chain::Ethereum).expect("Failed to create Permit2"); let owner = Bytes::from_str("0x2c6a3cd97c6283b95ac8c5a4459ebb0d5fd404f4").unwrap(); let spender = Bytes::from_str("0xba12222222228d8ba445958a75a0704d566bf2c8").unwrap(); @@ -257,12 +264,10 @@ mod tests { fn test_permit() { let anvil_account = Bytes::from_str("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266").unwrap(); let anvil_private_key = - B256::from_str("0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80") - .unwrap(); + "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80".to_string(); - let signer = - PrivateKeySigner::from_bytes(&anvil_private_key).expect("Failed to create signer"); - let permit2 = Permit2::new(signer, 1).expect("Failed to create Permit2"); + let permit2 = + Permit2::new(anvil_private_key, Chain::Ethereum).expect("Failed to create Permit2"); let token = Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(); let amount = BigUint::from(1000u64); diff --git a/src/encoding/evm/router_encoder.rs b/src/encoding/evm/router_encoder.rs index a3c74d7..e1f2f08 100644 --- a/src/encoding/evm/router_encoder.rs +++ b/src/encoding/evm/router_encoder.rs @@ -1,9 +1,7 @@ use std::str::FromStr; -use alloy::signers::local::PrivateKeySigner; -use alloy_primitives::ChainId; use num_bigint::BigUint; -use tycho_core::Bytes; +use tycho_core::{models::Chain, Bytes}; use crate::encoding::{ errors::EncodingError, @@ -16,8 +14,8 @@ use crate::encoding::{ #[allow(dead_code)] pub struct EVMRouterEncoder { strategy_selector: S, - signer: Option, - chain_id: ChainId, + signer: Option, + chain: Chain, router_address: String, } @@ -26,10 +24,10 @@ impl EVMRouterEncoder { pub fn new( strategy_selector: S, router_address: String, - signer: Option, - chain_id: ChainId, + signer: Option, + chain: Chain, ) -> Result { - Ok(EVMRouterEncoder { strategy_selector, signer, chain_id, router_address }) + Ok(EVMRouterEncoder { strategy_selector, signer, chain, router_address }) } } impl RouterEncoder for EVMRouterEncoder { @@ -50,7 +48,7 @@ impl RouterEncoder for EVMRouterEncoder { let strategy = self.strategy_selector.select_strategy( solution, self.signer.clone(), - self.chain_id, + self.chain, )?; let method_calldata = strategy.encode_strategy(solution.clone(), router_address)?; diff --git a/src/encoding/evm/strategy_encoder/encoder.rs b/src/encoding/evm/strategy_encoder/encoder.rs index 4a2d17a..dd5ba6f 100644 --- a/src/encoding/evm/strategy_encoder/encoder.rs +++ b/src/encoding/evm/strategy_encoder/encoder.rs @@ -1,10 +1,9 @@ use std::cmp::min; -use alloy::signers::local::PrivateKeySigner; -use alloy_primitives::{aliases::U24, map::HashSet, ChainId, U256, U8}; +use alloy_primitives::{aliases::U24, map::HashSet, U256, U8}; use alloy_sol_types::SolValue; use num_bigint::BigUint; -use tycho_core::Bytes; +use tycho_core::{models::Chain, Bytes}; use crate::encoding::{ errors::EncodingError, @@ -40,8 +39,8 @@ pub struct SplitSwapStrategyEncoder { } impl SplitSwapStrategyEncoder { - pub fn new(signer: PrivateKeySigner, chain_id: ChainId) -> Result { - Ok(Self { permit2: Permit2::new(signer, chain_id)? }) + pub fn new(signer_pk: String, chain: Chain) -> Result { + Ok(Self { permit2: Permit2::new(signer_pk, chain)? }) } } impl EVMStrategyEncoder for SplitSwapStrategyEncoder {} @@ -72,7 +71,7 @@ impl StrategyEncoder for SplitSwapStrategyEncoder { BigUint::ZERO }; - let tokens: Vec = solution + let mut tokens: Vec = solution .swaps .iter() .flat_map(|swap| vec![swap.token_in.clone(), swap.token_out.clone()]) @@ -80,6 +79,10 @@ impl StrategyEncoder for SplitSwapStrategyEncoder { .into_iter() .collect(); + // this is only to make the test deterministic (same index for the same token for different + // runs) + tokens.sort(); + let mut swaps = vec![]; for swap in solution.swaps.iter() { let registry = SWAP_ENCODER_REGISTRY.read().unwrap(); @@ -200,7 +203,6 @@ mod tests { use std::str::FromStr; use alloy::hex::encode; - use alloy_primitives::B256; use tycho_core::dto::ProtocolComponent; use super::*; @@ -210,9 +212,7 @@ mod tests { fn test_split_swap_strategy_encoder() { // Set up a mock private key for signing let private_key = - B256::from_str("4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033") - .expect("Invalid private key"); - let signer = PrivateKeySigner::from_bytes(&private_key).expect("Failed to create signer"); + "4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033".to_string(); let weth = Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(); let dai = Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap(); @@ -228,7 +228,7 @@ mod tests { split: 0f64, }; - let encoder = SplitSwapStrategyEncoder::new(signer, 1).unwrap(); + let encoder = SplitSwapStrategyEncoder::new(private_key, Chain::Ethereum).unwrap(); let solution = Solution { exact_out: false, given_token: weth, @@ -284,8 +284,8 @@ mod tests { "0000000000000000000000000000000000000000000000000000000000000000", "000000000000000000000000000000000000000000000000000000000005b", // Swap header - "00", // token in index - "01", // token out index + "01", // token in index + "00", // token out index "000000", // split // Swap data "5c2f5a71f67c01775180adc06909288b4c329308", // executor address diff --git a/src/encoding/evm/strategy_encoder/selector.rs b/src/encoding/evm/strategy_encoder/selector.rs index 7301552..b8e675c 100644 --- a/src/encoding/evm/strategy_encoder/selector.rs +++ b/src/encoding/evm/strategy_encoder/selector.rs @@ -1,5 +1,4 @@ -use alloy::signers::local::PrivateKeySigner; -use alloy_primitives::ChainId; +use tycho_core::models::Chain; use crate::encoding::{ errors::EncodingError, @@ -14,18 +13,18 @@ impl StrategySelector for EVMStrategySelector { fn select_strategy( &self, solution: &Solution, - signer: Option, - chain_id: ChainId, + signer: Option, + chain: Chain, ) -> Result, EncodingError> { if solution.straight_to_pool { Ok(Box::new(StraightToPoolStrategyEncoder {})) } else { - let signer = signer.ok_or_else(|| { + let signer_pk = signer.ok_or_else(|| { EncodingError::FatalError( "Signer is required for SplitSwapStrategyEncoder".to_string(), ) })?; - Ok(Box::new(SplitSwapStrategyEncoder::new(signer, chain_id).unwrap())) + Ok(Box::new(SplitSwapStrategyEncoder::new(signer_pk, chain).unwrap())) } } } diff --git a/src/encoding/evm/swap_encoder/encoders.rs b/src/encoding/evm/swap_encoder/encoders.rs index b0ae977..e950f56 100644 --- a/src/encoding/evm/swap_encoder/encoders.rs +++ b/src/encoding/evm/swap_encoder/encoders.rs @@ -1,7 +1,8 @@ use std::str::FromStr; -use alloy_primitives::{Address, Bytes as AlloyBytes}; +use alloy_primitives::{Address, Bytes as AlloyBytes, FixedBytes}; use alloy_sol_types::SolValue; +use tycho_core::keccak256; use crate::encoding::{ errors::EncodingError, @@ -12,10 +13,18 @@ use crate::encoding::{ swap_encoder::SwapEncoder, }; +pub trait EVMSwapEncoder: SwapEncoder { + fn executor_selector(&self) -> FixedBytes<4> { + let hash = keccak256("swap(uint256,bytes)".as_bytes()); + FixedBytes::<4>::from([hash[0], hash[1], hash[2], hash[3]]) + } +} + pub struct UniswapV2SwapEncoder { executor_address: String, } +impl EVMSwapEncoder for UniswapV2SwapEncoder {} impl UniswapV2SwapEncoder { fn get_zero_to_one(sell_token_address: Address, buy_token_address: Address) -> bool { sell_token_address < buy_token_address @@ -63,6 +72,7 @@ impl SwapEncoder for UniswapV2SwapEncoder { pub struct UniswapV3SwapEncoder { executor_address: String, } +impl EVMSwapEncoder for UniswapV3SwapEncoder {} impl UniswapV3SwapEncoder { fn get_zero_to_one(sell_token_address: Address, buy_token_address: Address) -> bool { @@ -135,6 +145,8 @@ pub struct BalancerV2SwapEncoder { vault_address: String, } +impl EVMSwapEncoder for BalancerV2SwapEncoder {} + impl SwapEncoder for BalancerV2SwapEncoder { fn new(executor_address: String) -> Self { Self { diff --git a/src/encoding/evm/utils.rs b/src/encoding/evm/utils.rs index 0afe763..3925d44 100644 --- a/src/encoding/evm/utils.rs +++ b/src/encoding/evm/utils.rs @@ -1,7 +1,7 @@ use alloy_primitives::{aliases::U24, Address, Keccak256, U256}; use alloy_sol_types::SolValue; use num_bigint::BigUint; -use tycho_core::Bytes; +use tycho_core::{models::Chain, Bytes}; use crate::encoding::errors::EncodingError; @@ -65,3 +65,10 @@ pub fn percentage_to_uint24(percentage: f64) -> U24 { let scaled = (percentage / 100.0) * (MAX_UINT24 as f64); U24::from(scaled.round()) } + +pub fn to_chain_id(chain: Chain) -> Result { + match chain { + Chain::Ethereum => Ok(1), + _ => Err(EncodingError::FatalError("Unsupported chain".to_string())), + } +} diff --git a/src/encoding/strategy_encoder.rs b/src/encoding/strategy_encoder.rs index 1204e6b..8d37fe1 100644 --- a/src/encoding/strategy_encoder.rs +++ b/src/encoding/strategy_encoder.rs @@ -1,6 +1,4 @@ -use alloy::signers::local::PrivateKeySigner; -use alloy_primitives::ChainId; -use tycho_core::Bytes; +use tycho_core::{models::Chain, Bytes}; use crate::encoding::{errors::EncodingError, models::Solution}; @@ -19,7 +17,7 @@ pub trait StrategySelector { fn select_strategy( &self, solution: &Solution, - signer: Option, - chain_id: ChainId, + signer_pk: Option, + chain_id: Chain, ) -> Result, EncodingError>; } diff --git a/src/encoding/swap_encoder.rs b/src/encoding/swap_encoder.rs index 65e94c3..1a1fc0c 100644 --- a/src/encoding/swap_encoder.rs +++ b/src/encoding/swap_encoder.rs @@ -1,6 +1,3 @@ -use alloy_primitives::FixedBytes; -use tycho_core::keccak256; - use crate::encoding::{ errors::EncodingError, models::{EncodingContext, Swap}, @@ -17,9 +14,4 @@ pub trait SwapEncoder: Sync + Send { encoding_context: EncodingContext, ) -> Result, EncodingError>; fn executor_address(&self) -> &str; - - fn executor_selector(&self) -> FixedBytes<4> { - let hash = keccak256("swap(uint256,bytes)".as_bytes()); - FixedBytes::<4>::from([hash[0], hash[1], hash[2], hash[3]]) - } }