feat(http-client): support custom request headers and a prpc bearer token

Author: kvinwang

dstack/http-client/src/lib.rs

@@ -34,6 +34,17 @@ pub async fn http_request(
     base: &str,
     path: &str,
     body: &[u8],
+) -> Result<(u16, Vec<u8>)> {
+    http_request_with_headers(method, base, path, body, &[]).await
+}
+
+/// Same as [`http_request`], with extra request headers (e.g. `Authorization`).
+pub async fn http_request_with_headers(
+    method: &str,
+    base: &str,
+    path: &str,
+    body: &[u8],
+    headers: &[(&str, &str)],
 ) -> Result<(u16, Vec<u8>)> {
     debug!("Sending HTTP request to {base}, path={path}");
     let mut response = if let Some(uds) = base.strip_prefix("unix:") {
@@ -44,24 +55,29 @@ pub async fn http_request(
         };
         let client: Client<UnixConnector, Full<Bytes>> = Client::unix();
         let unix_uri: hyper::Uri = Uri::new(uds, &path).into();
-        let req = Request::builder()
-            .method(method)
-            .uri(unix_uri)
-            .body(Full::new(Bytes::copy_from_slice(body)))?;
+        let mut builder = Request::builder().method(method).uri(unix_uri);
+        for (name, value) in headers {
+            builder = builder.header(*name, *value);
+        }
+        let req = builder.body(Full::new(Bytes::copy_from_slice(body)))?;
         client.request(req).await?
     } else if base.starts_with("vsock:") {
         let client = Client::vsock();
         let uri = mk_url(base, path).parse::<hyper::Uri>()?;
-        let req = Request::builder()
-            .method(method)
-            .uri(uri)
-            .body(Full::new(Bytes::copy_from_slice(body)))?;
+        let mut builder = Request::builder().method(method).uri(uri);
+        for (name, value) in headers {
+            builder = builder.header(*name, *value);
+        }
+        let req = builder.body(Full::new(Bytes::copy_from_slice(body)))?;
         client.request(req).await?
     } else {
         let uri = mk_url(base, path);
         let client = reqwest::Client::builder().build()?;
         let method = reqwest::Method::from_bytes(method.as_bytes())?;
         let mut request = client.request(method, uri);
+        for (name, value) in headers {
+            request = request.header(*name, *value);
+        }
         if !body.is_empty() {
             request = request.body(body.to_vec());
         }

dstack/http-client/src/prpc.rs

@@ -12,13 +12,15 @@ use serde::{de::DeserializeOwned, Serialize};
 pub struct PrpcClient {
     base_url: String,
     path_append: String,
+    auth_token: Option<String>,
 }
 
 impl PrpcClient {
     pub fn new(base_url: String) -> Self {
         Self {
             base_url,
             path_append: String::new(),
+            auth_token: None,
         }
     }
 
@@ -29,8 +31,16 @@ impl PrpcClient {
         Self {
             base_url: format!("unix:{socket_path}"),
             path_append: path,
+            auth_token: None,
         }
     }
+
+    /// Send `Authorization: Bearer <token>` with every request.
+    pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
+        let token = token.into();
+        self.auth_token = (!token.is_empty()).then_some(token);
+        self
+    }
 }
 
 impl RequestClient for PrpcClient {
@@ -41,7 +51,15 @@ impl RequestClient for PrpcClient {
     {
         let body = serde_json::to_vec(&body).context("Failed to serialize body")?;
         let path = format!("{}{path}?json", self.path_append);
-        let (status, body) = super::http_request("POST", &self.base_url, &path, &body).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) =
+            super::http_request_with_headers("POST", &self.base_url, &path, &body, &headers)
+                .await?;
         if status != 200 {
             anyhow::bail!("Invalid status code: {status}, path={path}");
         }