feat: Add an optional root-key argument to canister commands

CHANGELOG.md

@@ -14,6 +14,7 @@
 * feat: `icp canister logs` to display the current canister logs
   * use `--follow` to continuously poll for new logs. `--interval <n>` to poll every `n` seconds
 * feat: Support `k`, `m`, `b`, `t` suffixes in `.yaml` files when specifying cycles amounts
+* feat: Add an optional root-key argument to canister commands
 
 # v0.1.0
 

crates/icp-cli/src/commands/args.rs

@@ -13,7 +13,7 @@ use crate::options::{EnvironmentOpt, IdentityOpt, NetworkOpt};
 #[derive(Args, Debug)]
 pub(crate) struct CanisterCommandArgs {
     // Note: Could have flattened CanisterEnvironmentArg to avoid adding child field
-    /// Name or principal of canister to target
+    /// Name or principal of canister to target.
     /// When using a name an environment must be specified.
     pub(crate) canister: Canister,
 

crates/icp-cli/src/options.rs

@@ -1,4 +1,5 @@
-use clap::{ArgGroup, Args};
+use clap::error::ErrorKind;
+use clap::{ArgGroup, ArgMatches, Args, FromArgMatches};
 use icp::context::{EnvironmentSelection, NetworkSelection};
 use icp::identity::IdentitySelection;
 use icp::prelude::LOCAL;
@@ -63,22 +64,114 @@ impl From<EnvironmentOpt> for EnvironmentSelection {
     }
 }
 
+#[derive(Clone, Debug)]
+pub(crate) struct RootKey(pub Vec<u8>);
+
+fn parse_root_key(input: &str) -> Result<RootKey, String> {
+    let v = hex::decode(input).map_err(|e| format!("Invalid root key hex string: {e}"))?;
+    if v.len() != 133 {
+        Err(format!(
+            "Invalid root key. Expected 133 bytes but got {}",
+            v.len()
+        ))
+    } else {
+        Ok(RootKey(v))
+    }
+}
+
+#[derive(Clone, Debug)]
+enum NetworkTarget {
+    Url(Url),
+    Named(String),
+}
+
+fn parse_network_target(input: &str) -> Result<NetworkTarget, String> {
+    match Url::parse(input) {
+        Ok(url) => Ok(NetworkTarget::Url(url)),
+        Err(_) => Ok(NetworkTarget::Named(input.to_string())),
+    }
+}
+
 #[derive(Args, Clone, Debug, Default)]
 #[clap(group(ArgGroup::new("network-select").multiple(false)))]
