gw: Add auth API

Author: kvinwang

Cargo.lock

@@ -42,6 +42,7 @@ http-client = { workspace = true, features = ["prpc"] }
 sha2.workspace = true
 dstack-types.workspace = true
 serde-duration.workspace = true
+reqwest = { workspace = true, features = ["json"] }
 
 [target.'cfg(unix)'.dependencies]
 nix = { workspace = true, features = ["resource"] }

gateway/Cargo.toml

@@ -14,6 +14,11 @@ set_ulimit = true
 rpc_domain = ""
 run_in_dstack = true
 
+[core.auth]
+enabled = false
+url = "http://localhost/app-auth"
+timeout = "5s"
+
 [core.admin]
 enabled = false
 port = 8011

gateway/gateway.toml

@@ -121,6 +121,15 @@ pub struct Config {
     pub admin: AdminConfig,
     pub run_in_dstack: bool,
     pub sync: SyncConfig,
+    pub auth: AuthConfig,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct AuthConfig {
+    pub enabled: bool,
+    pub url: String,
+    #[serde(with = "serde_duration")]
+    pub timeout: Duration,
 }
 
 impl Config {

gateway/src/config.rs

@@ -6,6 +6,7 @@ use std::{
 };
 
 use anyhow::{bail, Context, Result};
+use auth_client::AuthClient;
 use certbot::{CertBot, WorkDir};
 use cmd_lib::run_cmd as cmd;
 use dstack_gateway_rpc::{
@@ -33,13 +34,16 @@ use crate::{
 
 mod sync_client;
 
+mod auth_client;
+
 #[derive(Clone)]
 pub struct Proxy {
     pub(crate) config: Arc<Config>,
     pub(crate) certbot: Option<Arc<CertBot>>,
     my_app_id: Option<Vec<u8>>,
     sync_tx: Sender<SyncEvent>,
     inner: Arc<Mutex<ProxyState>>,
+    auth_client: Arc<AuthClient>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -102,12 +106,14 @@ impl Proxy {
         let (sync_tx, sync_rx) = mpsc::channel(1);
         start_sync_task(Arc::downgrade(&inner), config.clone(), sync_rx);
         let certbot = start_certbot_task(&config).await?;
+        let auth_client = Arc::new(AuthClient::new(config.auth.clone()));
         Ok(Self {
             config,
             inner,
             certbot,
             my_app_id,
             sync_tx,
+            auth_client,
         })
     }
 }
@@ -607,6 +613,11 @@ impl GatewayRpc for RpcHandler {
         let app_info = ra
             .decode_app_info(false)
             .context("failed to decode app-info from attestation")?;
+        self.state
+            .auth_client
+            .ensure_app_authorized(&app_info)
+            .await
+            .context("App authorization failed")?;
         let app_id = hex::encode(&app_info.app_id);
         let instance_id = hex::encode(&app_info.instance_id);
 

gateway/src/main_service.rs

@@ -0,0 +1,31 @@
+use crate::config::AuthConfig;
+use anyhow::{Context, Result};
+use ra_tls::attestation::AppInfo;
+use reqwest::Client;
+
+pub(crate) struct AuthClient {
+    config: AuthConfig,
+    client: Client,
+}
+
+impl AuthClient {
+    pub(crate) fn new(config: AuthConfig) -> Self {
+        Self {
+            config,
+            client: reqwest::Client::new(),
+        }
+    }
+
+    pub(crate) async fn ensure_app_authorized(&self, app_info: &AppInfo) -> Result<()> {
+        if !self.config.enabled {
+            return Ok(());
+        }
+        let req = self.client.post(&self.config.url).json(app_info).send();
+        let res = tokio::time::timeout(self.config.timeout, req)
+            .await
+            .context("Auth timeout")?
+            .context("Failed to send request")?;
+        res.error_for_status().context("Request failed")?;
+        Ok(())
+    }
+}