Skip to content

Commit

Permalink
feat(builder): allow for multiple private keys
Browse files Browse the repository at this point in the history
  • Loading branch information
dancoombs committed Jun 18, 2024
1 parent 17d37a7 commit 9b60bf4
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 16 deletions.
47 changes: 38 additions & 9 deletions bin/rundler/src/cli/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

use std::{net::SocketAddr, time::Duration};

use anyhow::Context;
use anyhow::{bail, Context};
use clap::Args;
use rundler_builder::{
self, BloxrouteSenderArgs, BuilderEvent, BuilderEventKind, BuilderTask, BuilderTaskArgs,
Expand Down Expand Up @@ -57,13 +57,28 @@ pub struct BuilderArgs {
host: String,

/// Private key to use for signing transactions
/// DEPRECATED: Use `builder.private_keys` instead
///
/// If both `builder.private_key` and `builder.private_keys` are set, `builder.private_key` is appended
/// to `builder.private_keys`. Keys must be unique.
#[arg(
long = "builder.private_key",
name = "builder.private_key",
env = "BUILDER_PRIVATE_KEY"
)]
private_key: Option<String>,

/// Private keys to use for signing transactions
///
/// Cannot use both `builder.private_key` and `builder.aws_kms_key_ids` at the same time.
#[arg(
long = "builder.private_keys",
name = "builder.private_keys",
env = "BUILDER_PRIVATE_KEYS",
value_delimiter = ','
)]
private_keys: Vec<String>,

/// AWS KMS key IDs to use for signing transactions
#[arg(
long = "builder.aws_kms_key_ids",
Expand Down Expand Up @@ -284,17 +299,31 @@ impl BuilderArgs {
num_builders += common.num_builders_v0_7;
}

if self.private_key.is_some() {
if num_builders > 1 {
return Err(anyhow::anyhow!(
"Cannot use a private key with multiple builders. You may need to disable one of the entry points."
));
if (self.private_key.is_some() || !self.private_keys.is_empty())
&& !self.aws_kms_key_ids.is_empty()
{
bail!(
"Cannot use both builder.private_key(s) and builder.aws_kms_key_ids at the same time."
);
}

let mut private_keys = self.private_keys.clone();
if self.private_key.is_some() || !self.private_keys.is_empty() {
if let Some(pk) = &self.private_key {
private_keys.push(pk.clone());
}

if num_builders > private_keys.len() as u64 {
bail!(
"Found {} private keys, but need {} keys for the number of builders. You may need to disable one of the entry points.",
private_keys.len(), num_builders
);
}
} else if self.aws_kms_key_ids.len() < num_builders as usize {
return Err(anyhow::anyhow!(
bail!(
"Not enough AWS KMS key IDs for the number of builders. Need {} keys, found {}. You may need to disable one of the entry points.",
num_builders, self.aws_kms_key_ids.len()
));
);
}

let sender_args = self.sender_args(&chain_spec, &rpc_url)?;
Expand All @@ -304,7 +333,7 @@ impl BuilderArgs {
chain_spec,
unsafe_mode: common.unsafe_mode,
rpc_url,
private_key: self.private_key.clone(),
private_keys,
aws_kms_key_ids: self.aws_kms_key_ids.clone(),
aws_kms_region: common
.aws_region
Expand Down
25 changes: 19 additions & 6 deletions crates/builder/src/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,8 @@ pub struct Args {
/// True if using unsafe mode
pub unsafe_mode: bool,
/// Private key to use for signing transactions
/// If not provided, AWS KMS will be used
pub private_key: Option<String>,
/// If empty, AWS KMS will be used
pub private_keys: Vec<String>,
/// AWS KMS key ids to use for signing transactions
/// Only used if private_key is not provided
pub aws_kms_key_ids: Vec<String>,
Expand Down Expand Up @@ -152,6 +152,7 @@ where

let mut sender_handles = vec![];
let mut bundle_sender_actions = vec![];
let mut pk_iter = self.args.private_keys.clone().into_iter();

for ep in &self.args.entry_points {
match ep.version {
Expand All @@ -162,6 +163,7 @@ where
Arc::clone(&provider),
submit_provider.clone(),
ep_v0_6.clone(),
&mut pk_iter,
)
.await?;
sender_handles.extend(handles);
Expand All @@ -174,6 +176,7 @@ where
Arc::clone(&provider),
submit_provider.clone(),
ep_v0_7.clone(),
&mut pk_iter,
)
.await?;
sender_handles.extend(handles);
Expand Down Expand Up @@ -255,19 +258,21 @@ where
Box::new(self)
}

async fn create_builders_v0_6<C, E>(
async fn create_builders_v0_6<C, E, I>(
&self,
ep: &EntryPointBuilderSettings,
provider: Arc<EthersProvider<C>>,
submit_provider: Option<Arc<EthersProvider<C>>>,
ep_v0_6: E,
pk_iter: &mut I,
) -> anyhow::Result<(
Vec<JoinHandle<anyhow::Result<()>>>,
Vec<mpsc::Sender<BundleSenderAction>>,
)>
where
C: JsonRpcClient + 'static,
E: EntryPointProvider<v0_6::UserOperation> + Clone,
I: Iterator<Item = String>,
{
info!("Mempool config for ep v0.6: {:?}", ep.mempool_configs);
let mut sender_handles = vec![];
Expand All @@ -284,6 +289,7 @@ where
ep_v0_6.clone(),
self.args.sim_settings,
),
pk_iter,
)
.await?
} else {
Expand All @@ -298,6 +304,7 @@ where
self.args.sim_settings,
ep.mempool_configs.clone(),
),
pk_iter,
)
.await?
};
Expand All @@ -307,19 +314,21 @@ where
Ok((sender_handles, bundle_sender_actions))
}

