Merge pull request #230 from Dstack-TEE/gw-quote

gateway: Add api for evidences

Author: shelvenzhou

certbot/src/bot.rs

@@ -210,7 +210,7 @@ impl CertBot {
     }
 }
 
-fn read_pubkey(cert_pem: &str) -> Result<Vec<u8>> {
+pub fn read_pubkey(cert_pem: &str) -> Result<Vec<u8>> {
     let cert = read_pem(cert_pem)?;
     let public_key = cert.parse_x509().context("failed to parse x509 cert")?;
     Ok(public_key.tbs_certificate.public_key().raw.to_vec())

certbot/src/lib.rs

@@ -13,7 +13,7 @@
 //! to the documentation of individual structs and functions.
 
 pub use acme_client::AcmeClient;
-pub use bot::{CertBot, CertBotConfig};
+pub use bot::{read_pubkey, CertBot, CertBotConfig};
 pub use dns01_client::Dns01Client;
 pub use workdir::WorkDir;
 

certbot/src/workdir.rs

@@ -53,6 +53,10 @@ impl WorkDir {
         Ok(credentials.account_id)
     }
 
+    pub fn acme_account_quote_path(&self) -> PathBuf {
+        self.workdir.join("acme-account.quote")
+    }
+
     pub fn list_cert_public_keys(&self) -> Result<BTreeSet<Vec<u8>>> {
         crate::bot::list_cert_public_keys(self.backup_dir())
     }

gateway/rpc/proto/gateway_rpc.proto

@@ -25,7 +25,7 @@ message WireGuardPeer {
   string ip = 2;
   // The wireguard peer endpoint.
   string endpoint = 3;
-} 
+}
 
 // WireGuardConfig is the configuration of the WireGuard.
 message WireGuardConfig {
@@ -79,12 +79,25 @@ message HostInfo {
   uint64 num_connections = 7;
 }
 
+message QuotedPublicKey {
+  bytes public_key = 1;
+  string quote = 2;
+}
+
 // AcmeInfoResponse is the response for AcmeInfo.
 message AcmeInfoResponse {
   // The ACME account URI.
   string account_uri = 1;
   // The public key history of the certificate.
   repeated bytes hist_keys = 2;
+  // The quoted public key of the certificate.
+  repeated QuotedPublicKey quoted_hist_keys = 3;
+  // The quote of the ACME account URI.
+  string account_quote = 4;
+  // Active certificate
+  string active_cert = 5;
+  // The domain that serves ZT-HTTPS
+  string base_domain = 6;
 }
 
 // Get HostInfo for associated instance id.

gateway/src/main_service.rs

@@ -2,6 +2,7 @@ use std::{
     collections::{BTreeMap, BTreeSet},
     net::Ipv4Addr,
     ops::Deref,
+    path::Path,
     sync::{Arc, Mutex, MutexGuard, RwLock},
     time::{Duration, Instant, SystemTime, UNIX_EPOCH},
 };
@@ -12,11 +13,14 @@ use certbot::{CertBot, WorkDir};
 use cmd_lib::run_cmd as cmd;
 use dstack_gateway_rpc::{
     gateway_server::{GatewayRpc, GatewayServer},
-    AcmeInfoResponse, GatewayState, GuestAgentConfig, RegisterCvmRequest, RegisterCvmResponse,
-    WireGuardConfig, WireGuardPeer,
+    AcmeInfoResponse, GatewayState, GuestAgentConfig, QuotedPublicKey, RegisterCvmRequest,
+    RegisterCvmResponse, WireGuardConfig, WireGuardPeer,
 };
+use dstack_guest_agent_rpc::{dstack_guest_client::DstackGuestClient, RawQuoteArgs};
 use fs_err as fs;
+use http_client::prpc::PrpcClient;
 use ra_rpc::{CallContext, RpcCall, VerifiedAttestation};
+use ra_tls::attestation::QuoteContentType;
 use rand::seq::IteratorRandom;
 use rinja::Template as _;
 use safe_write::safe_write;
@@ -170,6 +174,51 @@ impl Proxy {
         }
         Ok(renewed)
     }
+
+    pub(crate) async fn acme_info(&self) -> Result<AcmeInfoResponse> {
+        let config = self.lock().config.clone();
+        let workdir = WorkDir::new(&config.certbot.workdir);
+        let account_uri = workdir.acme_account_uri().unwrap_or_default();
+        let keys = workdir.list_cert_public_keys().unwrap_or_default();
+        let agent = crate::dstack_agent().context("Failed to get dstack agent")?;
+        let account_quote = get_or_generate_quote(
+            &agent,
+            QuoteContentType::Custom("acme-account"),
+            account_uri.as_bytes(),
+            workdir.acme_account_quote_path(),
+        )
+        .await
+        .unwrap_or_default();
+
+        let mut quoted_hist_keys = vec![];
+        for cert_path in workdir.list_certs().unwrap_or_default() {
+            let cert_pem = fs::read_to_string(&cert_path).context("Failed to read key")?;
+            let pubkey = certbot::read_pubkey(&cert_pem).context("Failed to read pubkey")?;
+            let quote = get_or_generate_quote(
+                &agent,
+                QuoteContentType::Custom("zt-cert"),
+                &pubkey,
+                cert_path.display().to_string() + ".quote",
+            )
+            .await
+            .unwrap_or_default();
+            quoted_hist_keys.push(QuotedPublicKey {
+                public_key: pubkey,
+                quote,
+            });
+        }
+        let active_cert =
+            fs::read_to_string(workdir.cert_path()).context("Failed to read active cert")?;
+
+        Ok(AcmeInfoResponse {
+            account_uri,
+            hist_keys: keys.into_iter().collect(),
+            account_quote,
+            quoted_hist_keys,
+            active_cert,
+            base_domain: config.proxy.base_domain.clone(),
+        })
+    }
 }
 
 fn load_state(state_path: &str) -> Result<ProxyStateMut> {
@@ -681,14 +730,7 @@ impl GatewayRpc for RpcHandler {
     }
 
     async fn acme_info(self) -> Result<AcmeInfoResponse> {
-        let state = self.state.lock();
-        let workdir = WorkDir::new(&state.config.certbot.workdir);
-        let account_uri = workdir.acme_account_uri().unwrap_or_default();
-        let keys = workdir.list_cert_public_keys().unwrap_or_default();
-        Ok(AcmeInfoResponse {
-            account_uri,
-            hist_keys: keys.into_iter().collect(),
-        })
+        self.state.acme_info().await
     }
 
     async fn update_state(self, request: GatewayState) -> Result<()> {
@@ -725,6 +767,26 @@ impl GatewayRpc for RpcHandler {
     }
 }
 
+async fn get_or_generate_quote(
+    agent: &DstackGuestClient<PrpcClient>,
+    content_type: QuoteContentType<'_>,
+    payload: &[u8],
+    quote_path: impl AsRef<Path>,
+) -> Result<String> {
+    let quote_path = quote_path.as_ref();
+    if fs::metadata(quote_path).is_ok() {
+        return fs::read_to_string(quote_path).context("Failed to read quote");
+    }
+    let report_data = content_type.to_report_data(payload).to_vec();
+    let response = agent
+        .get_quote(RawQuoteArgs { report_data })
+        .await
+        .context("Failed to get quote")?;
+    let quote = serde_json::to_string(&response).context("Failed to serialize quote")?;
+    safe_write(quote_path, &quote).context("Failed to write quote")?;
+    Ok(quote)
+}
+
 impl RpcCall<Proxy> for RpcHandler {
     type PrpcService = GatewayServer<Self>;
 

gateway/src/proxy/tls_terminate.rs

@@ -13,6 +13,7 @@ use hyper_util::rt::tokio::TokioIo;
 use rustls::pki_types::pem::PemObject;
 use rustls::pki_types::{CertificateDer, PrivateKeyDer};
 use rustls::version::{TLS12, TLS13};
+use serde::Serialize;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::net::TcpStream;
 use tokio::time::timeout;
@@ -124,6 +125,22 @@ pub(crate) fn create_acceptor(config: &ProxyConfig) -> Result<TlsAcceptor> {
     Ok(acceptor)
 }
 
+fn json_response(body: &impl Serialize) -> Result<Response<String>> {
+    let body = serde_json::to_string(body).context("Failed to serialize response")?;
+    Response::builder()
+        .status(StatusCode::OK)
+        .header("Content-Type", "application/json")
+        .body(body)
+        .context("Failed to build response")
+}
+
+fn empty_response(status: StatusCode) -> Result<Response<String>> {
+    Response::builder()
+        .status(status)
+        .body(String::new())
+        .context("Failed to build response")
+}
+
 impl Proxy {
     /// Reload the TLS acceptor with fresh certificates
     pub fn reload_certificates(&self) -> Result<()> {
@@ -141,6 +158,61 @@ impl Proxy {
         Ok(())
     }
 
+    pub(crate) async fn handle_this_node(
+        &self,
+        inbound: TcpStream,
+        buffer: Vec<u8>,
+        port: u16,
+    ) -> Result<()> {
+        if port != 80 {
+            bail!("Only port 80 is supported for this node");
+        }
+        let stream = self.tls_accept(inbound, buffer).await?;
+        let io = TokioIo::new(stream);
+
+        let service = service_fn(|req: Request<Incoming>| async move {
+            // Only respond to GET / requests
+            if req.method() != hyper::Method::GET {
+                return empty_response(StatusCode::METHOD_NOT_ALLOWED);
+            }
+            if req.uri().path() == "/health" {
+                return empty_response(StatusCode::OK);
+            }
+            let path = req.uri().path().trim_start_matches("/.dstack");
+            match path {
+                "/index" => {
+                    let body = serde_json::json!({
+                        "type": "dstack gateway",
+                        "paths": [
+                            "/index",
+                            "/app-info",
+                            "/acme-info",
+                        ],
+                    });
+                    json_response(&body)
+                }
+                "/app-info" => {
+                    let agent = crate::dstack_agent().context("Failed to get dstack agent")?;
+                    let app_info = agent.info().await.context("Failed to get app info")?;
+                    json_response(&app_info)
+                }
+                "/acme-info" => {
+                    let acme_info = self.acme_info().await.context("Failed to get acme info")?;
+                    json_response(&acme_info)
+                }
+                _ => empty_response(StatusCode::NOT_FOUND),
+            }
+        });
+
+        http1::Builder::new()
+            .serve_connection(io, service)
+            .await
+            .context("Failed to serve HTTP connection")?;
+
+        Ok(())
+    }
+
+    /// Deprecated legacy endpoint
     pub(crate) async fn handle_health_check(
         &self,
         inbound: TcpStream,
@@ -158,17 +230,15 @@ impl Proxy {
         let service = service_fn(|req: Request<Incoming>| async move {
             // Only respond to GET / requests
             if req.method() != hyper::Method::GET || req.uri().path() != "/" {
-                return Ok::<_, anyhow::Error>(
-                    Response::builder()
-                        .status(StatusCode::METHOD_NOT_ALLOWED)
-                        .body(String::new())
-                        .unwrap(),
-                );
+                return Response::builder()
+                    .status(StatusCode::NOT_FOUND)
+                    .body(String::new())
+                    .context("Failed to build response");
             }
-            Ok(Response::builder()
+            Response::builder()
                 .status(StatusCode::OK)
                 .body(String::new())
-                .unwrap())
+                .context("Failed to build response")
         });
 
         http1::Builder::new()
@@ -214,6 +284,9 @@ impl Proxy {
         if app_id == "health" {
             return self.handle_health_check(inbound, buffer, port).await;
         }
+        if app_id == "gateway" {
+            return self.handle_this_node(inbound, buffer, port).await;
+        }
         let addresses = self
             .lock()
             .select_top_n_hosts(app_id)

guest-agent/dstack.toml

@@ -16,6 +16,7 @@ data_disks = ["/"]
 enabled = false
 quote_file = "quote.hex"
 event_log_file = "eventlog.json"
+sys_config_file = "sys-config.json"
 
 [internal-v0]
 address = "unix:/var/run/tappd.sock"

guest-agent/rpc/proto/agent_rpc.proto

@@ -162,6 +162,8 @@ message GetQuoteResponse {
   string event_log = 2;
   // Report data
   bytes report_data = 3;
+  // Hw config
+  string vm_config = 4;
 }
 
 message EmitEventArgs {

guest-agent/src/rpc_service.rs

@@ -212,17 +212,23 @@ impl DstackGuestRpc for InternalRpcHandler {
         }
         let report_data = pad64(&request.report_data).context("Report data is too long")?;
         if self.state.config().simulator.enabled {
-            return simulate_quote(self.state.config(), report_data);
+            return simulate_quote(
+                self.state.config(),
+                report_data,
+                &self.state.inner.vm_config,
+            );
         }
         let (_, quote) =
             tdx_attest::get_quote(&report_data, None).context("Failed to get quote")?;
         let event_log = read_event_logs().context("Failed to decode event log")?;
         let event_log =
             serde_json::to_string(&event_log).context("Failed to serialize event log")?;
+
         Ok(GetQuoteResponse {
             quote,
             event_log,
             report_data: report_data.to_vec(),
+            vm_config: self.state.inner.vm_config.clone(),
         })
     }
 
@@ -238,7 +244,11 @@ impl DstackGuestRpc for InternalRpcHandler {
     }
 }
 
-fn simulate_quote(config: &Config, report_data: [u8; 64]) -> Result<GetQuoteResponse> {
+fn simulate_quote(
+    config: &Config,
+    report_data: [u8; 64],
+    vm_config: &str,
+) -> Result<GetQuoteResponse> {
     let quote_file =
         fs::read_to_string(&config.simulator.quote_file).context("Failed to read quote file")?;
     let mut quote = hex::decode(quote_file.trim()).context("Failed to decode quote")?;
@@ -252,6 +262,7 @@ fn simulate_quote(config: &Config, report_data: [u8; 64]) -> Result<GetQuoteResp
         quote,
         event_log,
         report_data: report_data.to_vec(),
+        vm_config: vm_config.to_string(),
     })
 }
 
@@ -332,7 +343,11 @@ impl TappdRpc for InternalRpcHandlerV0 {
         let report_data =
             content_type.to_report_data_with_hash(&request.report_data, &request.hash_algorithm)?;
         if self.state.config().simulator.enabled {
-            let response = simulate_quote(self.state.config(), report_data)?;
+            let response = simulate_quote(
+                self.state.config(),
+                report_data,
+                &self.state.inner.vm_config,
+            )?;
             return Ok(TdxQuoteResponse {
                 quote: response.quote,
                 event_log: response.event_log,