Merge pull request #627 from Dstack-TEE/feat/remoteendpoint-unix-peer

ra-rpc: populate Unix peer creds for UDS endpoints

Author: kvinwang

guest-agent/src/server.rs

@@ -10,10 +10,11 @@ use crate::http_routes;
 use crate::rpc_service::{AppState, ExternalRpcHandler, InternalRpcHandler, InternalRpcHandlerV0};
 use crate::socket_activation::{ActivatedSockets, ActivatedUnixListener};
 use anyhow::{anyhow, Context, Result};
+use ra_rpc::rocket_helper::UnixPeerCredListener;
 use rocket::{
     fairing::AdHoc,
     figment::Figment,
-    listener::{Bind, DefaultListener},
+    listener::{unix::UnixListener, Bind, DefaultListener, Endpoint},
 };
 use rocket_vsock_listener::VsockListener;
 use sd_notify::{notify as sd_notify, NotifyState};
@@ -43,7 +44,7 @@ async fn run_internal_v0(
 
     if let Some(std_listener) = activated_socket {
         info!("Using systemd-activated socket for tappd.sock");
-        let listener = ActivatedUnixListener::new(std_listener)?;
+        let listener = UnixPeerCredListener::new(ActivatedUnixListener::new(std_listener)?);
         sock_ready_tx.send(()).ok();
         ignite
             .launch_on(listener)
@@ -52,14 +53,29 @@ async fn run_internal_v0(
     } else {
         let endpoint = DefaultListener::bind_endpoint(&ignite)
             .map_err(|err| anyhow!("Failed to get endpoint: {err}"))?;
-        let listener = DefaultListener::bind(&ignite)
-            .await
-            .map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
         sock_ready_tx.send(()).ok();
-        ignite
-            .launch_on(listener)
-            .await
-            .map_err(|err| anyhow!(err.to_string()))?;
+        match endpoint {
+            Endpoint::Unix(_) => {
+                let listener = UnixPeerCredListener::new(
+                    <UnixListener as Bind>::bind(&ignite)
+                        .await
+                        .map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?,
+                );
+                ignite
+                    .launch_on(listener)
+                    .await
+                    .map_err(|err| anyhow!(err.to_string()))?;
+            }
+            _ => {
+                let listener = DefaultListener::bind(&ignite)
+                    .await
+                    .map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
+                ignite
+                    .launch_on(listener)
+                    .await
+                    .map_err(|err| anyhow!(err.to_string()))?;
+            }
+        }
     }
     Ok(())
 }
@@ -80,7 +96,7 @@ async fn run_internal(
 
     if let Some(std_listener) = activated_socket {
         info!("Using systemd-activated socket for dstack.sock");
-        let listener = ActivatedUnixListener::new(std_listener)?;
+        let listener = UnixPeerCredListener::new(ActivatedUnixListener::new(std_listener)?);
         sock_ready_tx.send(()).ok();
         ignite
             .launch_on(listener)
@@ -89,14 +105,29 @@ async fn run_internal(
     } else {
         let endpoint = DefaultListener::bind_endpoint(&ignite)
             .map_err(|err| anyhow!("Failed to get endpoint: {err}"))?;
-        let listener = DefaultListener::bind(&ignite)
-            .await
-            .map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
         sock_ready_tx.send(()).ok();
-        ignite
-            .launch_on(listener)
-            .await
-            .map_err(|err| anyhow!(err.to_string()))?;
+        match endpoint {
+            Endpoint::Unix(_) => {
+                let listener = UnixPeerCredListener::new(
+                    <UnixListener as Bind>::bind(&ignite)
+                        .await
+                        .map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?,
+                );
+                ignite
+                    .launch_on(listener)
+                    .await
+                    .map_err(|err| anyhow!(err.to_string()))?;
+            }
+            _ => {
+                let listener = DefaultListener::bind(&ignite)
+                    .await
+                    .map_err(|err| anyhow!("Failed to bind on {endpoint}: {err}"))?;
+                ignite
+                    .launch_on(listener)
+                    .await
+                    .map_err(|err| anyhow!(err.to_string()))?;
+            }
+        }
     }
     Ok(())
 }

ra-rpc/src/lib.rs

@@ -22,12 +22,31 @@ pub mod client;
 #[cfg(feature = "openapi")]
 pub mod openapi;
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct UnixPeerCred {
+    /// Peer process ID (platform-independent representation)
+    pub pid: u64,
+    /// Peer user ID
+    pub uid: u64,
+    /// Peer group ID
+    pub gid: u64,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum RemoteEndpoint {
     Tcp(SocketAddr),
     Quic(SocketAddr),
-    Unix(PathBuf),
-    Vsock { cid: u32, port: u32 },
+    /// Unix domain socket endpoint.
+    ///
+    /// When available, `peer` can carry SO_PEERCRED (pid/uid/gid) of the caller.
+    Unix {
+        path: PathBuf,
+        peer: Option<UnixPeerCred>,
+    },
+    Vsock {
+        cid: u32,
+        port: u32,
+    },
     Other(String),
 }