Skip to content

Commit

Permalink
feat: private key store and getters/setters (#5)
Browse files Browse the repository at this point in the history 
* Add private key store and getters/setters

* cargo fmt

* feedback

* release key after submitting

* force key release and fix order udpating
  • Loading branch information
@zhongeric
zhongeric authored Aug 29, 2024
1 parent bf4c460 commit d8ebd18
Show file tree
Hide file tree
Showing 6 changed files with 350 additions and 66 deletions.
59 changes: 54 additions & 5 deletions src/executors/protect_executor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,28 @@ use anyhow::Result;
use artemis_core::executors::mempool_executor::SubmitTxToMempool;
use artemis_core::types::Executor;
use async_trait::async_trait;
use ethers::{providers::Middleware, types::U256};
use ethers::{
middleware::MiddlewareBuilder,
providers::Middleware,
signers::{LocalWallet, Signer},
types::U256,
};

use crate::strategies::keystore::KeyStore;

/// An executor that sends transactions to the mempool.
pub struct ProtectExecutor<M, N> {
client: Arc<M>,
sender_client: Arc<N>,
key_store: Arc<KeyStore>,
}

impl<M: Middleware, N: Middleware> ProtectExecutor<M, N> {
pub fn new(client: Arc<M>, sender_client: Arc<N>) -> Self {
pub fn new(client: Arc<M>, sender_client: Arc<N>, key_store: Arc<KeyStore>) -> Self {
Self {
client,
sender_client,
key_store,
}
}
}
Expand All @@ -30,11 +39,36 @@ impl<M, N> Executor<SubmitTxToMempool> for ProtectExecutor<M, N>
where
M: Middleware,
M::Error: 'static,
N: Middleware,
N: Middleware + 'static,
N::Error: 'static,
{
/// Send a transaction to the mempool.
async fn execute(&self, mut action: SubmitTxToMempool) -> Result<()> {
// Acquire a key from the key store
let (public_address, private_key) = self
.key_store
.acquire_key()
.await
.expect("Failed to acquire key");
info!("Acquired key: {}", public_address);

let chain_id = u64::from_str_radix(
&action
.tx
.chain_id()
.expect("Chain ID not found on transaction")
.to_string(),
10,
)
.expect("Failed to parse chain ID");

let wallet: LocalWallet = private_key
.as_str()
.parse::<LocalWallet>()
.unwrap()
.with_chain_id(chain_id);
let address = wallet.address();

let gas_usage_result = self
.client
.estimate_gas(&action.tx, None)
Expand Down Expand Up @@ -63,11 +97,26 @@ where
}
action.tx.set_gas_price(bid_gas_price);

let sender_client = self.sender_client.clone();
let nonce_manager = sender_client.nonce_manager(address);
let signer = nonce_manager.with_signer(wallet);

info!("Executing tx {:?}", action.tx);
self.sender_client
let result = signer
.send_transaction(action.tx, None)
.await
.map_err(|err| anyhow::anyhow!("Error sending transaction: {}", err))?;
.map_err(|err| anyhow::anyhow!("Error sending transaction: {}", err));

match self.key_store.release_key(public_address.clone()).await {
Ok(_) => {
info!("Released key: {}", public_address);
}
Err(e) => {
info!("Failed to release key: {}", e);
}
}
// Send result up the stack
result?;
Ok(())
}
}
62 changes: 55 additions & 7 deletions src/executors/public_1559_executor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,35 +4,68 @@ use tracing::info;
use anyhow::{Context, Result};
use artemis_core::types::Executor;
use async_trait::async_trait;
use ethers::{providers::Middleware, types::U256};
use ethers::{
middleware::MiddlewareBuilder,
providers::Middleware,
signers::{LocalWallet, Signer},
types::U256,
};

use crate::strategies::types::SubmitTxToMempoolWithExecutionMetadata;
use crate::strategies::{keystore::KeyStore, types::SubmitTxToMempoolWithExecutionMetadata};

/// An executor that sends transactions to the public mempool.
pub struct Public1559Executor<M, N> {
client: Arc<M>,
sender_client: Arc<N>,
key_store: Arc<KeyStore>,
}

impl<M: Middleware, N: Middleware> Public1559Executor<M, N> {
pub fn new(client: Arc<M>, sender_client: Arc<N>) -> Self {
pub fn new(client: Arc<M>, sender_client: Arc<N>, key_store: Arc<KeyStore>) -> Self {
Self {
client,
sender_client,
key_store,
}
}
}

#[async_trait]
impl<M, N> Executor<SubmitTxToMempoolWithExecutionMetadata> for Public1559Executor<M, N>
where
M: Middleware,
M: Middleware + 'static,
M::Error: 'static,
N: Middleware,
N: Middleware + 'static,
N::Error: 'static,
{
/// Send a transaction to the mempool.
async fn execute(&self, mut action: SubmitTxToMempoolWithExecutionMetadata) -> Result<()> {
// Acquire a key from the key store
let (public_address, private_key) = self
.key_store
.acquire_key()
.await
.expect("Failed to acquire key");
info!("Acquired key: {}", public_address);

let chain_id = u64::from_str_radix(
&action
.execution
.tx
.chain_id()
.expect("Chain ID not found on transaction")
.to_string(),
10,
)
.expect("Failed to parse chain ID");

let wallet: LocalWallet = private_key
.as_str()
.parse::<LocalWallet>()
.unwrap()
.with_chain_id(chain_id);
let address = wallet.address();

let gas_usage_result = self
.client
.estimate_gas(&action.execution.tx, None)
Expand Down Expand Up @@ -70,11 +103,26 @@ where

action.execution.tx.set_gas(gas_usage_result);

let sender_client = self.sender_client.clone();
let nonce_manager = sender_client.nonce_manager(address);
let signer = nonce_manager.with_signer(wallet);

info!("Executing tx {:?}", action.execution.tx);
self.sender_client
let result = signer
.send_transaction(action.execution.tx, None)
.await
.map_err(|err| anyhow::anyhow!("Error sending transaction: {}", err))?;
.map_err(|err| anyhow::anyhow!("Error sending transaction: {}", err));

match self.key_store.release_key(public_address.clone()).await {
Ok(_) => {
info!("Released key: {}", public_address);
}
Err(e) => {
info!("Failed to release key: {}", e);
}
}
// Send result up the stack
result?;
Ok(())
}
}
70 changes: 43 additions & 27 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,14 @@ use collectors::{
uniswapx_route_collector::UniswapXRouteCollector,
};
use ethers::{
prelude::MiddlewareBuilder,
providers::{Http, Provider, Ws},
signers::{LocalWallet, Signer},
};
use executors::protect_executor::ProtectExecutor;
use executors::public_1559_executor::Public1559Executor;
use std::collections::HashMap;
use std::sync::Arc;
use strategies::keystore::KeyStore;
use strategies::priority_strategy::UniswapXPriorityFill;
use strategies::{
types::{Action, Config, Event},
Expand All @@ -39,22 +39,26 @@ const MEV_BLOCKER: &str = "https://rpc.mevblocker.io/noreverts";
.required(true)
.args(&["private_key", "aws_secret_arn"])
))]
#[command(group(
ArgGroup::new("key_source")
.args(&["private_key", "private_key_file", "aws_secret_arn"])
))]
pub struct Args {
/// Ethereum node WS endpoint.
#[arg(long)]
pub wss: String,

/// Private key for sending txs.
#[arg(long)]
#[arg(long, group = "key_source")]
pub private_key: Option<String>,
/// public key for the bot that corresponds to the private key.
#[arg(long)]
pub bot_address: String,

/// Path to file containing mapping between public address and private key.
#[arg(long, group = "key_source")]
pub private_key_file: Option<String>,

/// AWS secret arn for fetching private key.
/// This is a secret manager arn that contains the private key as plain text.
#[arg(long)]
#[arg(long, group = "key_source")]
pub aws_secret_arn: Option<String>,

/// Percentage of profit to pay in gas.
Expand Down Expand Up @@ -96,11 +100,13 @@ async fn main() -> Result<()> {
let mevblocker_provider =
Provider::<Http>::try_from(MEV_BLOCKER).expect("could not instantiate HTTP Provider");

/// TODO: support an array of addresses
let pk = if let Some(aws_secret_arn) = args.aws_secret_arn {
let key_store = Arc::new(KeyStore::new());

if let Some(aws_secret_arn) = args.aws_secret_arn {
let config = aws_config::load_from_env().await;
let client = aws_sdk_secretsmanager::Client::new(&config);
let pk_mapping_json = client.get_secret_value()
let pk_mapping_json = client
.get_secret_value()
.secret_id(aws_secret_arn)
.send()
.await
Expand All @@ -109,23 +115,28 @@ async fn main() -> Result<()> {
.expect("secret string not found");
let pk_mapping = serde_json::from_str::<HashMap<String, String>>(&pk_mapping_json)
.expect("could not parse private key mapping");
pk_mapping.get(&args.bot_address).unwrap().clone()
// load into keystore
for (address, pk) in pk_mapping {
key_store.add_key(address, pk).await;
}
} else if let Some(pk_file) = args.private_key_file {
let pk_mapping_json = std::fs::read_to_string(pk_file).expect("could not read pk file");
let pk_mapping = serde_json::from_str::<HashMap<String, String>>(&pk_mapping_json)
.expect("could not parse private key mapping");
// load into keystore
for (address, pk) in pk_mapping {
key_store.add_key(address, pk).await;
}
} else {
args.private_key.clone().unwrap()
};

let wallet: LocalWallet = pk
.parse::<LocalWallet>()
.unwrap()
.with_chain_id(chain_id);
let address = wallet.address();
let pk = args.private_key.clone().unwrap();
let wallet: LocalWallet = pk.parse::<LocalWallet>().unwrap().with_chain_id(chain_id);
let address = wallet.address();
key_store.add_key(address.to_string(), pk).await;
}
info!("Key store initialized with {} keys", key_store.len().await);

let provider = Arc::new(provider.nonce_manager(address).with_signer(wallet.clone()));
let mevblocker_provider = Arc::new(
mevblocker_provider
.nonce_manager(address)
.with_signer(wallet),
);
let provider = Arc::new(provider);
let mevblocker_provider = Arc::new(mevblocker_provider);

// Set up engine.
let mut engine = Engine::default();
Expand Down Expand Up @@ -188,9 +199,14 @@ async fn main() -> Result<()> {
let protect_executor = Box::new(ProtectExecutor::new(
provider.clone(),
mevblocker_provider.clone(),
key_store.clone(),
));

let public_tx_executor = Box::new(Public1559Executor::new(provider.clone(), provider.clone()));
let public_tx_executor = Box::new(Public1559Executor::new(
provider.clone(),
provider.clone(),
key_store.clone(), // TODO: this should be the same as the protect executor
));

let protect_executor = ExecutorMap::new(protect_executor, |action| match action {
Action::SubmitTx(tx) => Some(tx),
Expand Down
Loading

0 comments on commit d8ebd18

Please sign in to comment.