controller: internal client support for TCP and TLS connections

Cargo.toml

@@ -58,6 +58,7 @@ path-absolutize = "3.1.1"
 tokio-tun = "0.11.2"
 tokio-listener = "0.3.1"
 trait-variant = "0.1.1"
+tokio-native-tls = "0.3.1"
 signal-hook = "0.3.17"
 
 [workspace.dependencies.uuid]

controller/Cargo.toml

@@ -15,6 +15,7 @@ cli-tables = { workspace = true }
 clap = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }
+tokio-native-tls = { workspace = true }
 
 [dependencies.krata]
 path = "../shared"

controller/bin/control.rs

@@ -54,7 +54,7 @@ async fn main() -> Result<()> {
 
     let args = ControllerArgs::parse();
     let stream = UnixStream::connect(&args.connection).await?;
-    let transport = KrataClientTransport::new(stream).await?;
+    let transport = KrataClientTransport::from_unix(stream).await?;
     let client = KrataClient::new(transport).await?;
 
     match args.command {

controller/src/client.rs

@@ -8,7 +8,7 @@ use krata::{
 use log::{trace, warn};
 use tokio::{
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
-    net::{unix, UnixStream},
+    net::{TcpStream, UnixStream},
     select,
     sync::{
         mpsc::{channel, Receiver, Sender},
@@ -16,6 +16,7 @@ use tokio::{
     },
     task::JoinHandle,
 };
+use tokio_native_tls::TlsStream;
 use tokio_stream::{wrappers::LinesStream, StreamExt};
 
 const QUEUE_MAX_LEN: usize = 100;
@@ -32,15 +33,15 @@ impl Drop for KrataClientTransport {
     }
 }
 
-impl KrataClientTransport {
+macro_rules! transport_new {
-    pub async fn new(stream: UnixStream) -> Result<Self> {
+    ($name:ident, $stream:ty, $processor:ident) => {
-        let (read, write) = stream.into_split();
+        pub async fn $name(stream: $stream) -> Result<Self> {
             let (tx_sender, tx_receiver) = channel::<Message>(QUEUE_MAX_LEN);
             let (rx_sender, rx_receiver) = channel::<Message>(QUEUE_MAX_LEN);
 
             let task = tokio::task::spawn(async move {
                 if let Err(error) =
-                KrataClientTransport::process_unix_stream(read, write, rx_sender, tx_receiver).await
+                    KrataClientTransport::$processor(stream, rx_sender, tx_receiver).await
                 {
                     warn!("failed to process krata transport messages: {}", error);
                 }
@@ -52,13 +53,17 @@ impl KrataClientTransport {
                 task,
             })
         }
+    };
+}
 
-    async fn process_unix_stream(
+macro_rules! transport_processor {
-        read: unix::OwnedReadHalf,
+    ($name:ident, $stream:ty) => {
-        mut write: unix::OwnedWriteHalf,
+        async fn $name(
+            stream: $stream,
             rx_sender: Sender<Message>,
             mut tx_receiver: Receiver<Message>,
         ) -> Result<()> {
+            let (read, mut write) = tokio::io::split(stream);
             let mut read = LinesStream::new(BufReader::new(read).lines());
             loop {
                 select! {
@@ -93,6 +98,17 @@ impl KrataClientTransport {
             }
             Ok(())
         }
+    };
+}
+
+impl KrataClientTransport {
+    transport_new!(from_unix, UnixStream, process_unix_stream);
+    transport_new!(from_tcp, TcpStream, process_tcp_stream);
+    transport_new!(from_tls_tcp, TlsStream<TcpStream>, process_tls_tcp_stream);
+
+    transport_processor!(process_unix_stream, UnixStream);
+    transport_processor!(process_tcp_stream, TcpStream);
+    transport_processor!(process_tls_tcp_stream, TlsStream<TcpStream>);
 }
 
 type RequestsMap = Arc<Mutex<HashMap<u64, oneshot::Sender<Response>>>>;