feat: Restructure VSOCK code as a starting point for improvements

Cargo.toml

@@ -199,9 +199,7 @@ members = [
     "rs/ic_os/sev/guest/testing",
     "rs/ic_os/sev/host",
     "rs/ic_os/utils",
-    "rs/ic_os/vsock/guest",
-    "rs/ic_os/vsock/host",
-    "rs/ic_os/vsock/vsock_lib",
+    "rs/ic_os/vsock",
     "rs/ingress_manager",
     "rs/interfaces",
     "rs/interfaces/adapter_client",

rs/ic_os/guest_upgrade/server/BUILD.bazel

@@ -16,7 +16,7 @@ rust_library(
         "//rs/ic_os/os_tools/guest_disk:guest_disk_lib",
         "//rs/ic_os/sev/attestation",
         "//rs/ic_os/sev/guest",
-        "//rs/ic_os/vsock/vsock_lib",
+        "//rs/ic_os/vsock:vsock_lib",
         "//rs/interfaces",
         "//rs/interfaces/registry",
         "//rs/protobuf",

rs/ic_os/guest_upgrade/server/Cargo.toml

@@ -20,7 +20,7 @@ thiserror = { workspace = true }
 tokio = { workspace = true }
 tokio-util = { workspace = true }
 tonic = { workspace = true, features = ["default"] }
-vsock_lib = { path = "../../vsock/vsock_lib" }
+vsock_lib = { path = "../../vsock" }
 x509-parser = { workspace = true }
 hyper = { workspace = true }
 hyper-util = { workspace = true }

rs/ic_os/guest_upgrade/tests/BUILD.bazel

@@ -14,7 +14,7 @@ rust_test(
         "//rs/ic_os/sev/attestation",
         "//rs/ic_os/sev/guest",
         "//rs/ic_os/sev/guest/testing",
-        "//rs/ic_os/vsock/vsock_lib",
+        "//rs/ic_os/vsock:vsock_lib",
         "//rs/interfaces/registry",
         "//rs/protobuf",
         "//rs/registry/fake",

rs/ic_os/guest_upgrade/tests/Cargo.toml

@@ -25,4 +25,4 @@ sev_guest = { path = "../../sev/guest" }
 sev_guest_testing = { path = "../../sev/guest/testing" }
 tempfile = { workspace = true }
 tokio = { workspace = true }
-vsock_lib = { path = "../../vsock/vsock_lib" }
+vsock_lib = { path = "../../vsock" }

rs/ic_os/release/BUILD.bazel

@@ -13,8 +13,8 @@ OBJECTS = {
     "setupos_tool": "//rs/ic_os/os_tools/setupos_tool:setupos_tool",
     "config_tool": "//rs/ic_os/config/tool:config_tool",
     "config_tool_dev": "//rs/ic_os/config/tool:config_tool_dev",
-    "vsock_guest": "//rs/ic_os/vsock/guest:vsock_guest",
-    "vsock_host": "//rs/ic_os/vsock/host:vsock_host",
+    "vsock_guest": "//rs/ic_os/vsock:vsock_guest",
+    "vsock_host": "//rs/ic_os/vsock:vsock_host",
     "metrics-proxy": "//rs/monitoring/metrics_proxy:metrics-proxy",
     "nss_icos": "//rs/ic_os/networking/nss_icos",
     "custom_metrics": "//rs/ic_os/metrics/custom_metrics:custom_metrics_bin",

rs/ic_os/vsock/vsock_lib/BUILD.bazel → rs/ic_os/vsock/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test")
+load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_library", "rust_test")
 
 package(default_visibility = ["//rs:ic-os-pkg"])
 
@@ -44,3 +44,30 @@ rust_test(
         "@crate_index//:vsock",
     ],
 )
+
+rust_binary(
+    name = "vsock_guest",
+    srcs = ["src/bin/guest.rs"],
+    crate_name = "vsock_guest",
+    target_compatible_with = [
+        "@platforms//os:linux",
+    ],
+    deps = [
+        # Keep sorted.
+        ":vsock_lib",
+        "@crate_index//:clap",
+    ],
+)
+
+rust_binary(
+    name = "vsock_host",
+    srcs = ["src/bin/host.rs"],
+    crate_name = "vsock_host",
+    target_compatible_with = [
+        "@platforms//os:linux",
+    ],
+    deps = [
+        # Keep sorted.
+        ":vsock_lib",
+    ],
+)

rs/ic_os/vsock/vsock_lib/Cargo.toml → rs/ic_os/vsock/Cargo.toml

@@ -7,6 +7,7 @@ edition.workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
+clap = { workspace = true } # only needed by bin, can't express in direct Cargo.toml
 mockall = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
@@ -21,5 +22,5 @@ rusb = { version = "0.9", features = ["vendored"] }
 sha2 = { workspace = true }
 tempfile = { workspace = true }
 vsock = "0.4"
-ic-http-utils = { path = "../../../http_utils" }
+ic-http-utils = { path = "../../http_utils" }
 tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }

rs/ic_os/vsock/guest/src/main.rs → rs/ic_os/vsock/src/bin/guest.rs

@@ -15,7 +15,11 @@ fn main() -> Result<(), String> {
     // Echo notify messages to the local GuestOS console so they are visible
     // in cloud environments where the host console is not accessible.
     if let Command::Notify(NotifyData { ref message, .. }) = command {
-        write_to_guest_console(message);
+        for path in ["/dev/tty1", "/dev/ttyS0"] {
+            if let Ok(mut tty) = OpenOptions::new().write(true).open(path) {
+                let _ = writeln!(tty, "\n{message}");
+            }
+        }
     }
 
     let payload = LinuxVSockClient::with_port(port).send_command(command)?;
@@ -30,14 +34,6 @@ fn main() -> Result<(), String> {
     Ok(())
 }
 
-fn write_to_guest_console(message: &str) {
-    for path in ["/dev/tty1", "/dev/ttyS0"] {
-        if let Ok(mut tty) = OpenOptions::new().write(true).open(path) {
-            let _ = writeln!(tty, "\n{message}");
-        }
-    }
-}
-
 #[derive(Debug, Parser)]
 #[clap(
     version = "1.0.0",

rs/ic_os/vsock/src/client.rs

@@ -0,0 +1,114 @@
+use crate::protocol::{Command, Response};
+use mockall::automock;
+
+#[automock]
+pub trait VSockClient {
+    fn send_command(&self, command: Command) -> Response;
+}
+
+#[cfg(target_os = "linux")]
+pub use linux::*;
+
+#[cfg(target_os = "linux")]
+mod linux {
+    use super::*;
+    use crate::protocol::{Request, Response};
+    use std::io::{Read, Write};
+    use vsock::{VMADDR_CID_HOST, VsockStream};
+
+    pub struct LinuxVSockClient {
+        port: u32,
+    }
+
+    impl LinuxVSockClient {
+        pub const DEFAULT_PORT: u32 = 19090;
+
+        pub fn with_port(port: u32) -> Self {
+            Self { port }
+        }
+    }
+
+    impl Default for LinuxVSockClient {
+        fn default() -> Self {
+            Self::with_port(Self::DEFAULT_PORT)
+        }
+    }
+
+    impl VSockClient for LinuxVSockClient {
+        fn send_command(&self, command: Command) -> Response {
+            let guest_cid = vsock::get_local_cid().map_err(|e| e.to_string())?;
+
+            let request = Request { guest_cid, command };
+
+            let mut stream = VsockStream::connect_with_cid_port(VMADDR_CID_HOST, self.port)
+                .map_err(|e| e.to_string())?;
+            // Set a long timeout, so HostOS has enough time to upgrade.
+            stream
+                .set_write_timeout(Some(std::time::Duration::from_secs(60 * 5)))
+                .map_err(|e| e.to_string())?;
+            stream
+                .set_read_timeout(Some(std::time::Duration::from_secs(60 * 5)))
+                .map_err(|e| e.to_string())?;
+
+            let json_request = serde_json::to_string(&request).map_err(|e| e.to_string())?;
+            stream
+                .write_all(json_request.as_bytes())
+                .map_err(|e| e.to_string())?;
+
+            // 64 KiB - generous for current responses (typically <1 KiB) while
+            // preventing unbounded allocation from a misbehaving host.
+            const MAX_RESPONSE_SIZE: u64 = 64 * 1024;
+            let mut response_str = String::new();
+            stream
+                .take(MAX_RESPONSE_SIZE)
+                .read_to_string(&mut response_str)
+                .map_err(|e| e.to_string())?;
+
+            parse_response(response_str.as_str())
+        }
+    }
+}
+
+/// Parse a response in a json string to a `Response` struct.
+pub fn parse_response(json_str: &str) -> Response {
+    if let Ok(response) = serde_json::from_str::<Response>(json_str) {
+        return response;
+    }
+    Err("Unable to parse host response: ".to_string() + json_str)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::protocol::{HostOSVsockVersion, Payload};
+
+    #[test]
+    fn test_parse_response() {
+        assert_eq!(
+            Ok(Payload::NoPayload),
+            parse_response("{\"Ok\":\"NoPayload\"}")
+        );
+        assert_eq!(
+            Ok(Payload::HostOSVersion("123".to_string())),
+            parse_response("{\"Ok\":{\"HostOSVersion\":\"123\"}}")
+        );
+        assert_eq!(
+            Ok(Payload::HostOSVsockVersion(HostOSVsockVersion {
+                major: 1,
+                minor: 0,
+                patch: 0,
+            })),
+            parse_response(
+                "{\"Ok\":{\"HostOSVsockVersion\":{\"major\":1,\"minor\":0,\"patch\":0}}}"
+            )
+        );
+        assert_eq!(
+            Err("Unable to parse host response: Error response".to_string()),
+            parse_response("Error response")
+        );
+
+        let json_str = r#"{"Ok":"NoPayload"#; // Missing closing brace
+        let response = parse_response(json_str);
+        assert!(response.is_err());
+    }
+}

rs/ic_os/vsock/src/lib.rs

@@ -0,0 +1,11 @@
+mod client;
+#[cfg(target_os = "linux")]
+pub use client::LinuxVSockClient;
+pub use client::{MockVSockClient, VSockClient};
+
+#[cfg(target_os = "linux")]
+mod server;
+#[cfg(target_os = "linux")]
+pub use server::run_server;
+
+pub mod protocol;