|
4 | 4 |
|
5 | 5 | use std::io; |
6 | 6 | use std::pin::Pin; |
7 | | -use std::sync::Arc; |
| 7 | +use std::sync::{Arc, OnceLock}; |
8 | 8 | use std::task::{Context, Poll}; |
9 | 9 |
|
10 | 10 | use anyhow::{anyhow, bail, Context as _, Result}; |
| 11 | +use dstack_guest_agent_rpc::v0::AppInfo; |
| 12 | +use dstack_types::AppCompose; |
11 | 13 | use hyper::body::Incoming; |
12 | 14 | use hyper::server::conn::http1; |
13 | 15 | use hyper::service::service_fn; |
14 | 16 | use hyper::{Request, Response, StatusCode}; |
15 | 17 | use hyper_util::rt::tokio::TokioIo; |
16 | 18 | use proxy_protocol::ProxyHeader; |
17 | 19 | use rustls::version::{TLS12, TLS13}; |
18 | | -use serde::Serialize; |
| 20 | +use serde::{Deserialize, Serialize}; |
19 | 21 | use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf}; |
20 | 22 | use tokio::net::TcpStream; |
21 | 23 | use tokio::time::timeout; |
@@ -206,7 +208,7 @@ impl Proxy { |
206 | 208 | "/app-info" => { |
207 | 209 | let agent = crate::dstack_agent().context("Failed to get dstack agent")?; |
208 | 210 | let app_info = agent.info().await.context("Failed to get app info")?; |
209 | | - json_response(&app_info) |
| 211 | + json_response(&public_app_info(app_info)) |
210 | 212 | } |
211 | 213 | "/acme-info" => { |
212 | 214 | let acme_info = self.acme_info(None).context("Failed to get acme info")?; |
@@ -618,6 +620,29 @@ impl AsyncWrite for MergedStream { |
618 | 620 | } |
619 | 621 | } |
620 | 622 |
|
| 623 | +/// The internal `Info` call returns the TCB info and VM config even when the |
| 624 | +/// app compose sets `public_tcbinfo = false`; the guest agent only withholds |
| 625 | +/// them on its external surface. `/app-info` is public, so withhold them here too. |
| 626 | +/// The choice is fixed for the CVM's lifetime, so it is decided once. |
| 627 | +fn public_app_info(mut info: AppInfo) -> AppInfo { |
| 628 | + static TCB_INFO_PUBLIC: OnceLock<bool> = OnceLock::new(); |
| 629 | + if !*TCB_INFO_PUBLIC.get_or_init(|| tcb_info_is_public(&info.tcb_info)) { |
| 630 | + info.tcb_info.clear(); |
| 631 | + info.vm_config.clear(); |
| 632 | + } |
| 633 | + info |
| 634 | +} |
| 635 | + |
| 636 | +fn tcb_info_is_public(tcb_info: &str) -> bool { |
| 637 | + #[derive(Deserialize)] |
| 638 | + struct TcbInfo { |
| 639 | + app_compose: String, |
| 640 | + } |
| 641 | + serde_json::from_str::<TcbInfo>(tcb_info) |
| 642 | + .and_then(|tcb| serde_json::from_str::<AppCompose>(&tcb.app_compose)) |
| 643 | + .is_ok_and(|compose| compose.public_tcbinfo) |
| 644 | +} |
| 645 | + |
621 | 646 | #[cfg(test)] |
622 | 647 | mod tests { |
623 | 648 | use super::*; |
@@ -659,4 +684,30 @@ mod tests { |
659 | 684 | let (remainder, _socket) = stream.into_socket_parts(); |
660 | 685 | assert!(remainder.is_empty(), "got {remainder:?}"); |
661 | 686 | } |
| 687 | + |
| 688 | + fn tcb_info_with_compose(extra: serde_json::Value) -> String { |
| 689 | + let mut compose = serde_json::json!({ |
| 690 | + "manifest_version": "3", |
| 691 | + "name": "gateway", |
| 692 | + "runner": "docker-compose", |
| 693 | + }); |
| 694 | + compose |
| 695 | + .as_object_mut() |
| 696 | + .expect("object") |
| 697 | + .extend(extra.as_object().expect("object").clone()); |
| 698 | + serde_json::json!({ "app_compose": compose.to_string() }).to_string() |
| 699 | + } |
| 700 | + |
| 701 | + #[test] |
| 702 | + fn tcb_info_is_private_when_the_compose_opts_out() { |
| 703 | + let tcb_info = tcb_info_with_compose(serde_json::json!({ "public_tcbinfo": false })); |
| 704 | + assert!(!tcb_info_is_public(&tcb_info)); |
| 705 | + } |
| 706 | + |
| 707 | + #[test] |
| 708 | + fn tcb_info_is_public_by_default() { |
| 709 | + assert!(tcb_info_is_public(&tcb_info_with_compose( |
| 710 | + serde_json::json!({}) |
| 711 | + ))); |
| 712 | + } |
662 | 713 | } |
0 commit comments