Dstack-TEE · kvinwang · Jul 27, 2026 · Jul 27, 2026 · Jul 27, 2026 · Jul 27, 2026
diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock
diff --git a/dstack/supervisor/client/Cargo.toml b/dstack/supervisor/client/Cargo.toml
@@ -27,11 +27,15 @@ serde.workspace = true
 http-body-util.workspace = true
 tracing-subscriber.workspace = true
 log.workspace = true
+libc.workspace = true
 fs-err.workspace = true
 futures.workspace = true
 
 supervisor.workspace = true
 http-client.workspace = true
 
+[dev-dependencies]
+tempfile.workspace = true
+
 [features]
 cli = ["dep:clap", "tokio/full"]
diff --git a/dstack/supervisor/client/src/lib.rs b/dstack/supervisor/client/src/lib.rs
@@ -11,6 +11,58 @@ use supervisor::{ProcessConfig, ProcessInfo, Response};
 
 pub use supervisor;
 
+#[cfg(unix)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct SocketIdentity {
+    device: u64,
+    inode: u64,
+}
+
+#[cfg(unix)]
+fn trusted_uds_identity(path: &Path) -> Result<SocketIdentity> {
+    use std::os::unix::fs::{FileTypeExt as _, MetadataExt as _};
+
+    let metadata = fs_err::symlink_metadata(path)
+        .with_context(|| format!("Failed to inspect supervisor socket {}", path.display()))?;
+    if !metadata.file_type().is_socket() {
+        anyhow::bail!(
+            "Supervisor endpoint is not a Unix socket: {}",
+            path.display()
+        );
+    }
+    let effective_uid = unsafe { libc::geteuid() };
+    if metadata.uid() != effective_uid {
+        anyhow::bail!("Supervisor socket is not owned by the current user");
+    }
+    if metadata.mode() & 0o022 != 0 {
+        anyhow::bail!("Supervisor socket is writable by another user");
+    }
+
+    let parent = path.parent().unwrap_or_else(|| Path::new("."));
+    let parent_metadata = fs_err::symlink_metadata(parent).with_context(|| {
+        format!(
+            "Failed to inspect supervisor socket directory {}",
+            parent.display()
+        )
+    })?;
+    if !parent_metadata.file_type().is_dir() {
+        anyhow::bail!("Supervisor socket parent is not a directory");
+    }
+    let parent_owned = parent_metadata.uid() == effective_uid;
+    let parent_sticky = parent_metadata.mode() & 0o1000 != 0;
+    if !parent_owned && !parent_sticky {
+        anyhow::bail!("Supervisor socket directory is not controlled by the current user");
+    }
+    if parent_metadata.mode() & 0o022 != 0 && !parent_sticky {
+        anyhow::bail!("Supervisor socket directory permits untrusted replacement");
+    }
+
+    Ok(SocketIdentity {
+        device: metadata.dev(),
+        inode: metadata.ino(),
+    })
+}
+
 #[derive(Debug, Clone)]
 pub struct SupervisorClient {
     base_url: Arc<String>,
@@ -31,29 +83,38 @@ impl SupervisorClient {
         detached: bool,
         auto_start: bool,
     ) -> Result<Self> {
-        let uri = format!("unix:{}", uds.as_ref().display());
+        let uds = uds.as_ref();
+        let uri = format!("unix:{}", uds.display());
         let client = Self::new(&uri);
-        if client.probe(Duration::from_millis(100)).await.is_ok() {
-            info!("Connected to supervisor at {uri}");
-            return Ok(client);
+        if fs_err::symlink_metadata(uds).is_ok() {
+            let identity = trusted_uds_identity(uds)?;
+            if client.probe(Duration::from_millis(100)).await.is_ok()
+                && trusted_uds_identity(uds)? == identity
+            {
+                info!("Connected to supervisor at {uri}");
+                return Ok(client);
+            }
         }
         if !auto_start {
             anyhow::bail!("Failed to connect to supervisor at {uri}");
         }
         info!("Failed to connect to supervisor at {uri}, trying to start supervisor");
-        // if the uds exists, remove it
-        if std::path::Path::new(uds.as_ref()).exists() {
-            fs_err::remove_file(uds.as_ref())?;
+        if fs_err::symlink_metadata(uds).is_ok() {
+            // Validate again immediately before removing a stale endpoint. Never
+            // delete a path that is not a trusted socket owned by this user.
+            trusted_uds_identity(uds)?;
+            fs_err::remove_file(uds)?;
         }
         let supervisor_path = supervisor_path.as_ref().to_path_buf();
-        let uds = uds.as_ref().to_path_buf();
+        let uds = uds.to_path_buf();
+        let supervisor_uds = uds.clone();
         let pid_file = pid_file.as_ref().to_path_buf();
         let log_file = log_file.as_ref().to_path_buf();
         std::thread::spawn(move || {
             // start supervisor
             let result = std::process::Command::new(supervisor_path)
                 .arg("--uds")
-                .arg(uds)
+                .arg(supervisor_uds)
                 .arg("--pid-file")
                 .arg(pid_file)
                 .arg("--log-file")
@@ -77,9 +138,13 @@ impl SupervisorClient {
         });
         // wait while ping returns pong
         for i in 1..=10 {
-            if client.probe(Duration::from_millis(100)).await.is_ok() {
-                info!("connected to supervisor at {uri}");
-                return Ok(client);
+            if let Ok(identity) = trusted_uds_identity(&uds) {
+                if client.probe(Duration::from_millis(100)).await.is_ok()
+                    && trusted_uds_identity(&uds).ok() == Some(identity)
+                {
+                    info!("connected to supervisor at {uri}");
+                    return Ok(client);
+                }
             }
             info!("waiting for supervisor at {uri} to start, attempt {i}");
             tokio::time::sleep(Duration::from_millis(100 * i)).await;
@@ -231,3 +296,37 @@ impl SupervisorClientSync {
         }
     }
 }
+
+#[cfg(all(test, unix))]
+mod tests {
+    use super::*;
+    use std::os::unix::fs::PermissionsExt as _;
+    use std::os::unix::net::UnixListener;
+
+    #[test]
+    fn trusted_uds_rejects_regular_file() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("supervisor.sock");
+        fs_err::write(&path, b"not a socket").unwrap();
+        assert!(trusted_uds_identity(&path).is_err());
+    }
+
+    #[test]
+    fn trusted_uds_accepts_owner_only_socket() {
+        let directory = tempfile::tempdir().unwrap();
+        fs_err::set_permissions(directory.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
+        let path = directory.path().join("supervisor.sock");
+        let _listener = UnixListener::bind(&path).unwrap();
+        fs_err::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
+        trusted_uds_identity(&path).expect("owner-only socket should be trusted");
+    }
+
+    #[test]
+    fn trusted_uds_rejects_socket_writable_by_others() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("supervisor.sock");
+        let _listener = UnixListener::bind(&path).unwrap();
+        fs_err::set_permissions(&path, std::fs::Permissions::from_mode(0o666)).unwrap();
+        assert!(trusted_uds_identity(&path).is_err());
+    }
+}