feat: impl run and attach to container

7 files changed, 87 insertions, 13 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 18187b2..0a00c07 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2188,6 +2188,8 @@ dependencies = [
  "serde_ini",
  "serde_json",
  "thiserror",
+ "tokio",
+ "wspty",
 ]
 
 [[package]]
diff --git a/zone_core/src/container.rs b/zone_core/src/container.rs
index 4ac9dbf..9296fb6 100644
--- a/zone_core/src/container.rs
+++ b/zone_core/src/container.rs
@@ -43,6 +43,19 @@ pub struct CloneOptions {
     pub user: String,
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone, Args)]
+pub struct WebSocketOptions {
+    pub id: u64,
+    pub template: String,
+    pub user: String,
+}
+
+impl From<WebSocketOptions> for String {
+    fn from(val: WebSocketOptions) -> Self {
+        format!("{}-{}-{}", val.user, val.template, val.id)
+    }
+}
+
 impl<T> FilterContainer for T
 where
     T: Iterator,
diff --git a/zone_core/src/lib.rs b/zone_core/src/lib.rs
index 852fa05..e612766 100644
--- a/zone_core/src/lib.rs
+++ b/zone_core/src/lib.rs
@@ -1,6 +1,8 @@
 use std::net::{IpAddr, Ipv4Addr};
 
-pub use crate::container::{CloneOptions, Container, ContainerOptions, ContainerStatus};
+pub use crate::container::{
+    CloneOptions, Container, ContainerOptions, ContainerStatus, WebSocketOptions,
+};
 pub use crate::error::{Error, Result};
 
 mod error;
diff --git a/zone_nspawn/Cargo.toml b/zone_nspawn/Cargo.toml
index ef9a4ad..ccf0e5b 100644
--- a/zone_nspawn/Cargo.toml
+++ b/zone_nspawn/Cargo.toml
@@ -19,3 +19,5 @@ serde = "1.0.136"
 serde_ini = "0.2.0"
 serde_json = "1.0.78"
 thiserror = "1.0.30"
+tokio = { version = "1.17.0", default-features = false, features = ["process"] }
+wspty = "0.1.1"
diff --git a/zone_nspawn/src/nspawn.rs b/zone_nspawn/src/nspawn.rs
index 8c97464..3e4e095 100644
--- a/zone_nspawn/src/nspawn.rs
+++ b/zone_nspawn/src/nspawn.rs
@@ -1,5 +1,10 @@
+use std::{ffi::OsStr, path::PathBuf, process::Command};
+
+use tokio::sync::mpsc::UnboundedReceiver;
+use wspty::{PtyCommand, PtyMaster};
+
 use crate::{Container, Error, Result};
-use std::{path::PathBuf, process::Command};
+
 #[derive(Default, Debug)]
 pub struct NSpawn;
 
@@ -39,6 +44,51 @@ impl NSpawn {
             .ok_or_else(|| Error::NSpawn(format!("Failed to create container: {:?}", self)))
     }
 
