diff --git a/Cargo.lock b/Cargo.lock index 4640a34f..17c9eb51 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5086,8 +5086,12 @@ checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" dependencies = [ "futures-util", "log", + "rustls", + "rustls-pki-types", "tokio", + "tokio-rustls", "tungstenite 0.29.0", + "webpki-roots 0.26.11", ] [[package]] @@ -5395,6 +5399,8 @@ dependencies = [ "httparse", "log", "rand 0.9.4", + "rustls", + "rustls-pki-types", "sha1", "thiserror 2.0.18", ] @@ -6133,6 +6139,24 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.9", +] + +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "weezl" version = "0.1.12" diff --git a/crates/workspace-server/Cargo.toml b/crates/workspace-server/Cargo.toml index 2b6dfff3..e4e58940 100644 --- a/crates/workspace-server/Cargo.toml +++ b/crates/workspace-server/Cargo.toml @@ -38,7 +38,7 @@ memory.workspace = true merge-request.workspace = true tokio = { workspace = true, features = ["fs", "macros", "net", "rt-multi-thread", "sync", "time"] } tower.workspace = true -tokio-tungstenite.workspace = true +tokio-tungstenite = { workspace = true, features = ["rustls-tls-webpki-roots"] } worker.workspace = true workspace-api.workspace = true workdir = { workspace = true, features = ["http-client"] } diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 6e6cdcbd..b31796c8 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -2984,7 +2984,33 @@ impl WorkdirHttpAuthorization for RemoteWorkdirAuthorization { } } -fn resolve_strict_remote_runtime_endpoint( +fn resolve_remote_addresses_with_timeout( + timeout: Duration, + resolver: F, +) -> Result, String> +where + F: FnOnce() -> Result, String> + Send + 'static, +{ + let (sender, receiver) = std::sync::mpsc::sync_channel(1); + std::thread::Builder::new() + .name("runtime-egress-dns".to_string()) + .spawn(move || { + let _ = sender.send(resolver()); + }) + .map_err(|_| "endpoint DNS resolver could not start".to_string())?; + receiver + .recv_timeout(timeout) + .map_err(|error| match error { + std::sync::mpsc::RecvTimeoutError::Timeout => { + "endpoint DNS resolution timed out".to_string() + } + std::sync::mpsc::RecvTimeoutError::Disconnected => { + "endpoint DNS resolver stopped unexpectedly".to_string() + } + })? +} + +pub(crate) fn resolve_strict_remote_runtime_endpoint( endpoint: &str, ) -> Result<(String, Vec), String> { let endpoint = @@ -3010,10 +3036,13 @@ fn resolve_strict_remote_runtime_endpoint( return Err("endpoint host is not public".to_string()); } let port = endpoint.port_or_known_default().unwrap_or(443); - let addresses = (host.as_str(), port) - .to_socket_addrs() - .map_err(|_| "endpoint DNS resolution failed".to_string())? - .collect::>(); + let resolution_host = host.clone(); + let addresses = resolve_remote_addresses_with_timeout(Duration::from_secs(3), move || { + (resolution_host.as_str(), port) + .to_socket_addrs() + .map(|addresses| addresses.collect::>()) + .map_err(|_| "endpoint DNS resolution failed".to_string()) + })?; if addresses.is_empty() || addresses .iter() @@ -4875,6 +4904,16 @@ mod tests { use std::sync::{Arc, Mutex}; use std::thread; + #[test] + fn strict_remote_runtime_dns_resolution_has_a_short_timeout() { + let error = resolve_remote_addresses_with_timeout(Duration::from_millis(1), || { + std::thread::sleep(Duration::from_millis(50)); + Ok(Vec::new()) + }) + .unwrap_err(); + assert_eq!(error, "endpoint DNS resolution timed out"); + } + #[test] fn strict_remote_runtime_egress_rejects_disallowed_endpoint_before_client_use() { let config = RemoteRuntimeConfig::new( diff --git a/crates/workspace-server/src/runtime_subscription.rs b/crates/workspace-server/src/runtime_subscription.rs index 88a3a47e..0397d7f4 100644 --- a/crates/workspace-server/src/runtime_subscription.rs +++ b/crates/workspace-server/src/runtime_subscription.rs @@ -10,12 +10,12 @@ use protocol::subscription::{ SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode, }; use tokio::sync::mpsc; -use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::{client_async_tls_with_config, connect_async}; use worker_runtime::auth::{CapabilityTokenSigner, capability_claims}; -use crate::hosts::RemoteRuntimeConfig; +use crate::hosts::{RemoteRuntimeConfig, resolve_strict_remote_runtime_endpoint}; const DOWNSTREAM_QUEUE_CAPACITY: usize = 256; const RECONNECT_DELAY: Duration = Duration::from_millis(100); @@ -918,10 +918,43 @@ async fn connect_runtime( .map_err(|error| format!("invalid Runtime authorization header: {error}"))?, ); } - connect_async(request) + if config.strict_public_egress { + let base_url = config.base_url.clone(); + let (_, addresses) = + tokio::task::spawn_blocking(move || resolve_strict_remote_runtime_endpoint(&base_url)) + .await + .map_err(|_| { + "Runtime subscription endpoint resolution task failed".to_string() + })??; + let stream = tokio::time::timeout(config.timeout, async move { + let mut last_error = None; + for address in addresses { + match tokio::net::TcpStream::connect(address).await { + Ok(stream) => return Ok(stream), + Err(error) => last_error = Some(error), + } + } + Err(last_error + .map(|error| error.to_string()) + .unwrap_or_else(|| "no validated Runtime address was available".to_string())) + }) .await + .map_err(|_| "Runtime subscription TCP connection timed out".to_string())??; + tokio::time::timeout( + config.timeout, + client_async_tls_with_config(request, stream, None, None), + ) + .await + .map_err(|_| "Runtime subscription TLS/WebSocket handshake timed out".to_string())? .map(|(socket, _)| socket) .map_err(|error| format!("failed to connect Runtime subscription endpoint: {error}")) + } else { + tokio::time::timeout(config.timeout, connect_async(request)) + .await + .map_err(|_| "Runtime subscription connection timed out".to_string())? + .map(|(socket, _)| socket) + .map_err(|error| format!("failed to connect Runtime subscription endpoint: {error}")) + } } fn runtime_endpoint(base_url: &str) -> String { let base = base_url.trim_end_matches('/'); diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index 74bc902a..7e3798e8 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -416,3 +416,19 @@ async fn embedded_runtime_uses_in_process_subscription_source() { )); server.abort(); } + +#[tokio::test] +async fn strict_runtime_subscription_rejects_private_endpoint_before_websocket_connect() { + let config = RemoteRuntimeConfig::new( + "runtime-private", + "Private Runtime", + "https://169.254.169.254", + None, + ) + .with_strict_public_egress(true); + let error = match connect_runtime(&config, "workspace-a").await { + Err(error) => error, + Ok(_) => panic!("private endpoint unexpectedly produced a WebSocket"), + }; + assert!(error.contains("endpoint host is not public"), "{error}"); +}