Merge branch 'eero/inject-deployment' into 'master'

[NODE-1089] Add deployment settings to `setupos-inject-configuration`

 

See merge request dfinity-lab/public/ic!15179

rs/ic_os/setupos-inject-configuration/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@rules_rust//rust:defs.bzl", "rust_binary")
+load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -8,6 +8,14 @@ DEPENDENCIES = [
     "@crate_index//:ipnet",
     "@crate_index//:tempfile",
     "@crate_index//:tokio",
+    "@crate_index//:serde",
+    "@crate_index//:serde_json",
+    "@crate_index//:serde_with",
+    "@crate_index//:url",
+]
+
+DEV_DEPENDENCIES = [
+    "@crate_index//:once_cell",
 ]
 
 MACRO_DEPENDENCIES = []
@@ -19,3 +27,11 @@ rust_binary(
     version = "0.1.0",
     deps = DEPENDENCIES,
 )
+
+rust_test(
+    name = "setupos_inject_configuration_test",
+    srcs = glob(["src/**/*.rs"]),
+    crate_root = "src/main.rs",
+    proc_macro_deps = MACRO_DEPENDENCIES,
+    deps = DEPENDENCIES + DEV_DEPENDENCIES,
+)

rs/ic_os/setupos-inject-configuration/Cargo.toml

