From b605b6df4c445e0e34808176fc0794f97369387a Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 25 Sep 2026 15:11:40 -0700 Subject: [PATCH 1/2] feat: expose TCP keepalive options through negotiated WISP extension --- server/Cargo.toml | 4 +- server/src/config.rs | 5 +- server/src/handle/wisp/keepalive.rs | 612 ++++++++++++++++++++++++++++ server/src/handle/wisp/mod.rs | 17 + wasmer/README.md | 7 + wasmer/config/config.toml | 2 +- wasmer/tcp-keepalive.md | 76 ++++ wasmer/test-native-keepalive.sh | 37 ++ wisp/src/extensions/mod.rs | 24 ++ wisp/src/mux/inner.rs | 30 +- wisp/src/packet.rs | 15 +- 11 files changed, 815 insertions(+), 14 deletions(-) create mode 100644 server/src/handle/wisp/keepalive.rs create mode 100644 wasmer/tcp-keepalive.md create mode 100755 wasmer/test-native-keepalive.sh diff --git a/server/Cargo.toml b/server/Cargo.toml index 1b220da..2a9590e 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -23,7 +23,7 @@ http-body-util = "0.1.2" hyper = { version = "1.4.1", features = ["server", "http1"] } hyper-util = { version = "0.1.7", features = ["tokio"] } lazy_static = "1.5.0" -libc = { version = "0.2.158", optional = true } +libc = "0.2.158" log = { version = "0.4.22", features = ["serde", "std"] } nix = { version = "0.30.1", features = ["term"] } pin-project-lite = "0.2.15" @@ -50,7 +50,7 @@ default = ["toml"] yaml = ["dep:serde_yaml"] toml = ["dep:toml"] -twisp = ["dep:pty-process", "dep:libc", "dep:shell-words"] +twisp = ["dep:pty-process", "dep:shell-words"] speed-limit = ["dep:async-speed-limit"] [build-dependencies] diff --git a/server/src/config.rs b/server/src/config.rs index 050eacc..b2f7513 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -136,6 +136,9 @@ pub enum ProtocolExtension { Udp, /// Wisp version 2 MOTD protocol extension. Motd, + /// Wasmer TCP keepalive socket options (private WISP v2 extension). + #[serde(rename = "tcp-keepalive")] + TcpKeepalive, /// Unofficial Wispnet-like protocol extension. Wispnet, } @@ -532,7 +535,7 @@ impl Config { config.wisp.allow_wsproxy = false; config.wisp.prefix.clear(); config.wisp.wisp_v2 = true; - config.wisp.extensions = vec![ProtocolExtension::Motd]; + config.wisp.extensions = vec![ProtocolExtension::Motd, ProtocolExtension::TcpKeepalive]; config.stream.allow_udp = false; config.stream.allow_wsproxy_udp = false; diff --git a/server/src/handle/wisp/keepalive.rs b/server/src/handle/wisp/keepalive.rs new file mode 100644 index 0000000..bd5e882 --- /dev/null +++ b/server/src/handle/wisp/keepalive.rs @@ -0,0 +1,612 @@ +//! Opt-in Wasmer TCP keepalive extension. See `wasmer/tcp-keepalive.md`. +use std::{ + collections::HashMap, + os::fd::{AsFd, AsRawFd, OwnedFd}, + sync::{Arc, Mutex}, +}; + +use async_trait::async_trait; +use bytes::{Buf, BufMut, Bytes, BytesMut}; +use tokio::sync::watch; +use wisp_mux::{ + extensions::{AnyProtocolExtension, ProtocolExtension, ProtocolExtensionBuilder}, + packet::StreamType, + stream::MuxStream, + ws::TransportWrite, + Role, WispError, +}; + +use crate::stream::ClientStream; + +const EXTENSION_ID: u8 = 0xF2; +const REQUEST_PACKET: u8 = 0xF2; +const REPLY_PACKET: u8 = 0xF3; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum Status { + Invalid = 1, + Unsupported = 2, + BadStream = 3, + Io = 4, +} + +#[derive(Debug, Clone)] +enum SocketState { + Connecting, + Connected(Arc), + Failed(Status), +} + +#[derive(Debug, Clone, Default)] +pub struct SocketRegistry(Arc>>>); + +impl SocketRegistry { + fn open(&self, id: u32, ty: StreamType) { + let state = if ty == StreamType::Tcp { + SocketState::Connecting + } else { + SocketState::Failed(Status::Unsupported) + }; + self.0.lock().unwrap().insert(id, watch::channel(state).0); + } + + fn close(&self, id: u32) { + if let Some(sender) = self.0.lock().unwrap().remove(&id) { + sender.send_replace(SocketState::Failed(Status::BadStream)); + } + } + + pub fn register_stream(&self, stream: &MuxStream) -> SocketRegistration { + let sockets = self.0.lock().unwrap(); + let sender = if stream.get_close_reason().is_none() { + sockets.get(&stream.get_stream_id()).cloned() + } else { + None + }; + SocketRegistration { + registry: self.clone(), + stream_id: stream.get_stream_id(), + sender, + } + } + + #[cfg(test)] + fn attach(&self, id: u32, stream: &ClientStream) -> SocketRegistration { + let registration = SocketRegistration { + registry: self.clone(), + stream_id: id, + sender: self.0.lock().unwrap().get(&id).cloned(), + }; + registration.attach(stream); + registration + } + + async fn socket(&self, id: u32) -> Result, Status> { + let mut receiver = self + .0 + .lock() + .unwrap() + .get(&id) + .ok_or(Status::BadStream)? + .subscribe(); + loop { + match receiver.borrow_and_update().clone() { + SocketState::Connected(fd) => return Ok(fd), + SocketState::Failed(status) => return Err(status), + SocketState::Connecting => {} + } + receiver.changed().await.map_err(|_| Status::BadStream)?; + } + } + + async fn request(&self, id: u32, operation: u8, option: u8, value: u32) -> Result { + if operation > 1 + || option > 3 + || (operation == 0 + && ((option == 0 && value > 1) + || (option != 0 && (value == 0 || value > i32::MAX as u32)))) + { + return Err(Status::Invalid); + } + let fd = self.socket(id).await?; + if operation == 0 { + socket_option(&fd, option, Some(value))?; + } + // ACK values and GET values both come from the socket, never a cache. + socket_option(&fd, option, None) + } +} + +/// Retire the duplicated descriptor when forwarding ends or is cancelled. +pub struct SocketRegistration { + registry: SocketRegistry, + stream_id: u32, + sender: Option>, +} +impl SocketRegistration { + pub fn attach(&self, stream: &ClientStream) { + let Some(sender) = &self.sender else { + return; + }; + sender.send_if_modified(|state| { + if !matches!(state, SocketState::Connecting) { + return false; + } + *state = match stream { + ClientStream::Tcp(stream) => match stream.as_fd().try_clone_to_owned() { + Ok(fd) => SocketState::Connected(Arc::new(fd)), + Err(error) => { + log::debug!("keepalive descriptor duplication failed: {error}"); + SocketState::Failed(Status::Io) + } + }, + _ => SocketState::Failed(Status::Unsupported), + }; + true + }); + } +} +impl Drop for SocketRegistration { + fn drop(&mut self) { + let Some(sender) = &self.sender else { + return; + }; + let mut sockets = self.registry.0.lock().unwrap(); + // A client can reuse a closed stream ID. An older forwarding task must + // never retire or attach its socket to that newer stream. + if sockets + .get(&self.stream_id) + .is_some_and(|current| current.same_channel(sender)) + { + sockets.remove(&self.stream_id); + } + sender.send_replace(SocketState::Failed(Status::BadStream)); + } +} + +#[derive(Debug, Clone)] +struct KeepaliveExtension(SocketRegistry); + +#[async_trait] +impl ProtocolExtension for KeepaliveExtension { + fn get_id(&self) -> u8 { + EXTENSION_ID + } + fn get_supported_packets(&self) -> &'static [u8] { + &[REQUEST_PACKET] + } + fn encode(&self) -> Bytes { + Bytes::from_static(&[1]) + } + fn on_stream_open(&mut self, id: u32, ty: StreamType) { + self.0.open(id, ty); + } + fn on_stream_close(&mut self, id: u32) { + self.0.close(id); + } + fn box_clone(&self) -> Box { + Box::new(self.clone()) + } + + async fn handle_stream_packet( + &mut self, + _: u8, + stream_id: u32, + mut packet: Bytes, + ) -> Result, WispError> { + if packet.len() < 4 { + return Err(WispError::PacketTooSmall); + } + let request_id = packet.get_u32_le(); + let result = if packet.len() == 6 { + let operation = packet.get_u8(); + let option = packet.get_u8(); + let value = packet.get_u32_le(); + self.0.request(stream_id, operation, option, value).await + } else { + Err(Status::Invalid) + }; + let mut reply = BytesMut::with_capacity(14); + reply.put_u8(REPLY_PACKET); + reply.put_u32_le(stream_id); + reply.put_u32_le(request_id); + match result { + Ok(value) => { + reply.put_u8(0); + reply.put_u32_le(value); + } + Err(status) => { + reply.put_u8(status as u8); + reply.put_u32_le(0); + } + } + Ok(Some(reply.freeze())) + } +} + +pub struct KeepaliveBuilder(pub SocketRegistry); +impl ProtocolExtensionBuilder for KeepaliveBuilder { + fn get_id(&self) -> u8 { + EXTENSION_ID + } + fn build_from_bytes( + &mut self, + bytes: Bytes, + role: Role, + ) -> Result { + if bytes.as_ref() != [1] { + return Err(WispError::ExtensionImplNotSupported); + } + self.build_to_extension(role) + } + fn build_to_extension(&mut self, role: Role) -> Result { + if !matches!(role, Role::Server) { + return Err(WispError::ExtensionImplNotSupported); + } + Ok(KeepaliveExtension(self.0.clone()).into()) + } +} + +#[cfg(not(target_os = "wasi"))] +fn socket_option(fd: &OwnedFd, option: u8, set: Option) -> Result { + let (level, name) = match option { + 0 => (libc::SOL_SOCKET, libc::SO_KEEPALIVE), + #[cfg(any(target_os = "macos", target_os = "ios"))] + 1 => (libc::IPPROTO_TCP, libc::TCP_KEEPALIVE), + #[cfg(any(target_os = "linux", target_os = "android", target_os = "freebsd"))] + 1 => (libc::IPPROTO_TCP, libc::TCP_KEEPIDLE), + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "linux", + target_os = "android", + target_os = "freebsd" + ))] + 2 => (libc::IPPROTO_TCP, libc::TCP_KEEPINTVL), + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "linux", + target_os = "android", + target_os = "freebsd" + ))] + 3 => (libc::IPPROTO_TCP, libc::TCP_KEEPCNT), + _ => return Err(Status::Unsupported), + }; + let mut value: libc::c_int = set.unwrap_or(0).try_into().map_err(|_| Status::Invalid)?; + let mut size = libc::socklen_t::try_from(size_of::()).map_err(|_| Status::Io)?; + // SAFETY: fd owns a live socket; both syscalls receive an initialized c_int + // and its exact size. Setting a timing option does not toggle SO_KEEPALIVE. + let result = unsafe { + if set.is_some() { + libc::setsockopt(fd.as_raw_fd(), level, name, (&raw const value).cast(), size) + } else { + libc::getsockopt( + fd.as_raw_fd(), + level, + name, + (&raw mut value).cast(), + &raw mut size, + ) + } + }; + if result != 0 { + let error = std::io::Error::last_os_error(); + log::debug!("keepalive socket option {option} failed: {error}"); + return Err(match error.raw_os_error() { + Some(libc::EINVAL) => Status::Invalid, + Some(code) + if code == libc::ENOPROTOOPT + || code == libc::EOPNOTSUPP + || code == libc::ENOSYS => + { + Status::Unsupported + } + _ => Status::Io, + }); + } + if option == 0 { + Ok(u32::from(value != 0)) + } else { + u32::try_from(value).map_err(|_| Status::Io) + } +} + +// Call the WASIX option ABI directly: old WASIX libc builds do not map the +// TCP_KEEP* POSIX options. This requires a runtime supporting options 27–29. +#[cfg(target_os = "wasi")] +fn socket_option(fd: &OwnedFd, option: u8, set: Option) -> Result { + #[link(wasm_import_module = "wasix_32v1")] + unsafe extern "C" { + fn sock_set_opt_flag(fd: u32, option: u8, value: u8) -> u16; + fn sock_get_opt_flag(fd: u32, option: u8, value: *mut u8) -> u16; + fn sock_set_opt_size(fd: u32, option: u8, value: u64) -> u16; + fn sock_get_opt_size(fd: u32, option: u8, value: *mut u64) -> u16; + } + let fd = fd.as_raw_fd() as u32; + let mut flag = 0_u8; + let mut size = 0_u64; + // SAFETY: the descriptor is live and output pointers have the ABI's size. + let errno = unsafe { + match (option, set) { + (0, Some(value)) => sock_set_opt_flag(fd, 12, value as u8), + (0, None) => sock_get_opt_flag(fd, 12, &raw mut flag), + (1..=3, Some(value)) => sock_set_opt_size(fd, 26 + option, value.into()), + (1..=3, None) => sock_get_opt_size(fd, 26 + option, &raw mut size), + _ => return Err(Status::Unsupported), + } + }; + match errno { + 0 => { + if option == 0 { + Ok(flag.into()) + } else { + size.try_into().map_err(|_| Status::Io) + } + } + 28 => Err(Status::Invalid), + 50 | 52 | 58 => Err(Status::Unsupported), + _ => { + log::debug!("keepalive WASIX option {option} failed: errno {errno}"); + Err(Status::Io) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures_util::{future::ready, sink, stream, StreamExt}; + use std::{pin::Pin, time::Duration}; + use tokio::{ + net::{TcpListener, TcpStream}, + sync::mpsc, + time::timeout, + }; + use wisp_mux::{packet::CloseReason, ws::TransportWrite, ServerMux, WispV2Handshake}; + + type TestWrite = Pin>; + type TestMux = ServerMux; + + async fn socket_pair() -> (ClientStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let (outgoing, incoming) = tokio::join!( + TcpStream::connect(listener.local_addr().unwrap()), + listener.accept() + ); + (ClientStream::Tcp(outgoing.unwrap()), incoming.unwrap().0) + } + + #[tokio::test] + async fn keepalive_kernel_readback_and_enable_state() { + let (stream, _peer) = socket_pair().await; + let registry = SocketRegistry::default(); + registry.open(7, StreamType::Tcp); + let registration = registry.attach(7, &stream); + assert_eq!(registry.request(7, 0, 0, 0).await, Ok(0)); + for (option, value) in [(1, 23), (2, 7), (3, 4)] { + assert_eq!(registry.request(7, 0, option, value).await, Ok(value)); + assert_eq!(registry.request(7, 1, option, 0).await, Ok(value)); + assert_eq!(registry.request(7, 1, 0, 0).await, Ok(0)); + } + assert_eq!(registry.request(7, 0, 0, 1).await, Ok(1)); + assert_eq!(registry.request(7, 0, 1, 31).await, Ok(31)); + assert_eq!(registry.request(7, 1, 0, 0).await, Ok(1)); + assert_eq!(registry.request(7, 1, 2, 0).await, Ok(7)); + assert_eq!(registry.request(7, 1, 3, 0).await, Ok(4)); + for (operation, option, value) in + [(0, 0, 2), (0, 1, 0), (0, 2, u32::MAX), (2, 1, 0), (1, 4, 0)] + { + assert_eq!( + registry.request(7, operation, option, value).await, + Err(Status::Invalid) + ); + } + // A real getsockopt failure must remain an error, never cached success. + let non_socket: OwnedFd = std::fs::File::open("/dev/null").unwrap().into(); + assert_eq!(socket_option(&non_socket, 0, None), Err(Status::Io)); + assert_eq!(registry.request(7, 1, 3, 0).await, Ok(4)); + drop(registration); + assert_eq!(registry.request(7, 1, 0, 0).await, Err(Status::BadStream)); + registry.open(8, StreamType::Udp); + assert_eq!(registry.request(8, 1, 0, 0).await, Err(Status::Unsupported)); + } + + fn info(version: Option) -> Bytes { + let mut bytes = BytesMut::from(&[5, 0, 0, 0, 0, 2, 0][..]); + if let Some(version) = version { + bytes.put_u8(EXTENSION_ID); + bytes.put_u32_le(1); + bytes.put_u8(version); + } + bytes.freeze() + } + + fn request(id: u32, sequence: u32, operation: u8, option: u8, value: u32) -> Bytes { + let mut bytes = BytesMut::new(); + bytes.put_u8(REQUEST_PACKET); + bytes.put_u32_le(id); + bytes.put_u32_le(sequence); + bytes.put_u8(operation); + bytes.put_u8(option); + bytes.put_u32_le(value); + bytes.freeze() + } + + fn connect(id: u32) -> Bytes { + let mut bytes = BytesMut::new(); + bytes.put_u8(1); + bytes.put_u32_le(id); + bytes.put_u8(1); + bytes.put_u16_le(80); + bytes.extend_from_slice(b"127.0.0.1"); + bytes.freeze() + } + + async fn next_reply(rx: &mut mpsc::UnboundedReceiver) -> Bytes { + timeout(Duration::from_secs(2), async { + loop { + let frame = rx.recv().await.unwrap(); + if frame[0] == REPLY_PACKET { + return frame; + } + } + }) + .await + .expect("keepalive response must not deadlock") + } + + async fn mux( + negotiated: bool, + ) -> ( + TestMux, + tokio::task::JoinHandle>, + SocketRegistry, + mpsc::UnboundedSender, + mpsc::UnboundedReceiver, + ) { + let (input, receiver) = mpsc::unbounded_channel(); + let (sender, mut output) = mpsc::unbounded_channel(); + let read = Box::pin(stream::unfold(receiver, |mut receiver| async move { + receiver.recv().await.map(|bytes| (Ok(bytes), receiver)) + })); + let write: TestWrite = Box::pin(sink::unfold( + sender, + |sender: mpsc::UnboundedSender, bytes| { + ready( + sender + .send(bytes) + .map(|()| sender) + .map_err(|_| WispError::WsImplSocketClosed), + ) + }, + )); + let registry = SocketRegistry::default(); + input.send(info(negotiated.then_some(1))).unwrap(); + let (mux, future) = ServerMux::new( + read, + write, + 128, + Some(WispV2Handshake::new(vec![KeepaliveBuilder( + registry.clone(), + ) + .into()])), + ) + .await + .unwrap() + .with_no_required_extensions(); + assert_eq!(output.recv().await.unwrap(), info(Some(1))); + assert_eq!(output.recv().await.unwrap()[0], 3); + assert_eq!(mux.get_extension_ids().contains(&EXTENSION_ID), negotiated); + (mux, tokio::spawn(future), registry, input, output) + } + + #[tokio::test] + async fn keepalive_negotiated_requests_wait_for_connect_and_continue_reading() { + let (mux, task, registry, input, mut output) = mux(true).await; + input.send(connect(9)).unwrap(); + input.send(request(9, 42, 0, 1, 29)).unwrap(); + let (_connect, mut stream) = mux.wait_for_stream().await.unwrap(); + assert!(timeout(Duration::from_millis(20), output.recv()) + .await + .is_err()); + let (socket, _peer) = socket_pair().await; + let registration = registry.register_stream(&stream); + registration.attach(&socket); + let reply = next_reply(&mut output).await; + assert_eq!( + reply.as_ref(), + &[REPLY_PACKET, 9, 0, 0, 0, 42, 0, 0, 0, 0, 29, 0, 0, 0] + ); + input.send(request(9, 43, 1, 1, 0)).unwrap(); + assert_eq!(next_reply(&mut output).await[10..], 29_u32.to_le_bytes()); + input.send(request(9, 44, 1, 0, 0)).unwrap(); + assert_eq!(next_reply(&mut output).await[10..], 0_u32.to_le_bytes()); + // Ordinary traffic must still be delivered after control packets. + input + .send(Bytes::from_static(&[2, 9, 0, 0, 0, 72, 73])) + .unwrap(); + assert_eq!( + timeout(Duration::from_secs(2), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(), + Bytes::from_static(b"HI") + ); + input.send(request(9, 45, 0, 0, 2)).unwrap(); + assert_eq!(next_reply(&mut output).await[9], Status::Invalid as u8); + timeout(Duration::from_secs(2), stream.close(CloseReason::Voluntary)) + .await + .unwrap() + .unwrap(); + input.send(request(9, 46, 1, 0, 0)).unwrap(); + assert_eq!(next_reply(&mut output).await[9], Status::BadStream as u8); + // Reusing a closed stream ID cannot let an old forwarding task attach + // its socket to or retire the new stream. + input.send(connect(9)).unwrap(); + let (_, replacement) = mux.wait_for_stream().await.unwrap(); + let (replacement_socket, _replacement_peer) = socket_pair().await; + let replacement_registration = registry.register_stream(&replacement); + replacement_registration.attach(&replacement_socket); + registration.attach(&socket); + drop(registration); + input.send(request(9, 47, 0, 1, 37)).unwrap(); + assert_eq!(next_reply(&mut output).await[10..], 37_u32.to_le_bytes()); + drop(replacement_registration); + drop(input); + timeout(Duration::from_secs(2), task) + .await + .unwrap() + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn keepalive_connect_failure_wakes_pending_request_without_write_lock_deadlock() { + let (mux, task, _registry, input, mut output) = mux(true).await; + input.send(connect(11)).unwrap(); + input.send(request(11, 1, 0, 0, 1)).unwrap(); + let (_, stream) = mux.wait_for_stream().await.unwrap(); + assert!(timeout(Duration::from_millis(20), output.recv()) + .await + .is_err()); + timeout( + Duration::from_secs(2), + stream.close(CloseReason::ServerStreamUnreachable), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(next_reply(&mut output).await[9], Status::BadStream as u8); + drop(input); + timeout(Duration::from_secs(2), task) + .await + .unwrap() + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn keepalive_requires_negotiation_and_exact_version() { + let (_mux, task, registry, input, _output) = mux(false).await; + input.send(request(1, 1, 0, 0, 1)).unwrap(); + assert!(matches!( + timeout(Duration::from_secs(2), task) + .await + .unwrap() + .unwrap(), + Err(WispError::InvalidPacketType(REQUEST_PACKET)) + )); + assert!(registry.0.lock().unwrap().is_empty()); + let mut builder = KeepaliveBuilder(SocketRegistry::default()); + for version in [ + Bytes::new(), + Bytes::from_static(&[2]), + Bytes::from_static(&[1, 0]), + ] { + assert!(builder.build_from_bytes(version, Role::Server).is_err()); + } + } +} diff --git a/server/src/handle/wisp/mod.rs b/server/src/handle/wisp/mod.rs index 289c72b..17a1c1b 100644 --- a/server/src/handle/wisp/mod.rs +++ b/server/src/handle/wisp/mod.rs @@ -1,3 +1,4 @@ +mod keepalive; #[cfg(feature = "twisp")] pub mod twisp; pub mod utils; @@ -243,6 +244,7 @@ async fn handle_stream( muxstream: MuxStream, id: String, event: Arc, + keepalive: keepalive::SocketRegistration, #[cfg(feature = "twisp")] twisp_map: twisp::TwispMap, #[cfg(feature = "speed-limit")] read_limit: async_speed_limit::Limiter< async_speed_limit::clock::StandardClock, @@ -263,6 +265,7 @@ async fn handle_stream( return; }; + keepalive.attach(&stream); debug!("[{id}] [{uuid}] stream resolved and connected: {resolved_stream:?}"); if let Some(client) = CLIENTS.lock().await.get(&id) { @@ -319,6 +322,18 @@ pub async fn handle_wisp(stream: WispResult, is_v2: bool, id: String) -> anyhow: } } + let keepalive = keepalive::SocketRegistry::default(); + let mut extensions = extensions; + if CONFIG + .wisp + .extensions + .contains(&crate::config::ProtocolExtension::TcpKeepalive) + { + if let Some(extensions) = &mut extensions { + extensions.add_extension(keepalive::KeepaliveBuilder(keepalive.clone()).into()); + } + } + #[cfg(feature = "speed-limit")] let read_limiter = async_speed_limit::Limiter::builder(CONFIG.wisp.read_limit) .refill(Duration::from_secs(1)) @@ -383,11 +398,13 @@ pub async fn handle_wisp(stream: WispResult, is_v2: bool, id: String) -> anyhow: }); while let Some((connect, stream)) = mux.wait_for_stream().await { + let registration = keepalive.register_stream(&stream); set.spawn(handle_stream( connect, stream, id.clone(), event.clone(), + registration, #[cfg(feature = "twisp")] twisp_map.clone(), #[cfg(feature = "speed-limit")] diff --git a/wasmer/README.md b/wasmer/README.md index 8329a28..ca6edac 100644 --- a/wasmer/README.md +++ b/wasmer/README.md @@ -79,3 +79,10 @@ multicast, and non-global targets. It permits direct global IP addresses because browser WASIX clients connect to the address selected by their DNS resolver. Add WISP authentication or enforce access at the deployment edge before exposing it publicly; an unauthenticated WISP endpoint is still an outbound proxy. + +## TCP keepalive options + +The package enables the negotiated [Wasmer TCP keepalive extension](tcp-keepalive.md). +It exposes real socket setters/getters to compatible browser clients. Its WASIX +deployment host must support the new idle/interval/count socket option imports; +building this package alone does not update the host runtime. diff --git a/wasmer/config/config.toml b/wasmer/config/config.toml index 4a86e42..b50f09e 100644 --- a/wasmer/config/config.toml +++ b/wasmer/config/config.toml @@ -8,7 +8,7 @@ log_level = "INFO" prefix = "" wisp_v2 = true allow_wsproxy = false -extensions = ["motd"] +extensions = ["motd", "tcp-keepalive"] [stream] allow_udp = false diff --git a/wasmer/tcp-keepalive.md b/wasmer/tcp-keepalive.md new file mode 100644 index 0000000..3bed4c3 --- /dev/null +++ b/wasmer/tcp-keepalive.md @@ -0,0 +1,76 @@ +# Wasmer TCP keepalive extension (version 1) + +This private WISP v2 extension controls the outbound TCP socket owned by the +proxy. Enable it with `extensions = ["motd", "tcp-keepalive"]` in `[wisp]`, or +`WISP_PROTOCOL_EXTENSIONS=motd,tcp-keepalive`. The Wasmer package enables it; +ordinary server defaults do not. Clients must explicitly negotiate it. + +The extension ID is `0xF2`, with the one-byte metadata payload `[1]` in both INFO +packets. Other versions or metadata lengths are rejected. IDs `0xF0` and `0xF1` +are already used by TWisp and Wispnet. This is a private extension, not part of +the standard WISP protocol. + +Requests use packet type `0xF2`, replies `0xF3`. Both use the normal WISP header: +packet type `u8`, followed by the target stream ID `u32` little-endian. Stream +zero is not a TCP stream. An unnegotiated request is an invalid WISP packet and +closes the multiplexor; it is never silently accepted. + +The request payload is exactly 10 bytes: + +| Field | Type | Meaning | +| --- | --- | --- | +| request ID | u32 LE | Client correlation ID, echoed in the reply | +| operation | u8 | 0: set; 1: get | +| option | u8 | 0: enabled; 1: idle seconds; 2: interval seconds; 3: probe count | +| value | u32 LE | Setter value; ignored by get | + +Enabled accepts only 0 or 1. Timing/count setters accept positive values no +larger than `INT_MAX`; the host OS can impose additional limits. Changing an +individual timing option preserves the other timings and the enabled state. + +The reply payload is exactly 9 bytes: + +| Field | Type | Meaning | +| --- | --- | --- | +| request ID | u32 LE | Request being acknowledged | +| status | u8 | 0: success; 1: invalid; 2: unsupported; 3: bad/closed stream; 4: OS I/O error | +| value | u32 LE | Actual socket value on success, otherwise zero | + +Both SET acknowledgements and GET replies read back the socket with +`getsockopt` (or the corresponding WASIX import). Enabled is normalized to 0/1, +since some kernels return their socket option bitmask. Native error codes are +mapped to the statuses above; errors are never converted to successful ACKs. +A setter may have taken effect even if its subsequent readback fails, as with +other nontransactional socket APIs. + +A request for a stream that is still resolving/connecting waits for that +connection. The stream is registered when CONNECT is processed, before its +connection task starts. Failure or closure wakes the waiting request with +status 3. UDP and Wispnet forwarding streams return status 2. A duplicated, +owned descriptor keeps the socket identity stable during a request and is +retired when forwarding ends. Registrations retain stream identity even when a +client reuses a closed stream ID. Requests are processed in wire order, and replies +are sent only after the operation has completed. The transport write lock is +not held while waiting for a connection, so failure/close processing can run. + +Native builds use POSIX socket options on Linux, Android, macOS, iOS and FreeBSD. +Other native platforms return unsupported for timing options they cannot map. +The WASIX build calls `wasix_32v1.sock_set_opt_size` / `sock_get_opt_size` with +option IDs 27 (idle), 28 (interval), and 29 (count), and the flag imports with +option 12 for enabled. These direct imports avoid the old libc TCP_KEEP* mapping, +which returns ENOSYS. **The deployed proxy's Wasmer runtime also needs support +for these new option IDs and real socket operations.** Updating only the browser +SDK or the proxy package cannot add support to an older deployment host. + +## Validation + +`./wasmer/test-native-keepalive.sh` runs the focused native tests against the +source tree. It creates an ignored native workspace because the committed lock +file selects WASIX-only crate versions unavailable on crates.io. The tests +cover real kernel readback, enabled-state preservation, real OS errors, +negotiation/version rejection, requests before connection completion, failed +connections without deadlock, reused stream IDs, and both repeated requests and ordinary data after an +extension packet. + +`./wasmer/build.sh` builds the WASIX package module. Building it does not publish +the package or update any deployed proxy. diff --git a/wasmer/test-native-keepalive.sh b/wasmer/test-native-keepalive.sh new file mode 100755 index 0000000..dd919aa --- /dev/null +++ b/wasmer/test-native-keepalive.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +readonly NATIVE_DIR="${ROOT_DIR}/.wasmer-build/native" + +# Keep the production WASIX lockfile/registry settings intact. Sources are +# symlinked so native checks always exercise the current working tree. +python3 - "${ROOT_DIR}" "${NATIVE_DIR}" <<'PY' +from pathlib import Path +import shutil +import sys + +root, native = map(Path, sys.argv[1:]) +native.mkdir(parents=True, exist_ok=True) +manifest = (root / "Cargo.toml").read_text() +start = manifest.index("[patch.crates-io]") +end = manifest.index("[workspace.lints", start) +(native / "Cargo.toml").write_text(manifest[:start] + manifest[end:]) +for name in ("server", "wisp"): + package = native / name + package.mkdir(exist_ok=True) + shutil.copyfile(root / name / "Cargo.toml", package / "Cargo.toml") + for entry in ("src", "build.rs"): + source = root / name / entry + destination = package / entry + if source.exists() and not destination.exists(): + destination.symlink_to(source) +PY + +# Cargo discovers .cargo/config from the cwd, even with --manifest-path. +# Run outside the repository to avoid its WASIX registry replacement. +cd "${TMPDIR:-/tmp}" +# Match the production lock: newer vergen 9.1 uses an incompatible vergen-lib +# with vergen-git2 1.0.7. This pin only affects the ignored native test workspace. +cargo +nightly update --manifest-path "${NATIVE_DIR}/Cargo.toml" -p vergen --precise 9.0.6 +cargo +nightly test --manifest-path "${NATIVE_DIR}/Cargo.toml" -p epoxy-server keepalive -- "$@" diff --git a/wisp/src/extensions/mod.rs b/wisp/src/extensions/mod.rs index 4086b51..0d0c722 100644 --- a/wisp/src/extensions/mod.rs +++ b/wisp/src/extensions/mod.rs @@ -15,6 +15,7 @@ use async_trait::async_trait; use bytes::{BufMut, Bytes}; use crate::{ + packet::StreamType, ws::{PayloadMut, TransportRead, TransportWrite}, Role, WispError, }; @@ -131,6 +132,29 @@ pub trait ProtocolExtension: std::fmt::Debug + Sync + Send + 'static { Ok(()) } + /// Handle a stream control request without holding the transport write lock. + /// Return a complete response frame, or `None` to use `handle_packet`. + async fn handle_stream_packet( + &mut self, + packet_type: u8, + stream_id: u32, + packet: Bytes, + ) -> Result, WispError> { + let _ = (packet_type, stream_id, packet); + Ok(None) + } + + /// Called before a new stream is handed to the application. Cloned extension + /// instances must share any state used by this hook and the packet handler. + fn on_stream_open(&mut self, stream_id: u32, stream_type: StreamType) { + let _ = (stream_id, stream_type); + } + + /// Called when a stream closes, including multiplexor shutdown. + fn on_stream_close(&mut self, stream_id: u32) { + let _ = stream_id; + } + /// Clone the protocol extension. fn box_clone(&self) -> Box; diff --git a/wisp/src/mux/inner.rs b/wisp/src/mux/inner.rs index 428104d..cf5f388 100644 --- a/wisp/src/mux/inner.rs +++ b/wisp/src/mux/inner.rs @@ -172,6 +172,7 @@ pub(crate) struct MuxInner>, tx: LockedWebSocketWrite, flow_stream_types: Box<[u8]>, + extensions: Vec, mux: M, @@ -212,11 +213,12 @@ impl> MuxInner> MuxInner> MuxInner>, WispError> = async { - if let Some(msg) = rx.next().await { + while let Some(msg) = rx.next().await { match MaybeExtensionPacket::decode(msg?, &mut extensions, &mut rx, &mut tx) .await? { - MaybeExtensionPacket::Packet(x) => Ok(Some(WsEvent::WispMessage(x))), - MaybeExtensionPacket::ExtensionHandled => Ok(None), + MaybeExtensionPacket::Packet(x) => { + return Ok(Some(WsEvent::WispMessage(x))) + } + MaybeExtensionPacket::ExtensionHandled => {} } - } else { - Ok(None) } + Ok(None) } .await; ret.transpose().map(|x| (x, (rx, tx, extensions))) @@ -310,6 +316,9 @@ impl> MuxInner { if let Some(stream) = self.streams.remove(&id) { + for extension in &mut self.extensions { + extension.on_stream_close(id); + } Self::close_stream(stream, close); let pkt = Packet { stream_id: id, @@ -373,6 +382,9 @@ impl> MuxInner> MuxInner { + if packet.stream_id == 0 || self.streams.contains_key(&packet.stream_id) { + return Err(WispError::InvalidStreamId(packet.stream_id)); + } let stream = self.add_stream(packet.stream_id, connect.stream_type); self.mux.handle_connect_packet(stream, connect)?; Ok(false) @@ -416,6 +431,9 @@ impl> MuxInner { packet_type: PacketType::decode(packet, ty)?, })) } else { - tx.lock().await; - let mut handle = tx.get_handle(); for extension in extensions { if extension.get_supported_packets().contains(&ty) { - extension.handle_packet(ty, packet, rx, &mut handle).await?; + let response = extension + .handle_stream_packet(ty, stream_id, packet.clone()) + .await?; + tx.lock().await; + let mut handle = tx.get_handle(); + if let Some(response) = response { + handle.send(response).await?; + } else { + extension.handle_packet(ty, packet, rx, &mut handle).await?; + } return Ok(Self::ExtensionHandled); } } - drop(handle); Err(WispError::InvalidPacketType(ty)) } From ebacdab640affa75dab630055e386e0bc4e008b7 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 25 Sep 2026 15:15:14 -0700 Subject: [PATCH 2/2] fix: duplicate WASIX socket descriptors through fd_dup2 --- server/src/handle/wisp/keepalive.rs | 29 +++++++++++++++++++++++++++-- wasmer/tcp-keepalive.md | 4 +++- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/server/src/handle/wisp/keepalive.rs b/server/src/handle/wisp/keepalive.rs index bd5e882..de939d4 100644 --- a/server/src/handle/wisp/keepalive.rs +++ b/server/src/handle/wisp/keepalive.rs @@ -1,7 +1,7 @@ //! Opt-in Wasmer TCP keepalive extension. See `wasmer/tcp-keepalive.md`. use std::{ collections::HashMap, - os::fd::{AsFd, AsRawFd, OwnedFd}, + os::fd::{AsRawFd, OwnedFd}, sync::{Arc, Mutex}, }; @@ -134,7 +134,7 @@ impl SocketRegistration { return false; } *state = match stream { - ClientStream::Tcp(stream) => match stream.as_fd().try_clone_to_owned() { + ClientStream::Tcp(stream) => match duplicate_socket(stream) { Ok(fd) => SocketState::Connected(Arc::new(fd)), Err(error) => { log::debug!("keepalive descriptor duplication failed: {error}"); @@ -165,6 +165,31 @@ impl Drop for SocketRegistration { } } +#[cfg(not(target_os = "wasi"))] +fn duplicate_socket(stream: &tokio::net::TcpStream) -> std::io::Result { + use std::os::fd::AsFd; + stream.as_fd().try_clone_to_owned() +} + +#[cfg(target_os = "wasi")] +fn duplicate_socket(stream: &tokio::net::TcpStream) -> std::io::Result { + use std::os::fd::FromRawFd; + #[link(wasm_import_module = "wasix_32v1")] + unsafe extern "C" { + fn fd_dup2(fd: u32, min_result_fd: u32, cloexec: u8, result: *mut u32) -> u16; + } + let mut duplicated = 0_u32; + // Rust's BorrowedFd::try_clone_to_owned is an unsupported WASI stub. + // WASIX provides real duplication, including close-on-exec semantics. + // SAFETY: the input fd is live and result points to one writable WASI fd. + let errno = unsafe { fd_dup2(stream.as_raw_fd() as u32, 0, 1, &raw mut duplicated) }; + if errno != 0 { + return Err(std::io::Error::from_raw_os_error(errno.into())); + } + // SAFETY: fd_dup2 returned a new descriptor owned solely by this value. + Ok(unsafe { OwnedFd::from_raw_fd(duplicated as i32) }) +} + #[derive(Debug, Clone)] struct KeepaliveExtension(SocketRegistry); diff --git a/wasmer/tcp-keepalive.md b/wasmer/tcp-keepalive.md index 3bed4c3..128f0c3 100644 --- a/wasmer/tcp-keepalive.md +++ b/wasmer/tcp-keepalive.md @@ -58,7 +58,9 @@ Other native platforms return unsupported for timing options they cannot map. The WASIX build calls `wasix_32v1.sock_set_opt_size` / `sock_get_opt_size` with option IDs 27 (idle), 28 (interval), and 29 (count), and the flag imports with option 12 for enabled. These direct imports avoid the old libc TCP_KEEP* mapping, -which returns ENOSYS. **The deployed proxy's Wasmer runtime also needs support +which returns ENOSYS. WASIX descriptor duplication uses the existing `fd_dup2` +import because Rust's WASI descriptor-cloning implementation is unsupported. +**The deployed proxy's Wasmer runtime also needs support for these new option IDs and real socket operations.** Updating only the browser SDK or the proxy package cannot add support to an older deployment host.