diff --git a/packages/local-replica/src/gateway.rs b/packages/local-replica/src/gateway.rs index 0a3c7d5..bbb04e8 100644 --- a/packages/local-replica/src/gateway.rs +++ b/packages/local-replica/src/gateway.rs @@ -23,10 +23,32 @@ use tracing_subscriber::{ Registry as TracingRegistry, layer::SubscriberExt, reload, util::SubscriberInitExt, }; +pub enum ReplicaUrl { + Mainnet(Url), + PocketIc(Url), +} + +impl ReplicaUrl { + pub fn new_remote(url: Url) -> Self { + ReplicaUrl::Mainnet(url) + } + + pub fn new_pocket_ic(url: Url) -> Self { + ReplicaUrl::PocketIc(url) + } + + pub fn into_url(&self) -> &Url { + match self { + ReplicaUrl::Mainnet(url) => url, + ReplicaUrl::PocketIc(url) => url, + } + } +} + pub async fn start_gateway( listen_ip_addr: &str, listen_port: u16, - replica_url: &Url, + replica_url: &ReplicaUrl, shutdown_token: CancellationToken, ) -> Result< ( @@ -37,7 +59,7 @@ pub async fn start_gateway( > { let listen_addr = format!("{listen_ip_addr}:{listen_port}"); - let gateway_args = vec![ + let mut gateway_args = vec![ "", "--domain", "localhost", @@ -45,12 +67,27 @@ pub async fn start_gateway( listen_ip_addr, "--domain-canister-id-from-query-params", "--domain-canister-id-from-referer", - "--ic-unsafe-root-key-fetch", "--listen-plain", &listen_addr, ]; - let (router, tasks) = - create_http_gateway_router(gateway_args, replica_url, shutdown_token.clone()).await?; + + match replica_url { + ReplicaUrl::Mainnet(url) => { + gateway_args.push("--ic-use-discovery"); + gateway_args.push("--ic-url"); + gateway_args.push(url.as_str()); + } + ReplicaUrl::PocketIc(_) => { + gateway_args.push("--ic-unsafe-root-key-fetch"); + } + } + + let (router, tasks) = create_http_gateway_router( + gateway_args, + &replica_url.into_url(), + shutdown_token.clone(), + ) + .await?; tasks.start(); diff --git a/packages/local-replica/src/main.rs b/packages/local-replica/src/main.rs index 84a77a7..f8ec5d5 100644 --- a/packages/local-replica/src/main.rs +++ b/packages/local-replica/src/main.rs @@ -1,9 +1,14 @@ mod gateway; mod pocket_ic; +use clap::Parser; +use ic_gateway::ic_bn_lib::reqwest::Url; use tokio_util::sync::CancellationToken; -use crate::{gateway::start_gateway, pocket_ic::start_pocket_ic}; +use crate::{ + gateway::{ReplicaUrl, start_gateway}, + pocket_ic::start_pocket_ic, +}; const PACKAGE_DIR: &str = env!("CARGO_MANIFEST_DIR"); const POCKET_IC_SERVER_BIN_PATH: &str = "bin/pocket-ic"; @@ -11,20 +16,39 @@ const POCKET_IC_SERVER_BIN_PATH: &str = "bin/pocket-ic"; const LOCAL_REPLICA_HTTP_LISTEN_IP_ADDR: &str = "127.0.0.1"; const LOCAL_REPLICA_HTTP_LISTEN_PORT: u16 = 4943; +#[derive(Parser, Debug)] +#[command(name = "replica")] +#[command(about = "Local replica server with HTTP gateway", long_about = None)] +struct Args { + /// Replica URL to use (if set, the PocketIC server won't be started) + #[arg(long, value_name = "URL")] + replica_url: Option, +} + #[tokio::main] async fn main() -> Result<(), Box> { + let args = Args::parse(); let shutdown_token = CancellationToken::new(); - // Start PocketIC server - let pic_path = format!("{PACKAGE_DIR}/{POCKET_IC_SERVER_BIN_PATH}"); - let (pic, pic_url) = start_pocket_ic(&pic_path).await; - println!("PocketIC Server URL: {}", pic_url); + // Determine replica URL and optionally start PocketIC server + let (replica_url, pic_handle) = if let Some(url_str) = args.replica_url { + // Use provided replica URL, don't start PocketIC + let url = Url::parse(&url_str)?; + println!("Using replica URL: {}", url_str); + (ReplicaUrl::new_remote(url), None) + } else { + // Start PocketIC server if no replica URL is provided + let pic_path = format!("{PACKAGE_DIR}/{POCKET_IC_SERVER_BIN_PATH}"); + let (pic, pic_url) = start_pocket_ic(&pic_path).await; + println!("PocketIC Server URL: {}", pic_url); + (ReplicaUrl::new_pocket_ic(pic_url), Some(pic)) + }; // Setup gateway let (router, listener) = start_gateway( LOCAL_REPLICA_HTTP_LISTEN_IP_ADDR, LOCAL_REPLICA_HTTP_LISTEN_PORT, - &pic_url, + &replica_url, shutdown_token.clone(), ) .await?; @@ -40,7 +64,10 @@ async fn main() -> Result<(), Box> { tokio::signal::ctrl_c().await.ok(); println!("\nShutting down..."); shutdown_token.cancel(); - pic.drop().await; + if let Some(pic) = pic_handle { + pic.drop().await; + println!("PocketIC server stopped"); + } }; axum::serve(listener, router)