@@ -9,3 +9,10 @@ clap = { version = "4.4.3", features = ["derive"] }
 ipnet = "2.8.0"
 tempfile = "3.8.0"
 tokio = { version = "1.32.0", features = ["full"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+url = { version = "2.1.1", features = ["serde"] }
+serde_with = "1.6.2"
+
+[dev-dependencies]
+once_cell = "1.8"

rs/ic_os/setupos-inject-configuration/src/deployment.rs

@@ -0,0 +1,89 @@
+use serde_with::{serde_as, DisplayFromStr};
+use url::Url;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
+pub struct DeploymentJson {
+    pub deployment: Deployment,
+    pub logging: Logging,
+    pub nns: Nns,
+    pub dns: Dns,
+    pub resources: Resources,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
+pub struct Deployment {
+    pub name: String,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
+pub struct Logging {
+    pub hosts: String,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
+pub struct Nns {
+    pub url: Url,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
+pub struct Dns {
+    pub name_servers: String,
+}
+
+#[serde_as]
+#[derive(Serialize, Deserialize, PartialEq, Debug)]
+pub struct Resources {
+    #[serde_as(as = "DisplayFromStr")]
+    pub memory: u32,
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use once_cell::sync::Lazy;
+
+    const DEPLOYMENT_STR: &str = r#"{
+  "deployment": {
+    "name": "mainnet"
+  },
+  "logging": {
+    "hosts": "elasticsearch-node-0.mercury.dfinity.systems:443 elasticsearch-node-1.mercury.dfinity.systems:443 elasticsearch-node-2.mercury.dfinity.systems:443 elasticsearch-node-3.mercury.dfinity.systems:443"
+  },
+  "nns": {
+    "url": "https://dfinity.org/"
+  },
+  "dns": {
+    "name_servers": "2606:4700:4700::1111 2606:4700:4700::1001 2001:4860:4860::8888 2001:4860:4860::8844"
+  },
+  "resources": {
+    "memory": "490"
+  }
+}"#;
+
+    static DEPLOYMENT_STRUCT: Lazy<DeploymentJson> = Lazy::new(|| {
+        DeploymentJson {
+            deployment: Deployment { name: "mainnet".to_string() },
+            logging: Logging { hosts: "elasticsearch-node-0.mercury.dfinity.systems:443 elasticsearch-node-1.mercury.dfinity.systems:443 elasticsearch-node-2.mercury.dfinity.systems:443 elasticsearch-node-3.mercury.dfinity.systems:443".to_string() },
+            nns: Nns { url: Url::parse("https://dfinity.org").unwrap() },
+            dns: Dns { name_servers: "2606:4700:4700::1111 2606:4700:4700::1001 2001:4860:4860::8888 2001:4860:4860::8844".to_string() },
+            resources: Resources { memory: 490 },
+        }
+    });
+
+    #[test]
+    fn read_deployment() {
+        let parsed_deployment: DeploymentJson = { serde_json::from_str(DEPLOYMENT_STR).unwrap() };
+
+        assert_eq!(*DEPLOYMENT_STRUCT, parsed_deployment);
+    }
+
+    #[test]
+    fn write_deployment() {
+        let written_deployment =
+            serde_json::to_string_pretty::<DeploymentJson>(&DEPLOYMENT_STRUCT).unwrap();
+
+        assert_eq!(DEPLOYMENT_STR, written_deployment);
+    }
+}

rs/ic_os/setupos-inject-configuration/src/main.rs

@@ -11,7 +11,10 @@ use loopdev::{create_loop_device, detach_loop_device};
 use sysmount::{mount, umount};
 use tempfile::tempdir;
 use tokio::fs;
+use url::Url;
 
+mod deployment;
+use deployment::DeploymentJson;
 mod loopdev;
 mod sysmount;
 
@@ -31,19 +34,33 @@ struct Cli {
 
     #[arg(long, value_delimiter = ',')]
     public_keys: Option<Vec<String>>,
+
+    #[command(flatten)]
+    deployment: DeploymentConfig,
 }
 
 #[derive(Args)]
-#[group(required = true)]
 struct NetworkConfig {
     #[arg(long)]
     ipv6_prefix: Option<Ipv6Net>,
 
     #[arg(long)]
     ipv6_gateway: Option<Ipv6Net>,
 
-    #[arg(long, conflicts_with_all = ["ipv6_prefix", "ipv6_gateway"])]
-    ipv6_address: Option<Ipv6Net>,
+    #[arg(long)]
+    mgmt_mac: Option<String>,
+}
+
+#[derive(Args)]
+struct DeploymentConfig {
+    #[arg(long)]
+    nns_url: Option<Url>,
+
+    #[arg(long, allow_hyphen_values = true)]
+    nns_public_key: Option<String>,
+
+    #[arg(long)]
+    memory_gb: Option<u32>,
 }
 
 #[tokio::main]
@@ -74,23 +91,7 @@ async fn main() -> Result<(), Error> {
         .context("failed to print previous config")?;
 
     // Update config.ini
-    let cfg = match (
-        cli.network.ipv6_prefix,
-        cli.network.ipv6_gateway,
-        cli.network.ipv6_address,
-    ) {
-        // PrefixAndGateway
-        (Some(ipv6_prefix), Some(ipv6_gateway), None) => {
-            Config::PrefixAndGateway(ipv6_prefix, ipv6_gateway)
-        }
-
-        // Address
-        (None, None, Some(ipv6_address)) => Config::Address(ipv6_address),
-
-        _ => panic!("invalid network arguments"),
-    };
-
-    write_config(&target_dir.path().join("config.ini"), &cfg)
+    write_config(&target_dir.path().join("config.ini"), &cli.network)
         .await
         .context("failed to write config file")?;
 
@@ -117,6 +118,41 @@ async fn main() -> Result<(), Error> {
             .context("failed to write public keys")?;
     }
 
+    // Unmount partition
+    umount(&target_dir.path().to_string_lossy())
+        .await
+        .context("failed to unmount partition")?;
+
+    let data_partition_path = format!("{device_path}p4");
+
+    // Mount data partition
+    let target_dir = tempdir().context("failed to create temporary dir")?;
+
+    mount(
+        &data_partition_path,                 // source
+        &target_dir.path().to_string_lossy(), // target
+    )
+    .await
+    .context("failed to mount partition")?;
+
+    // Print previous deployment.json
+    println!("Previous deployment.json:\n---");
+    print_file_contents(&target_dir.path().join("deployment.json"))
+        .await
+        .context("failed to print previous deployment config")?;
+
+    // Update deployment.json
+    update_deployment(&target_dir.path().join("deployment.json"), &cli.deployment)
+        .await
+        .context("failed to write deployment config file")?;
+
+    // Update NNS key
+    if let Some(public_key) = cli.deployment.nns_public_key {
+        let mut f = File::create(target_dir.path().join("nns_public_key.pem"))
+            .context("failed to create nns key file")?;
+        write!(&mut f, "{public_key}")?;
+    }
+
     // Unmount partition
     umount(&target_dir.path().to_string_lossy())
         .await
@@ -130,11 +166,6 @@ async fn main() -> Result<(), Error> {
     Ok(())
 }
 
-enum Config {
-    PrefixAndGateway(Ipv6Net, Ipv6Net),
-    Address(Ipv6Net),
-}
-
 async fn print_file_contents(path: &Path) -> Result<(), Error> {
     let s = fs::read_to_string(path).await;
 
@@ -145,24 +176,28 @@ async fn print_file_contents(path: &Path) -> Result<(), Error> {
     Ok(())
 }
 
-async fn write_config(path: &Path, cfg: &Config) -> Result<(), Error> {
+async fn write_config(path: &Path, cfg: &NetworkConfig) -> Result<(), Error> {
     let mut f = File::create(path).context("failed to create config file")?;
 
-    match cfg {
-        Config::PrefixAndGateway(ipv6_prefix, ipv6_gateway) => {
-            writeln!(
-                &mut f,
-                "ipv6_prefix={}",
-                ipv6_prefix.addr().to_string().trim_end_matches("::")
-            )?;
-
-            writeln!(&mut f, "ipv6_subnet=/{}", ipv6_prefix.prefix_len())?;
-            writeln!(&mut f, "ipv6_gateway={}", ipv6_gateway.addr())?;
-        }
-
-        Config::Address(ipv6_address) => {
-            writeln!(&mut f, "ipv6_address={}", ipv6_address.addr())?;
-        }
+    let NetworkConfig {
+        ipv6_prefix,
+        ipv6_gateway,
+        mgmt_mac,
+    } = cfg;
+
+    if let (Some(ipv6_prefix), Some(ipv6_gateway)) = (ipv6_prefix, ipv6_gateway) {
+        writeln!(
+            &mut f,
+            "ipv6_prefix={}",
+            ipv6_prefix.addr().to_string().trim_end_matches("::")
+        )?;
+
+        writeln!(&mut f, "ipv6_subnet=/{}", ipv6_prefix.prefix_len())?;
+        writeln!(&mut f, "ipv6_gateway={}", ipv6_gateway.addr())?;
+    }
+
+    if let Some(mgmt_mac) = mgmt_mac {
+        writeln!(&mut f, "mgmt_mac={}", mgmt_mac)?;
     }
 
     Ok(())
@@ -172,8 +207,31 @@ async fn write_public_keys(path: &Path, ks: Vec<String>) -> Result<(), Error> {
     let mut f = File::create(path).context("failed to create public keys file")?;
 
     for k in ks {
-        writeln!(&mut f, "{k}",)?;
+        writeln!(&mut f, "{k}")?;
     }
 
     Ok(())
 }
+
+async fn update_deployment(path: &Path, cfg: &DeploymentConfig) -> Result<(), Error> {
+    let mut deployment_json = {
+        let f = File::open(path).context("failed to open deployment config file")?;
+        let deployment_json: DeploymentJson = serde_json::from_reader(f)?;
+
+        deployment_json
+    };
+
+    if let Some(nns_url) = &cfg.nns_url {
+        deployment_json.nns.url = nns_url.clone();
+    }
+
+    if let Some(memory) = cfg.memory_gb {
+        deployment_json.resources.memory = memory;
+    }
+
+    let mut f = File::create(path).context("failed to open deployment config file")?;
+    let output = serde_json::to_string_pretty(&deployment_json)?;
+    write!(&mut f, "{output}")?;
+
+    Ok(())
+}