async fn create_builders_v0_7<C, E>(
async fn create_builders_v0_7<C, E, I>(
&self,
ep: &EntryPointBuilderSettings,
provider: Arc<EthersProvider<C>>,
submit_provider: Option<Arc<EthersProvider<C>>>,
ep_v0_7: E,
pk_iter: &mut I,
) -> anyhow::Result<(
Vec<JoinHandle<anyhow::Result<()>>>,
Vec<mpsc::Sender<BundleSenderAction>>,
)>
where
C: JsonRpcClient + 'static,
E: EntryPointProvider<v0_7::UserOperation> + Clone,
I: Iterator<Item = String>,
{
info!("Mempool config for ep v0.7: {:?}", ep.mempool_configs);
let mut sender_handles = vec![];
Expand All @@ -336,6 +345,7 @@ where
ep_v0_7.clone(),
self.args.sim_settings,
),
pk_iter,
)
.await?
} else {
Expand All @@ -350,6 +360,7 @@ where
self.args.sim_settings,
ep.mempool_configs.clone(),
),
pk_iter,
)
.await?
};
Expand All @@ -359,13 +370,14 @@ where
Ok((sender_handles, bundle_sender_actions))
}

async fn create_bundle_builder<UO, E, S, C>(
async fn create_bundle_builder<UO, E, S, C, I>(
&self,
index: u64,
provider: Arc<EthersProvider<C>>,
submit_provider: Option<Arc<EthersProvider<C>>>,
entry_point: E,
simulator: S,
pk_iter: &mut I,
) -> anyhow::Result<(
JoinHandle<anyhow::Result<()>>,
mpsc::Sender<BundleSenderAction>,
Expand All @@ -376,10 +388,11 @@ where
E: EntryPointProvider<UO> + Clone,
S: Simulator<UO = UO>,
C: JsonRpcClient + 'static,
I: Iterator<Item = String>,
{
let (send_bundle_tx, send_bundle_rx) = mpsc::channel(1);

let signer = if let Some(pk) = &self.args.private_key {
let signer = if let Some(pk) = pk_iter.next() {
info!("Using local signer");
BundlerSigner::Local(
LocalSigner::connect(
Expand Down
5 changes: 4 additions & 1 deletion docs/cli.md
Original file line number Diff line number Diff line change
Expand Up @@ -166,10 +166,13 @@ List of command line options for configuring the Builder.
- *Only required when running in distributed mode*
- `--builder.private_key`: Private key to use for signing transactions
- env: *BUILDER_PRIVATE_KEY*
- *Always used if provided. If not provided builder.aws_kms_key_ids is used*
- **DEPRECATED**: Use `--builder.private_keys` instead. If both used this is added to the list.
- `--builder.private_keys`: Private keys to use for signing transactions, separated by `,`
- env: *BUILDER_PRIVATE_KEYS*
- `--builder.aws_kms_key_ids`: AWS KMS key IDs to use for signing transactions (comma-separated)
- env: *BUILDER_AWS_KMS_KEY_IDS*
- *Only required if BUILDER_PRIVATE_KEY is not provided*
- *Cannot use `builder.private_keys` and `builder.aws_kms_key_ids` at the same time*
- `--builder.redis_uri`: Redis URI to use for KMS leasing (default: `""`)
- env: *BUILDER_REDIS_URI*
- *Only required when AWS_KMS_KEY_IDS are provided*
Expand Down

0 comments on commit 9b60bf4

Please sign in to comment.