Merge pull request #180 from Dstack-TEE/gw-app-auth

gw: Add authority to run

Author: kvinwang

.gitignore

@@ -5,3 +5,4 @@
 generated/
 node_modules/
 /.cargo
+.venv

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

@@ -1,4 +1,5 @@
 /.app-compose.json
+/.prelaunch.sh
 /.env
 /.venv
 /.app_env

gateway/dstack-app/.gitignore

@@ -1,5 +1,28 @@
 #!/bin/bash
 
+APP_COMPOSE_FILE=""
+
+usage() {
+  echo "Usage: $0 [-c <app compose file>]"
+  echo "  -c  App compose file"
+}
+
+while getopts "c:h" opt; do
+  case $opt in
+    c)
+      APP_COMPOSE_FILE=$OPTARG
+      ;;
+    h)
+      usage
+      exit 0
+      ;;
+    \?)
+      usage
+      exit 1
+      ;;
+  esac
+done
+
 # Check if .env exists
 if [ -f ".env" ]; then
   # Load variables from .env
@@ -57,6 +80,9 @@ GATEWAY_SERVING_ADDR=0.0.0.0:9204
 GUEST_AGENT_ADDR=127.0.0.1:9206
 WG_ADDR=0.0.0.0:9202
 
+# The token used to launch the App
+APP_LAUNCH_TOKEN=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 32 | head -n 1)
+
 EOF
   echo "Please edit the .env file and set the required variables, then run this script again."
   exit 1
@@ -72,6 +98,7 @@ required_env_vars=(
   "WG_ADDR"
   "GATEWAY_APP_ID"
   "MY_URL"
+  "APP_LAUNCH_TOKEN"
   # "BOOTNODE_URL"
 )
 
@@ -112,17 +139,36 @@ WG_ENDPOINT=$PUBLIC_IP:$WG_PORT
 MY_URL=$MY_URL
 BOOTNODE_URL=$BOOTNODE_URL
 SUBNET_INDEX=$SUBNET_INDEX
+APP_LAUNCH_TOKEN=$APP_LAUNCH_TOKEN
 EOF
 
-$CLI compose \
-  --docker-compose "$COMPOSE_TMP" \
-  --name dstack-gateway \
-  --kms \
-  --env-file .app_env \
-  --public-logs \
-  --public-sysinfo \
-  --no-instance-id \
-  --output .app-compose.json
+if [ -n "$APP_COMPOSE_FILE" ]; then
+  cp "$APP_COMPOSE_FILE" .app-compose.json
+else
+
+  EXPECTED_TOKEN_HASH=$(echo -n "$APP_LAUNCH_TOKEN" | sha256sum | cut -d' ' -f1)
+  cat >.prelaunch.sh <<EOF
+ACTUAL_TOKEN_HASH=\$(echo -n "\$APP_LAUNCH_TOKEN" | sha256sum | cut -d' ' -f1)
+if [ "$EXPECTED_TOKEN_HASH" != "\$ACTUAL_TOKEN_HASH" ]; then
+    echo "Error: Incorrect APP_LAUNCH_TOKEN, please make sure set the correct APP_LAUNCH_TOKEN in env"
+    reboot
+    exit 1
+else
+    echo "APP_LAUNCH_TOKEN checked OK"
+fi
+EOF
+
+  $CLI compose \
+    --docker-compose "$COMPOSE_TMP" \
+    --name dstack-gateway \
+    --kms \
+    --env-file .app_env \
+    --public-logs \
+    --public-sysinfo \
+    --no-instance-id \
+    --prelaunch-script .prelaunch.sh \
+    --output .app-compose.json
+fi
 
 # Remove the temporary file as it is no longer needed
 rm "$COMPOSE_TMP"

gateway/dstack-app/deploy-to-vmm.sh

@@ -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(())
+    }
+}

gateway/src/main_service/auth_client.rs

@@ -426,7 +426,7 @@ def create_app_compose(self,
             "secure_time": True,
         }
         if prelaunch_script:
-            app_compose["prelaunch_script"] = prelaunch_script
+            app_compose["pre_launch_script"] = open(prelaunch_script, 'rb').read().decode('utf-8')
 
         compose_file = json.dumps(app_compose, indent=4).encode('utf-8')
         compose_hash = hashlib.sha256(compose_file).hexdigest()