From 31cf3044a4490d7da6ca0fd8d166b842c797a503 Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@edera.dev>
Date: Tue, 5 Mar 2024 20:47:53 -0800
Subject: [PATCH] controller: implement support for dialing urls for transport

---
 controller/Cargo.toml         |  1 +
 controller/bin/control.rs     |  7 +++---
 controller/src/client.rs      | 40 ++++++++++++++++++++++++++++++++++-
 scripts/krata-debug-common.sh |  2 +-
 shared/src/lib.rs             |  3 +++
 5 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/controller/Cargo.toml b/controller/Cargo.toml
index a672fc4..e4c616c 100644
--- a/controller/Cargo.toml
+++ b/controller/Cargo.toml
@@ -16,6 +16,7 @@ clap = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }
 tokio-native-tls = { workspace = true }
+url = { workspace = true }
 
 [dependencies.krata]
 path = "../shared"
diff --git a/controller/bin/control.rs b/controller/bin/control.rs
index f617200..f1aab14 100644
--- a/controller/bin/control.rs
+++ b/controller/bin/control.rs
@@ -8,12 +8,12 @@ use kratactl::{
     client::{KrataClient, KrataClientTransport},
     console::XenConsole,
 };
-use tokio::net::UnixStream;
+use url::Url;
 
 #[derive(Parser, Debug)]
 #[command(version, about)]
 struct ControllerArgs {
-    #[arg(long, default_value = "/var/lib/krata/daemon.socket")]
+    #[arg(short, long, default_value = "unix:///var/lib/krata/daemon.socket")]
     connection: String,
 
     #[command(subcommand)]
@@ -53,8 +53,7 @@ async fn main() -> Result<()> {
     env_logger::Builder::from_env(Env::default().default_filter_or("warn")).init();
 
     let args = ControllerArgs::parse();
-    let stream = UnixStream::connect(&args.connection).await?;
-    let transport = KrataClientTransport::from_unix(stream).await?;
+    let transport = KrataClientTransport::dial(Url::parse(&args.connection)?).await?;
     let client = KrataClient::new(transport).await?;
 
     match args.command {
diff --git a/controller/src/client.rs b/controller/src/client.rs
index ef031f3..1e06dd1 100644
--- a/controller/src/client.rs
+++ b/controller/src/client.rs
@@ -4,6 +4,7 @@ use anyhow::{anyhow, Result};
 use krata::{
     control::{Message, Request, RequestBox, Response},
     stream::{ConnectionStreams, StreamContext},
+    KRATA_DEFAULT_TCP_PORT, KRATA_DEFAULT_TLS_PORT,
 };
 use log::{trace, warn};
 use tokio::{
@@ -16,8 +17,9 @@ use tokio::{
     },
     task::JoinHandle,
 };
-use tokio_native_tls::TlsStream;
+use tokio_native_tls::{native_tls::TlsConnector, TlsStream};
 use tokio_stream::{wrappers::LinesStream, StreamExt};
+use url::{Host, Url};
 
 const QUEUE_MAX_LEN: usize = 100;
 
@@ -106,6 +108,42 @@ impl KrataClientTransport {
     transport_new!(from_tcp, TcpStream, process_tcp_stream);
     transport_new!(from_tls_tcp, TlsStream<TcpStream>, process_tls_tcp_stream);
 
+    pub async fn dial(url: Url) -> Result<KrataClientTransport> {
+        match url.scheme() {
+            "unix" => {
+                let stream = UnixStream::connect(url.path()).await?;
+                Ok(KrataClientTransport::from_unix(stream).await?)
+            }
+
+            "tcp" => {
+                let address = format!(
+                    "{}:{}",
+                    url.host().unwrap_or(Host::Domain("localhost")),
+                    url.port().unwrap_or(KRATA_DEFAULT_TCP_PORT)
+                );
+                let stream = TcpStream::connect(address).await?;
+                Ok(KrataClientTransport::from_tcp(stream).await?)
+            }
+
+            "tls" | "tls-insecure" => {
+                let insecure = url.scheme() == "tls-insecure";
+                let host = format!("{}", url.host().unwrap_or(Host::Domain("localhost")));
+                let address = format!("{}:{}", host, url.port().unwrap_or(KRATA_DEFAULT_TLS_PORT));
+                let stream = TcpStream::connect(address).await?;
+                let mut connector = TlsConnector::builder();
+                if insecure {
+                    connector.danger_accept_invalid_certs(true);
+                }
+                let connector = connector.build()?;
+                let connector = tokio_native_tls::TlsConnector::from(connector);
+                let stream = connector.connect(&host, stream).await?;
+                Ok(KrataClientTransport::from_tls_tcp(stream).await?)
+            }
+
+            _ => Err(anyhow!("unsupported url scheme: {}", url.scheme())),
+        }
+    }
+
     transport_processor!(process_unix_stream, UnixStream);
     transport_processor!(process_tcp_stream, TcpStream);
     transport_processor!(process_tls_tcp_stream, TlsStream<TcpStream>);
diff --git a/scripts/krata-debug-common.sh b/scripts/krata-debug-common.sh
index 1c7f9df..ea56ff6 100644
--- a/scripts/krata-debug-common.sh
+++ b/scripts/krata-debug-common.sh
@@ -22,7 +22,7 @@ build_and_run() {
   if [ "${KRATA_BUILD_INITRD}" = "1" ]
   then
     ./initrd/build.sh -q
-    sudo cp "target/initrd/initrd" "/var/lib/krata/default/initrd"
+    sudo cp "initrd/target/initrd" "/var/lib/krata/default/initrd"
   fi
   RUST_TARGET="$(./scripts/detect-rust-target.sh)"
   ./scripts/cargo.sh build ${CARGO_BUILD_FLAGS} --bin "${EXE_TARGET}"
diff --git a/shared/src/lib.rs b/shared/src/lib.rs
index fe2ee99..43c5b0a 100644
--- a/shared/src/lib.rs
+++ b/shared/src/lib.rs
@@ -2,3 +2,6 @@ pub mod control;
 pub mod ethtool;
 pub mod launchcfg;
 pub mod stream;
+
+pub const KRATA_DEFAULT_TCP_PORT: u16 = 4350;
+pub const KRATA_DEFAULT_TLS_PORT: u16 = 4353;