Unify dstack agent address getter

Author: kvinwang

Cargo.lock

@@ -133,3 +133,12 @@ impl KeyProviderInfo {
 }
 
 pub mod shared_filenames;
+
+/// Get the address of the dstack agent
+pub fn dstack_agent_address() -> String {
+    // Check env DSTACK_AGENT_ADDRESS
+    if let Ok(address) = std::env::var("DSTACK_AGENT_ADDRESS") {
+        return address;
+    }
+    "unix:/var/run/dstack.sock".into()
+}

dstack-types/src/lib.rs

@@ -40,6 +40,7 @@ ra-tls.workspace = true
 dstack-guest-agent-rpc.workspace = true
 http-client = { workspace = true, features = ["prpc"] }
 sha2.workspace = true
+dstack-types.workspace = true
 
 [target.'cfg(unix)'.dependencies]
 nix = { workspace = true, features = ["resource"] }

gateway/Cargo.toml

@@ -1,4 +1,4 @@
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{anyhow, Context, Result};
 use clap::Parser;
 use config::{Config, TlsConfig};
 use dstack_guest_agent_rpc::{dstack_guest_client::DstackGuestClient, GetTlsKeyArgs};
@@ -49,11 +49,8 @@ fn set_max_ulimit() -> Result<()> {
 }
 
 fn dstack_agent() -> Result<DstackGuestClient<PrpcClient>> {
-    let uds_path = "/var/run/dstack.sock";
-    if !std::fs::exists(uds_path).unwrap_or_default() {
-        bail!("dstack socket({uds_path}) not found");
-    }
-    let http_client = PrpcClient::new_unix(uds_path.to_string(), "/prpc".to_string());
+    let address = dstack_types::dstack_agent_address();
+    let http_client = PrpcClient::new(address);
     Ok(DstackGuestClient::new(http_client))
 }
 

gateway/src/main.rs

@@ -33,8 +33,13 @@ pub async fn http_request(
 ) -> Result<(u16, Vec<u8>)> {
     debug!("Sending HTTP request to {base}, path={path}");
     let mut response = if base.starts_with("unix:") {
+        let path = if path.starts_with("/") {
+            path.to_string()
+        } else {
+            format!("/{path}")
+        };
         let client: Client<UnixConnector, Full<Bytes>> = Client::unix();
-        let unix_uri: hyper::Uri = Uri::new(base.strip_prefix("unix:").unwrap(), path).into();
+        let unix_uri: hyper::Uri = Uri::new(base.strip_prefix("unix:").unwrap(), &path).into();
         let req = Request::builder()
             .method(method)
             .uri(unix_uri)

http-client/src/lib.rs

@@ -37,6 +37,7 @@ x509-parser = { workspace = true, features = ["verify"] }
 ring.workspace = true
 safe-write.workspace = true
 serde_json.workspace = true
+dstack-types.workspace = true
 
 [features]
 default = []

kms/Cargo.toml

@@ -241,7 +241,8 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig) -> Result<()> {
 }
 
 fn dstack_client() -> DstackGuestClient<PrpcClient> {
-    let http_client = PrpcClient::new_unix("/var/run/dstack.sock".into(), "/prpc".into());
+    let address = dstack_types::dstack_agent_address();
+    let http_client = PrpcClient::new(address);
     DstackGuestClient::new(http_client)
 }