+    pub async fn spawn<O, C, S>(
+        opts: O,
+        cmd: C,
+        kill_rx: UnboundedReceiver<()>,
+    ) -> Result<PtyMaster>
+    where
+        O: IntoIterator<Item = S>,
+        C: IntoIterator<Item = S>,
+        S: AsRef<OsStr>,
+    {
+        let base_opts = ["--quiet", "--wait", "--collect", "--service-type=exec"];
+
+        let mut proc = tokio::process::Command::new("systemd-run");
+
+        proc.args(base_opts)
+            .args(opts)
+            .args(cmd)
+            .env("TERM", "xterm-256color");
+
+        PtyCommand::from(proc)
+            .run(kill_rx)
+            .await
+            .map_err(Error::from)
+    }
+
+    pub async fn attach(&self, name: String, kill_rx: UnboundedReceiver<()>) -> Result<PtyMaster> {
+        let opts = [
+            &format!("{}={}", "--machine", name),
+            "--pty",
+            "--send-sighup",
+        ];
+        let cmd = ["/usr/bin/login", "-H", "-f", "root"];
+        Self::spawn(opts, cmd, kill_rx).await
+    }
+
+    pub async fn run(&self, name: String, kill_rx: UnboundedReceiver<()>) -> Result<PtyMaster> {
+        let opts = [
+            &format!("{}={}", "--machine", name),
+            "--pty",
+            "--send-sighup",
+        ];
+        let cmd = ["/usr/bin/login", "-H", "-f", "root"];
+        Self::spawn(opts, cmd, kill_rx).await
+    }
+
     pub fn shutdown(name: String) -> Result<()> {
         Command::new("machinectl")
             .arg("poweroff")
diff --git a/zoned/src/http.rs b/zoned/src/http.rs
index 4ed289a..b3d5149 100644
--- a/zoned/src/http.rs
+++ b/zoned/src/http.rs
@@ -7,7 +7,7 @@ use axum::{
 };
 use std::{process::Command, sync::Arc};
 use tracing::{info, instrument, warn};
-use zone_core::{CloneOptions, Container, ContainerOptions, ContainerStatus, FilterContainer};
+use zone_core::{CloneOptions, Container, ContainerOptions, ContainerStatus, FilterContainer, WebSocketOptions};
 
 use crate::{ws, Error, Result, State};
 
@@ -105,13 +105,14 @@ async fn clone_container(
 #[instrument(ret, skip_all)]
 async fn ws_upgrade(
     ws: WebSocketUpgrade,
-    user_agent: Option<TypedHeader<headers::UserAgent>>,
+    Json(options): Json<WebSocketOptions>,
     Extension(state): Extension<Arc<State>>,
+    user_agent: Option<TypedHeader<headers::UserAgent>>,
 ) -> impl IntoResponse {
     let agent = user_agent.map_or("Unknown".to_string(), |u| u.to_string());
     info!(%agent, "Client connected");
 
-    ws.on_upgrade(|socket| ws::handler(socket, state))
+    ws.on_upgrade(|socket| ws::handler(socket, options, state))
 }
 
 #[cfg(test)]
diff --git a/zoned/src/ws.rs b/zoned/src/ws.rs
index 4d441f5..b7714da 100644
--- a/zoned/src/ws.rs
+++ b/zoned/src/ws.rs
@@ -9,11 +9,11 @@ use serde::Deserialize;
 use std::sync::Arc;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
-    process::Command,
     sync::mpsc,
 };
 use tracing::{instrument, warn};
-use wspty::{PtyCommand, PtyMaster};
+use wspty::PtyMaster;
+use zone_core::WebSocketOptions;
 
 use crate::{Result, State};
 
@@ -24,15 +24,16 @@ struct WindowSize {
 }
 
 #[instrument(err, skip_all)]
-pub async fn handler(ws_stream: WebSocket, _state: Arc<State>) -> Result<()> {
+pub async fn handler(
+    ws_stream: WebSocket,
+    options: WebSocketOptions,
+    state: Arc<State>,
+) -> Result<()> {
     let (sender, receiver) = ws_stream.split();
     let (kill_tx, kill_rx) = mpsc::unbounded_channel();
     let (tx, rx) = mpsc::channel(1024);
 
-    let mut cmd = Command::new("bash");
-    cmd.arg("-l").env("TERM", "xterm-256color");
-
-    let pty = PtyCommand::from(cmd).run(kill_rx).await?;
+    let pty = state.nspawn.attach(options.into(), kill_rx).await?;
 
     tokio::select! {
         res = msg_handler(receiver, pty.clone(), tx.clone(), kill_tx) => res,
@@ -79,7 +80,10 @@ async fn msg_handler(
 }
 
 #[instrument(err, skip_all)]
-async fn pty_handler(mut pty_read: PtyMaster, tx: mpsc::Sender<Message>) -> Result<()> {
+async fn pty_handler<R>(mut pty_read: R, tx: mpsc::Sender<Message>) -> Result<()>
+where
+    R: AsyncReadExt + Sized + Unpin,
+{
     let mut buffer = BytesMut::with_capacity(1024);
     buffer.resize(1024, 0u8);
     loop {