Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: private key store and getters/setters #5

Merged
merged 5 commits into from
Aug 29, 2024
Merged

feat: private key store and getters/setters #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b06f62a
Add private key store and getters/setters
zhongeric Aug 29, 2024
4b80527
cargo fmt
zhongeric Aug 29, 2024
3a9b5ca
feedback
zhongeric Aug 29, 2024
a5cc248
release key after submitting
zhongeric Aug 29, 2024
a2b0da4
force key release and fix order udpating
zhongeric Aug 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 54 additions & 5 deletions src/executors/protect_executor.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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