feat(cli): authenticate to a VMM with a bearer token

Author: kvinwang

dstack/crates/dstack-cli-core/src/vmm.rs

@@ -11,7 +11,7 @@
 use anyhow::{anyhow, bail, Result};
 use dstack_vmm_rpc::vmm_client::VmmClient;
 use dstack_vmm_rpc::{Id, StatusRequest, StatusResponse, VmConfiguration};
-use http_client::http_request;
+use http_client::http_request_with_headers;
 use http_client::prpc::PrpcClient;
 
 /// default local VMM control socket (created by `dstackup install`).
@@ -20,25 +20,46 @@ pub const DEFAULT_HOST: &str = "unix:/var/run/dstack/vmm.sock";
 /// a connection to a VMM — local unix socket or remote http endpoint.
 pub struct Vmm {
     rpc: VmmClient<PrpcClient>,
-    /// base string usable with [`http_request`] for non-prpc endpoints.
+    /// base string usable with [`http_request_with_headers`] for non-prpc endpoints.
     base: String,
+    /// bearer token sent with every request when the VMM has `[auth]` enabled.
+    auth_token: Option<String>,
 }
 
 impl Vmm {
     /// connect to a VMM addressed by `host`:
     /// `unix:/path/to/vmm.sock` (local) or `http(s)://host:port` (remote).
     pub fn connect(host: &str) -> Result<Self> {
+        Self::connect_with_token(host, None)
+    }
+
+    /// like [`Vmm::connect`], additionally sending `Authorization: Bearer
+    /// <token>` with every request — required when the VMM has `[auth]`
+    /// enabled. An empty or `None` token sends no header.
+    pub fn connect_with_token(host: &str, token: Option<&str>) -> Result<Self> {
         let host = host.trim();
+        let token = token.map(str::trim).filter(|t| !t.is_empty());
+        let with_token = |client: PrpcClient| match token {
+            Some(token) => client.with_bearer_token(token),
+            None => client,
+        };
         if let Some(sock) = host.strip_prefix("unix:") {
-            let rpc = VmmClient::new(PrpcClient::new_unix(sock.to_string(), "/prpc".to_string()));
+            let client = PrpcClient::new_unix(sock.to_string(), "/prpc".to_string());
+            let rpc = VmmClient::new(with_token(client));
             Ok(Self {
                 rpc,
                 base: format!("unix:{sock}"),
+                auth_token: token.map(str::to_string),
             })
         } else if host.starts_with("http://") || host.starts_with("https://") {
             let base = host.trim_end_matches('/').to_string();
-            let rpc = VmmClient::new(PrpcClient::new(format!("{base}/prpc")));
-            Ok(Self { rpc, base })
+            let client = PrpcClient::new(format!("{base}/prpc"));
+            let rpc = VmmClient::new(with_token(client));
+            Ok(Self {
+                rpc,
+                base,
+                auth_token: token.map(str::to_string),
+            })
         } else {
             bail!(
                 "unsupported host '{host}': expected unix:/path/to/vmm.sock or http(s)://host:port"
@@ -112,7 +133,14 @@ impl Vmm {
     /// socket and over the localhost HTTP endpoint written by `dstackup install`.
     pub async fn logs(&self, id: &str, lines: u32) -> Result<String> {
         let path = format!("/logs?id={id}&follow=false&ansi=false&lines={lines}");
-        let (status, body) = http_request("GET", &self.base, &path, b"").await?;
+        let auth_header;
+        let mut headers: Vec<(&str, &str)> = Vec::new();
+        if let Some(token) = &self.auth_token {
+            auth_header = format!("Bearer {token}");
+            headers.push(("Authorization", auth_header.as_str()));
+        }
+        let (status, body) =
+            http_request_with_headers("GET", &self.base, &path, b"", &headers).await?;
         if status != 200 {
             bail!("vmm /logs returned status {status}");
         }

dstack/crates/dstack-cli/src/main.rs

@@ -32,7 +32,9 @@ struct Cli {
     #[arg(long, global = true, value_name = "DIR")]
     prefix: Option<String>,
 
-    /// auth token for a remote VMM.
+    /// auth token for a VMM with `[auth]` enabled (sent as `Authorization:
+    /// Bearer`). Falls back to `DSTACK_VMM_TOKEN`, then the token written by
+    /// `dstackup install`.
     #[arg(long, global = true)]
     token: Option<String>,
 
@@ -105,20 +107,31 @@ enum Command {
 #[tokio::main]
 async fn main() -> Result<()> {
     let cli = Cli::parse();
-    // remote-auth wiring lands with the TLS+token transport.
-    let _ = &cli.token;
     let defaults = LocalDefaults::read(cli.prefix.as_deref());
     let use_local_defaults = cli.host.is_none();
     let host = cli
         .host
         .clone()
         .or_else(|| defaults.as_ref().and_then(|d| d.client_url.clone()))
         .unwrap_or_else(|| DEFAULT_HOST.to_string());
+    // auth token: --token, then DSTACK_VMM_TOKEN, then the token file written
+    // by `dstackup install` (local defaults only).
+    let token = cli
+        .token
+        .clone()
+        .or_else(|| std::env::var("DSTACK_VMM_TOKEN").ok())
+        .filter(|t| !t.trim().is_empty())
+        .or_else(|| {
+            use_local_defaults
+                .then(|| defaults.as_ref().and_then(LocalDefaults::token))
+                .flatten()
+        });
+    let token = token.as_deref();
     let json = cli.json;
 
     match cli.command {
-        Command::Apps => cmd_apps(&host, json).await,
-        Command::Logs { id, lines } => cmd_logs(&host, &id, lines).await,
+        Command::Apps => cmd_apps(&host, token, json).await,
+        Command::Logs { id, lines } => cmd_logs(&host, token, &id, lines).await,
         Command::Deploy {
             compose,
             compose_file,
@@ -149,6 +162,7 @@ async fn main() -> Result<()> {
             };
             cmd_deploy(
                 &host,
+                token,
                 &compose,
                 &name,
                 image.as_deref(),
@@ -178,6 +192,7 @@ fn resolve_compose_arg(positional: Option<String>, flagged: Option<String>) -> R
 
 struct LocalDefaults {
     client_url: Option<String>,
+    client_token_path: Option<String>,
     image: Option<String>,
     allowlist_path: Option<String>,
 }
@@ -197,6 +212,11 @@ impl LocalDefaults {
                 .and_then(|x| x.as_str())
                 .filter(|s| !s.is_empty())
                 .map(str::to_string),
+            client_token_path: v
+                .get("client_token_path")
+                .and_then(|x| x.as_str())
+                .filter(|s| !s.is_empty())
+                .map(str::to_string),
             image: v
                 .get("image")
                 .and_then(|x| x.as_str())
@@ -213,6 +233,14 @@ impl LocalDefaults {
     fn allowlist_path(&self) -> Option<String> {
         self.allowlist_path.clone()
     }
+
+    /// read the VMM API token from the file recorded by `dstackup install`.
+    fn token(&self) -> Option<String> {
+        let path = self.client_token_path.as_ref()?;
+        let token = std::fs::read_to_string(path).ok()?;
+        let token = token.trim();
+        (!token.is_empty()).then(|| token.to_string())
+    }
 }
 
 #[cfg(test)]
@@ -307,6 +335,7 @@ mod tests {
 #[allow(clippy::too_many_arguments)]
 async fn cmd_deploy(
     host: &str,
+    token: Option<&str>,
     compose_path: &str,
     name: &str,
     image: Option<&str>,
@@ -339,7 +368,7 @@ async fn cmd_deploy(
         ..Default::default()
     };
 
-    let vmm = Vmm::connect(host)?;
+    let vmm = Vmm::connect_with_token(host, token)?;
     let hash = vmm.get_compose_hash(&cfg).await?;
     let app_id = short(&hash, 40);
     cfg.app_id = Some(app_id.clone());
@@ -429,8 +458,8 @@ fn stub(name: &str) -> Result<()> {
     )
 }
 
-async fn cmd_apps(host: &str, json: bool) -> Result<()> {
-    let vmm = Vmm::connect(host)?;
+async fn cmd_apps(host: &str, token: Option<&str>, json: bool) -> Result<()> {
+    let vmm = Vmm::connect_with_token(host, token)?;
     let resp = vmm.status().await?;
     if json {
         let arr: Vec<_> = resp
@@ -470,8 +499,8 @@ async fn cmd_apps(host: &str, json: bool) -> Result<()> {
     Ok(())
 }
 
-async fn cmd_logs(host: &str, id: &str, lines: u32) -> Result<()> {
-    let vmm = Vmm::connect(host)?;
+async fn cmd_logs(host: &str, token: Option<&str>, id: &str, lines: u32) -> Result<()> {
+    let vmm = Vmm::connect_with_token(host, token)?;
     let logs = vmm.logs(id, lines).await?;
     print!("{logs}");
     Ok(())