fix: bind Runtime WebSockets to egress policy

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"

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"] }

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<F>(
+    timeout: Duration,
+    resolver: F,
+) -> Result<Vec<SocketAddr>, String>
+where
+    F: FnOnce() -> Result<Vec<SocketAddr>, 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<SocketAddr>), 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::<Vec<_>>();
+    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::<Vec<_>>())
+            .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(

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('/');

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}");
+}