feat: allow local gateway to point to mainnet

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,20 +59,35 @@ 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",
         "--domain",
         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();
 

packages/local-replica/src/main.rs

@@ -1,30 +1,54 @@
 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";
 
 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<String>,
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    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<dyn std::error::Error>> {
         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)