-pub(crate) struct NetworkOpt {
-    /// Name of the network to target, conflicts with environment argument
-    #[arg(long, short = 'n', env = "ICP_NETWORK", group = "network-select", help_heading = heading::NETWORK_PARAMETERS)]
-    network: Option<String>,
+pub(crate) struct NetworkOptInner {
+    /// Name or URL of the network to target, conflicts with environment argument
+    #[arg(long, short = 'n', env = "ICP_NETWORK", group = "network-select", help_heading = heading::NETWORK_PARAMETERS, value_parser = parse_network_target)]
+    network: Option<NetworkTarget>,
+
+    /// The root key to use if connecting to a network by URL.
+    /// Required when using `--network <URL>`.
+    #[arg(long, short = 'k', requires = "network", help_heading = heading::NETWORK_PARAMETERS, value_parser = parse_root_key)]
+    root_key: Option<RootKey>,
+}
+
+// This is wrapper around NetworkOptInner that will do some additional
+// validation to only allow --root-key when the network is a url.
+#[derive(Clone, Debug, Default)]
+pub(crate) enum NetworkOpt {
+    Url(Url, RootKey),
+
+    Name(String),
+
+    #[default]
+    None,
+}
+
+impl FromArgMatches for NetworkOpt {
+    fn from_arg_matches(matches: &ArgMatches) -> Result<Self, clap::Error> {
+        let inner = NetworkOptInner::from_arg_matches(matches)?;
+
+        match (inner.network, inner.root_key) {
+            // Case: We have a URL, so we REQUIRE the root key
+            (Some(NetworkTarget::Url(url)), Some(key)) => Ok(NetworkOpt::Url(url, key)),
+
+            // ERROR Case: URL provided but missing root key
+            (Some(NetworkTarget::Url(_)), None) => Err(clap::Error::raw(
+                ErrorKind::MissingRequiredArgument,
+                "`--root-key` is required when `--network` is a URL.\n",
+            )),
+
+            // Case: Named network (root key should be empty)
+            (Some(NetworkTarget::Named(name)), None) => Ok(NetworkOpt::Name(name)),
+
+            // ERROR case: Name provided with a root key
+            (Some(NetworkTarget::Named(_)), Some(_)) => Err(clap::Error::raw(
+                ErrorKind::MissingRequiredArgument,
+                "`--root-key` is only valid when `--network` is a URL.\n",
+            )),
+
+            // Case: No network specified
+            (None, None) => Ok(NetworkOpt::None),
+
+            // Case: Should be impossible, --root-key is passed without a network argument
+            (None, Some(_)) => {
+                panic!("Invalid cli arg combination: --root-key without a --network <NETWORK>")
+            }
+        }
+    }
+
+    fn update_from_arg_matches(&mut self, matches: &ArgMatches) -> Result<(), clap::Error> {
+        // For simple wrappers, we can just replace the current state
+        *self = Self::from_arg_matches(matches)?;
+        Ok(())
+    }
+}
+
+impl Args for NetworkOpt {
+    fn augment_args(cmd: clap::Command) -> clap::Command {
+        NetworkOptInner::augment_args(cmd)
+    }
+    fn augment_args_for_update(cmd: clap::Command) -> clap::Command {
+        NetworkOptInner::augment_args_for_update(cmd)
+    }
 }
 
 impl From<NetworkOpt> for NetworkSelection {
     fn from(v: NetworkOpt) -> Self {
-        match v.network {
-            Some(network) => match Url::parse(&network) {
-                Ok(url) => NetworkSelection::Url(url),
-                Err(_) => NetworkSelection::Named(network),
-            },
-            None => NetworkSelection::Default,
+        match v {
+            NetworkOpt::Url(url, RootKey(key)) => NetworkSelection::Url(url, key),
+            NetworkOpt::Name(name) => NetworkSelection::Named(name),
+            NetworkOpt::None => NetworkSelection::Default,
         }
     }
 }

crates/icp-cli/tests/canister_call_root_key_tests.rs

@@ -0,0 +1,139 @@
+use indoc::formatdoc;
+use predicates::ord::eq;
+use predicates::str::{PredicateStrExt, contains};
+
+use crate::common::{ENVIRONMENT_RANDOM_PORT, NETWORK_RANDOM_PORT, TestContext};
+use icp::fs::write_string;
+
+mod common;
+
+#[tokio::test]
+async fn canister_call_with_url_and_root_key() {
+    let ctx = TestContext::new();
+
+    // Setup project
+    let project_dir = ctx.create_project_dir("icp");
+
+    // Use vendored WASM
+    let wasm = ctx.make_asset("example_icp_mo.wasm");
+
+    // Project manifest
+    let pm = formatdoc! {r#"
+        canisters:
+          - name: my-canister
+            build:
+              steps:
+                - type: script
+                  command: cp '{wasm}' "$ICP_WASM_OUTPUT_PATH"
+
+        {NETWORK_RANDOM_PORT}
+        {ENVIRONMENT_RANDOM_PORT}
+    "#};
+
+    write_string(&project_dir.join("icp.yaml"), &pm).expect("failed to write project manifest");
+
+    // Start network
+    let _g = ctx.start_network_in(&project_dir, "random-network").await;
+    ctx.ping_until_healthy(&project_dir, "random-network");
+
+    // Deploy canister
+    ctx.icp()
+        .current_dir(&project_dir)
+        .args([
+            "deploy",
+            "my-canister",
+            "--environment",
+            "random-environment",
+        ])
+        .assert()
+        .success();
+
+    // Test calling with Candid text format
+    ctx.icp()
+        .current_dir(&project_dir)
+        .args([
+            "canister",
+            "call",
+            "--environment",
+            "random-environment",
+            "my-canister",
+            "greet",
+            "(\"world\")",
+        ])
+        .assert()
+        .success()
+        .stdout(eq("(\"Hello, world!\")").trim());
+
+    // Get the network information so we can call the network directly
+    let assert = ctx
+        .icp()
+        .current_dir(&project_dir)
+        .args([
+            "network",
+            "status",
+            "--environment",
+            "random-environment",
+            "--json",
+        ])
+        .assert()
+        .success();
+    let output = assert.get_output();
+
+    let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap();
+    let gateway_url = json["gateway_url"].as_str().expect("Should be a string");
+    let root_key = json["root_key"].as_str().expect("Should be a string");
+
+    // Get the canister information so we can call the network directly
+    let assert = ctx
+        .icp()
+        .current_dir(&project_dir)
+        .args([
+            "canister",
+            "status",
+            "--environment",
+            "random-environment",
+            "my-canister",
+            "--id-only",
+        ])
+        .assert()
+        .success();
+
+    let output = assert.get_output();
+    let canister_id =
+        String::from_utf8(output.stdout.clone()).expect("canister id should be a valid string");
+    let canister_id = canister_id.trim();
+
+    // Test calling with with url from external directory
+    ctx.icp()
+        .args([
+            "canister",
+            "call",
+            "--network",
+            gateway_url,
+            "--root-key",
+            root_key,
+            canister_id,
+            "greet",
+            "(\"world\")",
+        ])
+        .assert()
+        .success()
+        .stdout(eq("(\"Hello, world!\")").trim());
+
+    // Test calling with with url from external directory with bad root key
+    ctx.icp()
+        .args([
+            "canister",
+            "call",
+            "--network",
+            gateway_url,
+            "--root-key",
+            "badbadbad", // This is an invalid root key
+            canister_id,
+            "greet",
+            "(\"world\")",
+        ])
+        .assert()
+        .failure()
+        .stderr(contains("invalid value 'badbadbad' for '--root-key"));
+}

crates/icp-cli/tests/cycles_tests.rs

@@ -2,7 +2,7 @@ use indoc::formatdoc;
 use predicates::str::contains;
 
 use crate::common::{ENVIRONMENT_RANDOM_PORT, NETWORK_RANDOM_PORT, TestContext, clients};
-use icp::{fs::write_string, prelude::IC_MAINNET_NETWORK_API_URL};
+use icp::fs::write_string;
 
 mod common;
 
@@ -169,14 +169,7 @@ async fn cycles_mint_on_ic() {
     // Run mint command with --network
     ctx.icp()
         .current_dir(&project_dir)
-        .args([
-            "cycles",
-            "mint",
-            "--icp",
-            "1",
-            "--network",
-            IC_MAINNET_NETWORK_API_URL,
-        ])
+        .args(["cycles", "mint", "--icp", "1", "--network", "ic"])
         .assert()
         .stderr(contains(
             "Error: Insufficient funds: 1.00010000 ICP required, 0 ICP available.",

crates/icp/src/context/mod.rs

@@ -31,7 +31,7 @@ pub enum NetworkSelection {
     /// Use a named network
     Named(String),
     /// Use a network by URL
-    Url(Url),
+    Url(Url, Vec<u8>),
 }
 
 /// Selection type for environments - similar to IdentitySelection
@@ -180,13 +180,13 @@ impl Context {
                 }
             }
             NetworkSelection::Default => Err(GetNetworkError::DefaultNetwork),
-            NetworkSelection::Url(url) => Ok(crate::Network {
+            NetworkSelection::Url(url, root_key) => Ok(crate::Network {
                 name: url.to_string(),
                 configuration: crate::network::Configuration::Connected {
                     connected: crate::network::Connected {
                         api_url: url.clone(),
                         http_gateway_url: Some(url.clone()),
-                        root_key: IC_ROOT_KEY.to_vec(),
+                        root_key: root_key.to_vec(),
                     },
                 },
             }),
@@ -402,7 +402,7 @@ impl Context {
         match (environment, network) {
             // Error: Both environment and network specified
             (EnvironmentSelection::Named(_), NetworkSelection::Named(_))
-            | (EnvironmentSelection::Named(_), NetworkSelection::Url(_)) => {
+            | (EnvironmentSelection::Named(_), NetworkSelection::Url(_, _)) => {
                 Err(GetAgentError::EnvironmentAndNetworkSpecified)
             }
 
@@ -428,7 +428,7 @@ impl Context {
 
             // Network specified
             (EnvironmentSelection::Default, NetworkSelection::Named(_))
-            | (EnvironmentSelection::Default, NetworkSelection::Url(_)) => {
+            | (EnvironmentSelection::Default, NetworkSelection::Url(_, _)) => {
                 Ok(self.get_agent_for_network(identity, network).await?)
             }
         }
@@ -446,13 +446,13 @@ impl Context {
                 match (environment, network) {
                     // Error: Both environment and network specified
                     (EnvironmentSelection::Named(_), NetworkSelection::Named(_))
-                    | (EnvironmentSelection::Named(_), NetworkSelection::Url(_)) => {
+                    | (EnvironmentSelection::Named(_), NetworkSelection::Url(_, _)) => {
                         Err(GetCanisterIdError::CanisterEnvironmentAndNetworkSpecified)
                     }
 
                     // Error: Canister by name with explicit network but no environment
                     (EnvironmentSelection::Default, NetworkSelection::Named(_))
-                    | (EnvironmentSelection::Default, NetworkSelection::Url(_)) => {
+                    | (EnvironmentSelection::Default, NetworkSelection::Url(_, _)) => {
                         Err(GetCanisterIdError::AmbiguousCanisterName)
                     }