diff --git a/Cargo.lock b/Cargo.lock index 69bb8f7c..aebb1dc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -135,6 +135,15 @@ version = "1.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +[[package]] +name = "arc-swap" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f3647c145568cec02c42054e07bdf9a5a698e15b466fb2341bfc393cd24aa5" +dependencies = [ + "rustversion", +] + [[package]] name = "arraydeque" version = "0.5.1" @@ -343,6 +352,26 @@ dependencies = [ "serde", ] +[[package]] +name = "bincode" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36eaf5d7b090263e8150820482d5d93cd964a81e4019913c972f4edcc6edb740" +dependencies = [ + "bincode_derive", + "serde", + "unty", +] + +[[package]] +name = "bincode_derive" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf95709a440f45e986983918d0e8a1f30a9b1df04918fc828670606804ac3c09" +dependencies = [ + "virtue", +] + [[package]] name = "bitflags" version = "1.3.2" @@ -599,9 +628,12 @@ version = "0.5.6" dependencies = [ "anyhow", "bon", + "bytes", "enum_dispatch", "fs-err", "hickory-resolver 0.24.4", + "http", + "http-body-util", "instant-acme", "path-absolutize", "rand 0.8.5", @@ -1457,6 +1489,8 @@ name = "dstack-gateway" version = "0.5.6" dependencies = [ "anyhow", + "arc-swap", + "base64 0.22.1", "bytes", "certbot", "clap", @@ -1465,13 +1499,16 @@ dependencies = [ "dstack-guest-agent-rpc", "dstack-kms-rpc", "dstack-types", + "flate2", "fs-err", "futures", "git-version", "hex", "hickory-resolver 0.24.4", + "http-body-util", "http-client", "hyper", + "hyper-rustls", "hyper-util", "insta", "ipnet", @@ -1486,6 +1523,7 @@ dependencies = [ "rand 0.8.5", "reqwest", "rinja", + "rmp-serde", "rocket", "rustls", "safe-write", @@ -1495,10 +1533,15 @@ dependencies = [ "sha2 0.10.9", "shared_child", "smallvec", + "tdx-attest", + "tempfile", "tokio", "tokio-rustls", "tracing", "tracing-subscriber", + "uuid", + "wavekv", + "x509-parser", ] [[package]] @@ -1672,6 +1715,17 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "dstack-port-forward" +version = "0.5.6" +dependencies = [ + "anyhow", + "nix 0.29.0", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "dstack-types" version = "0.5.6" @@ -1778,6 +1832,7 @@ dependencies = [ "clap", "dirs", "dstack-kms-rpc", + "dstack-port-forward", "dstack-types", "dstack-vmm-rpc", "fatfs", @@ -3447,7 +3502,7 @@ checksum = "2044d8bd5489b199890c3dbf38d4c8f50f3a5a38833986808b14e2367fe267fa" dependencies = [ "aes 0.7.5", "base64 0.13.1", - "bincode", + "bincode 1.3.3", "crossterm", "hmac 0.11.0", "pbkdf2", @@ -3712,6 +3767,7 @@ dependencies = [ "cfg-if", "cfg_aliases", "libc", + "memoffset 0.9.1", ] [[package]] @@ -6504,6 +6560,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -6850,6 +6907,12 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "unty" +version = "0.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae" + [[package]] name = "url" version = "2.5.7" @@ -6903,6 +6966,12 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "virtue" +version = "0.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" + [[package]] name = "void" version = "1.0.2" @@ -7026,6 +7095,29 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wavekv" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf9b73bc556dfdb7ef33617a9d477b803198db43ea3df25463efaf43d4986fe8" +dependencies = [ + "anyhow", + "bincode 2.0.1", + "chrono", + "crc32fast", + "dashmap", + "fs-err", + "futures", + "hex", + "rmp-serde", + "serde", + "serde-human-bytes", + "serde_json", + "sha2 0.10.9", + "tokio", + "tracing", +] + [[package]] name = "web-sys" version = "0.3.83" diff --git a/gateway/src/proxy/tls_passthough.rs b/gateway/src/proxy/tls_passthough.rs index e0939d6f..b334069d 100644 --- a/gateway/src/proxy/tls_passthough.rs +++ b/gateway/src/proxy/tls_passthough.rs @@ -2,6 +2,7 @@ // // SPDX-License-Identifier: BUSL-1.1 +use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; use std::sync::atomic::Ordering; @@ -34,8 +35,23 @@ impl AppAddress { } } +fn select_app_address(items: &[Box<[u8]>], known_apps: &BTreeMap>) -> Result { + let mut fallback = None; + for data in items { + if let Ok(addr) = AppAddress::parse(data) { + if known_apps.contains_key(&addr.app_id) { + return Ok(addr); + } + if fallback.is_none() { + fallback = Some(addr); + } + } + } + fallback.context("no app address found in txt record") +} + /// resolve app address by sni -async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result { +async fn resolve_app_address(prefix: &str, sni: &str, compat: bool, state: &Proxy) -> Result { let txt_domain = format!("{prefix}.{sni}"); let resolver = hickory_resolver::AsyncResolver::tokio_from_system_conf() .context("failed to create dns resolver")?; @@ -65,11 +81,8 @@ async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result Box<[u8]> { + s.to_vec().into_boxed_slice() + } + + async fn create_test_proxy() -> (Proxy, TempDir) { + let figment = load_config_figment(None); + let mut config = figment.focus("core").extract::().unwrap(); + let temp_dir = TempDir::new().expect("failed to create temp dir"); + config.sync.data_dir = temp_dir.path().to_string_lossy().to_string(); + let proxy = Proxy::new(ProxyOptions { + config, + my_app_id: None, + tls_config: TlsConfig { + certs: "".to_string(), + key: "".to_string(), + mutual: MutualConfig { ca_certs: "".to_string() }, + }, + }) + .await + .expect("failed to create proxy"); + (proxy, temp_dir) + } #[tokio::test] async fn test_resolve_app_address() { + let (state, _dir) = create_test_proxy().await; let app_addr = resolve_app_address( "_dstack-app-address", "3327603e03f5bd1f830812ca4a789277fc31f577.app.dstack.org", false, + &state, ) .await .unwrap(); assert_eq!(app_addr.app_id, "3327603e03f5bd1f830812ca4a789277fc31f577"); assert_eq!(app_addr.port, 8090); } + + #[test] + fn test_select_app_address_prefers_local() { + let items = vec![boxed(b"aaaaaa:443"), boxed(b"bbbbbb:8080")]; + let mut apps = BTreeMap::new(); + apps.insert("bbbbbb".to_string(), BTreeSet::new()); + let addr = select_app_address(&items, &apps).unwrap(); + assert_eq!(addr.app_id, "bbbbbb"); + assert_eq!(addr.port, 8080); + } + + #[test] + fn test_select_app_address_fallback_when_none_local() { + let items = vec![boxed(b"aaaaaa:443"), boxed(b"bbbbbb:8080")]; + let addr = select_app_address(&items, &BTreeMap::new()).unwrap(); + assert_eq!(addr.app_id, "aaaaaa"); + assert_eq!(addr.port, 443); + } }