diff --git a/Cargo.lock b/Cargo.lock index 5207fce6f..4e92167d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7616,6 +7616,7 @@ dependencies = [ "fastrand", "futures-util", "ipnetwork", + "jmux-proto", "libsql", "mcp-proxy", "network-scanner", diff --git a/crates/agent-tunnel-proto/src/lib.rs b/crates/agent-tunnel-proto/src/lib.rs index 17af4919e..32bc5ce14 100644 --- a/crates/agent-tunnel-proto/src/lib.rs +++ b/crates/agent-tunnel-proto/src/lib.rs @@ -29,6 +29,9 @@ pub use session::{ConnectRequest, ConnectResponse, MAX_SESSION_MESSAGE_SIZE}; pub use stream::{ControlStream, FramedRecv, FramedSend, SessionStream}; pub use version::{ALPN_PROTOCOL, CURRENT_PROTOCOL_VERSION, MIN_SUPPORTED_VERSION, validate_protocol_version}; +/// Maximum time the Gateway keeps an Agent route online without a liveness message. +pub const AGENT_OFFLINE_TIMEOUT_SECS: u64 = 90; + /// Current wall-clock time in milliseconds since UNIX epoch. pub fn current_time_millis() -> u64 { u64::try_from( diff --git a/crates/agent-tunnel/src/registry.rs b/crates/agent-tunnel/src/registry.rs index eaa92df38..5587804f9 100644 --- a/crates/agent-tunnel/src/registry.rs +++ b/crates/agent-tunnel/src/registry.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; -use agent_tunnel_proto::{DomainAdvertisement, current_time_millis}; +use agent_tunnel_proto::{AGENT_OFFLINE_TIMEOUT_SECS, DomainAdvertisement, current_time_millis}; use ipnetwork::Ipv4Network; use parking_lot::RwLock; use serde::Serialize; @@ -13,7 +13,7 @@ use uuid::Uuid; use crate::routing::RouteTarget; /// Duration after which an agent is considered offline if no heartbeat has been received. -pub const AGENT_OFFLINE_TIMEOUT: Duration = Duration::from_secs(90); +pub const AGENT_OFFLINE_TIMEOUT: Duration = Duration::from_secs(AGENT_OFFLINE_TIMEOUT_SECS); /// Tracks route advertisements received from an agent. /// diff --git a/crates/jmux-proxy/src/event.rs b/crates/jmux-proxy/src/event.rs index 5ecfe1845..212405848 100644 --- a/crates/jmux-proxy/src/event.rs +++ b/crates/jmux-proxy/src/event.rs @@ -66,8 +66,8 @@ pub enum EventOutcome { /// Complete audit information for one traffic item's lifecycle. /// -/// A single `TrafficEvent` is emitted exactly once per JMUX traffic item when -/// it ends (successfully or with error). +/// A single `TrafficEvent` is emitted exactly once when an eligible JMUX traffic item ends. +/// An item is eligible only when direct target resolution provides a concrete IP address. /// /// # Timestamp semantics /// @@ -94,7 +94,7 @@ pub enum EventOutcome { /// - connect failure: the last IP that was attempted /// - `target_port`: the destination port. /// -/// DNS failures do **not** produce an event because `target_ip` is unknown. +/// DNS failures and connection attempts handled by connector overrides do **not** produce an event because `target_ip` is unknown. #[derive(Clone, Debug)] pub struct TrafficEvent { /// How the traffic item's lifecycle ended. @@ -142,13 +142,13 @@ pub struct TrafficEvent { /// Type-erased traffic audit callback. /// -/// Invoked exactly once per JMUX traffic item at end-of-lifecycle. The callback -/// itself is **synchronous**; perform any asynchronous work by spawning within -/// the callback (e.g., `tokio::spawn`) or by sending to an internal channel. +/// Invoked exactly once at the end of each eligible JMUX traffic item. +/// An item is eligible only when direct target resolution provides a concrete IP address. +/// The callback itself is **synchronous**; perform asynchronous work by spawning within the callback (e.g., `tokio::spawn`) or by sending to an internal channel. /// /// # Exactly-once /// -/// - Each traffic item yields exactly one event. +/// - Each eligible traffic item yields exactly one event. /// - Emitted at cleanup time, not during operation. /// - Guarded to prevent duplicate emission. /// - No aggregation—each event stands alone. @@ -162,7 +162,7 @@ pub struct TrafficEvent { /// /// ```rust,ignore /// let proxy = JmuxProxy::new(reader, writer) -/// .with_traffic_event_callback(|event| { +/// .with_outgoing_traffic_event_callback(|event| { /// // Log quickly... /// tracing::info!( /// outcome = ?event.outcome, diff --git a/crates/jmux-proxy/src/lib.rs b/crates/jmux-proxy/src/lib.rs index 7373fce85..4c5c7d00f 100644 --- a/crates/jmux-proxy/src/lib.rs +++ b/crates/jmux-proxy/src/lib.rs @@ -12,7 +12,10 @@ mod id_allocator; use std::collections::{HashMap, HashSet}; use std::convert::TryFrom; +use std::future::Future; use std::io; +use std::net::IpAddr; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::time::SystemTime; @@ -22,7 +25,6 @@ use bytes::Bytes; use jmux_proto::{ChannelData, DistantChannelId, Header, LocalChannelId, Message, ReasonCode}; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio::net::TcpStream; -use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; use tokio::sync::{Notify, mpsc, oneshot}; use tokio::task::JoinHandle; use tokio_util::codec::FramedRead; @@ -53,6 +55,15 @@ pub type ApiResponseReceiver = oneshot::Receiver; pub type ApiRequestSender = mpsc::Sender; pub type ApiRequestReceiver = mpsc::Receiver; +// A supertrait is required because trait objects may include only one non-auto trait. +trait TargetStream: AsyncRead + AsyncWrite + Unpin + Send {} + +impl TargetStream for T where T: AsyncRead + AsyncWrite + Unpin + Send {} + +type ErasedTargetStream = Box; +type TargetConnectorOverrideFuture = Pin>> + Send>>; +type TargetConnectorOverride = Arc TargetConnectorOverrideFuture + Send + Sync>; + #[derive(Debug)] pub enum JmuxApiRequest { OpenChannel { @@ -84,6 +95,7 @@ pub struct JmuxProxy { jmux_reader: Box, jmux_writer: Box, traffic_callback: Option, + target_connector_override: Option, } impl JmuxProxy { @@ -98,6 +110,7 @@ impl JmuxProxy { jmux_reader, jmux_writer, traffic_callback: None, + target_connector_override: None, } } @@ -113,11 +126,33 @@ impl JmuxProxy { self } + /// Overrides the default direct TCP connector when applicable. + /// + /// Return `Ok(Some(stream))` to use the override, `Ok(None)` to delegate to the default connector, or an error to reject the connection without falling back. + /// Connection attempts handled by the override do not emit outgoing traffic events because their resolved target IP is unknown. + #[must_use] + pub fn with_target_connector_override(mut self, connector: C) -> Self + where + C: Fn(DestinationUrl) -> F + Send + Sync + 'static, + F: Future>> + Send + 'static, + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + { + self.target_connector_override = Some(Arc::new(move |destination_url| { + let connect = connector(destination_url); + + Box::pin(async move { + connect + .await + .map(|stream| stream.map(|stream| Box::new(stream) as ErasedTargetStream)) + }) + })); + self + } + /// Configures an outgoing-traffic callback for lifecycle event monitoring. /// - /// The provided callback will be invoked exactly once per outgoing stream at the end of its - /// lifecycle, providing comprehensive audit information including connection metadata, - /// byte counts, timing, and termination classification. + /// The provided callback is invoked exactly once at the end of each outgoing stream whose resolved target IP is known. + /// It provides connection metadata, byte counts, timing, and termination classification. /// /// # Event Emission /// @@ -128,6 +163,7 @@ impl JmuxProxy { /// /// Events are **NOT** emitted for: /// - DNS resolution failures (no concrete IP address available) + /// - Connection attempts handled by a target connector override (no concrete IP address available) /// - Internal JMUX protocol errors before stream establishment /// /// For hostnames with multiple IP addresses, connection attempts follow a Happy Eyeballs @@ -135,7 +171,7 @@ impl JmuxProxy { /// /// # Callback Contract /// - /// - **Exactly once**: Each traffic item generates precisely one event, protected by atomic guards + /// - **Exactly once**: Each eligible traffic item generates one event, protected by atomic guards /// - **At stream end**: Events are emitted during cleanup, not during operation /// - **Synchronous**: The callback is called synchronously from JMUX task contexts /// - **Thread safe**: Must be `Send + Sync + 'static` for multi-threaded access @@ -186,6 +222,7 @@ async fn run_proxy_impl(proxy: JmuxProxy, span: Span) -> anyhow::Result<()> { jmux_reader, jmux_writer, traffic_callback, + target_connector_override, } = proxy; let (msg_to_send_tx, msg_to_send_rx) = mpsc::channel::(JMUX_MESSAGE_MPSC_CHANNEL_SIZE); @@ -206,6 +243,7 @@ async fn run_proxy_impl(proxy: JmuxProxy, span: Span) -> anyhow::Result<()> { msg_to_send_tx, api_request_rx, traffic_callback, + target_connector_override, parent_span: span, } .spawn(); @@ -255,7 +293,7 @@ struct JmuxChannelCtx { // Traffic audit metadata target_host: String, /// Target server resolved address IP - target_ip: Option, + target_ip: Option, /// Target server port target_port: u16, /// Time the connection with target peer was established at @@ -303,7 +341,7 @@ impl JmuxCtx { fn unregister(&mut self, id: LocalChannelId, traffic_callback: &Option, is_abnormal_error: bool) { if let Some(channel) = self.channels.remove(&id) { // Emit audit event if we have a callback and haven't already emitted. - // For now, we only emit an event when the IP address is known = on the "server side". + // Streams without a known target IP do not emit an event. if let Some(callback) = traffic_callback && let Some(target_ip) = channel.target_ip && !channel.audit_emitted.swap(true, Ordering::SeqCst) @@ -344,7 +382,6 @@ type DataReceiver = mpsc::Receiver; type DataSender = mpsc::Sender; type InternalMessageSender = mpsc::Sender; -#[derive(Debug)] enum InternalMessage { Eof { id: LocalChannelId, @@ -353,7 +390,13 @@ enum InternalMessage { // Boxing reduces enum size from 224 bytes to ~16 bytes // (clippy::large_enum_variant) channel: Box, - stream: TcpStream, + stream: ErasedTargetStream, + }, + StreamResolutionFailed { + id: LocalChannelId, + distant_id: DistantChannelId, + reason_code: ReasonCode, + description: String, }, AbnormalTermination { id: LocalChannelId, @@ -427,6 +470,7 @@ struct JmuxSchedulerTask { msg_to_send_tx: MessageSender, api_request_rx: ApiRequestReceiver, traffic_callback: Option, + target_connector_override: Option, parent_span: Span, } @@ -448,6 +492,7 @@ async fn scheduler_task_impl(task: JmuxSc msg_to_send_tx, mut api_request_rx, traffic_callback, + target_connector_override, parent_span, } = task; @@ -501,7 +546,8 @@ async fn scheduler_task_impl(task: JmuxSc error!(%error, "Couldn't send leftover bytes"); } - let (reader, writer) = stream.into_split(); + let stream = Box::new(stream) as ErasedTargetStream; + let (reader, writer) = tokio::io::split(stream); DataWriterTask { writer, @@ -626,7 +672,7 @@ async fn scheduler_task_impl(task: JmuxSc debug!("Channel accepted"); }); - let (reader, writer) = stream.into_split(); + let (reader, writer) = tokio::io::split(stream); DataWriterTask { writer, @@ -652,6 +698,18 @@ async fn scheduler_task_impl(task: JmuxSc .spawn(channel_span) .detach(); } + InternalMessage::StreamResolutionFailed { + id, + distant_id, + reason_code, + description, + } => { + jmux_ctx.id_allocator.free(id); + msg_to_send_tx + .send(Message::open_failure(distant_id, reason_code, description)) + .await + .context("couldn't send OPEN FAILURE message through mpsc channel")?; + } } } msg = jmux_stream.next() => { @@ -758,8 +816,8 @@ async fn scheduler_task_impl(task: JmuxSc channel, destination_url: msg.destination_url, internal_msg_tx: internal_msg_tx.clone(), - msg_to_send_tx: msg_to_send_tx.clone(), traffic_callback: traffic_callback.clone(), + target_connector_override: target_connector_override.clone(), } .spawn() .detach(); @@ -958,7 +1016,7 @@ async fn scheduler_task_impl(task: JmuxSc // ---------------------- // struct DataReaderTask { - reader: OwnedReadHalf, + reader: tokio::io::ReadHalf, local_id: LocalChannelId, distant_id: DistantChannelId, window_size_updated: Arc, @@ -1087,7 +1145,7 @@ impl DataReaderTask { // ---------------------- // struct DataWriterTask { - writer: OwnedWriteHalf, + writer: tokio::io::WriteHalf, data_rx: DataReceiver, /// Tracks bytes written into the stream. bytes_tx: Arc, @@ -1123,6 +1181,8 @@ impl DataWriterTask { bytes_tx.fetch_add(data.len() as u64, Ordering::SeqCst); } + + let _ = writer.shutdown().await; } .instrument(span), ); @@ -1137,8 +1197,14 @@ struct StreamResolverTask { channel: JmuxChannelCtx, destination_url: DestinationUrl, internal_msg_tx: InternalMessageSender, - msg_to_send_tx: MessageSender, traffic_callback: Option, + target_connector_override: Option, +} + +struct StreamResolutionFailure { + error: anyhow::Error, + reason_code: ReasonCode, + description: String, } impl StreamResolverTask { @@ -1162,96 +1228,130 @@ impl StreamResolverTask { mut channel, destination_url, internal_msg_tx, - msg_to_send_tx, traffic_callback, + target_connector_override, } = self; let scheme = destination_url.scheme(); let host = destination_url.host(); let port = destination_url.port(); - match scheme { - "tcp" => { - // Perform DNS resolution first to get concrete IP addresses. - let socket_addrs = match tokio::net::lookup_host((host, port)).await { - Ok(addrs) => addrs, - Err(error) => { - debug!(?error, "DNS resolution failed"); - // No event emission for DNS failures - cannot determine target IP. - msg_to_send_tx - .send(Message::open_failure( - channel.distant_id, - ReasonCode::from(error.kind()), - error.to_string(), - )) - .await - .context("couldn't send OPEN FAILURE message through mpsc channel")?; - anyhow::bail!("couldn't resolve {host}:{port}: {error}"); - } - }; - - // Try connecting to each resolved address (Happy Eyeballs style). - let mut last_error = None; - - for socket_addr in socket_addrs { - match TcpStream::connect(socket_addr).await { - Ok(stream) => { - // Update channel with resolved target IP and connect time. - channel.target_ip = Some(socket_addr.ip()); - channel.connect_at = SystemTime::now(); - - internal_msg_tx - .send(InternalMessage::StreamResolved { - channel: Box::new(channel), - stream, - }) - .await - .context("couldn't send back resolved stream through internal mpsc channel")?; + let resolution = if scheme != "tcp" { + let description = format!("unsupported scheme: {scheme}"); + Err(StreamResolutionFailure { + error: anyhow::anyhow!(description.clone()), + reason_code: ReasonCode::GENERAL_FAILURE, + description, + }) + } else if let Some(connector) = target_connector_override { + match connector(destination_url.clone()).await { + Ok(Some(stream)) => Ok(stream), + Ok(None) => Self::connect_direct(host, port, &mut channel, traffic_callback.as_ref()).await, + Err(error) => Err(StreamResolutionFailure { + error: error.context(format!("couldn't connect to {host}:{port}")), + reason_code: ReasonCode::GENERAL_FAILURE, + description: "target connection failed".to_owned(), + }), + } + } else { + Self::connect_direct(host, port, &mut channel, traffic_callback.as_ref()).await + }; - return Ok(()); - } - Err(error) => { - debug!(?error, ?socket_addr, "TcpStream::connect failed"); - last_error = Some((socket_addr, error)); - } - } - } + match resolution { + Ok(stream) => { + channel.connect_at = SystemTime::now(); + internal_msg_tx + .send(InternalMessage::StreamResolved { + channel: Box::new(channel), + stream, + }) + .await + .ok() + .context("couldn't send back resolved stream through internal mpsc channel") + } + Err(StreamResolutionFailure { + error, + reason_code, + description, + }) => { + internal_msg_tx + .send(InternalMessage::StreamResolutionFailed { + id: channel.local_id, + distant_id: channel.distant_id, + reason_code, + description, + }) + .await + .ok() + .context("couldn't report stream resolution failure through internal mpsc channel")?; + Err(error) + } + } + } - // All connection attempts failed - emit ConnectFailure for the last attempted address. - if let Some((failed_addr, error)) = last_error { - // Emit ConnectFailure event - we always have a concrete IP at this point. - if let Some(callback) = &traffic_callback { - let connect_and_disconnect_time = SystemTime::now(); - - callback(TrafficEvent { - outcome: EventOutcome::ConnectFailure, - protocol: TransportProtocol::Tcp, - target_host: channel.target_host.clone(), - target_ip: failed_addr.ip(), - target_port: failed_addr.port(), - connect_at: connect_and_disconnect_time, - disconnect_at: connect_and_disconnect_time, - active_duration: std::time::Duration::ZERO, - bytes_tx: 0, - bytes_rx: 0, - }); - } + async fn connect_direct( + host: &str, + port: u16, + channel: &mut JmuxChannelCtx, + traffic_callback: Option<&TrafficCallback>, + ) -> Result { + let socket_addrs = match tokio::net::lookup_host((host, port)).await { + Ok(addrs) => addrs, + Err(error) => { + debug!(?error, "DNS resolution failed"); + // No event emission for DNS failures - cannot determine target IP. + return Err(StreamResolutionFailure { + reason_code: ReasonCode::from(error.kind()), + description: error.to_string(), + error: anyhow::Error::new(error).context(format!("couldn't resolve {host}:{port}")), + }); + } + }; - msg_to_send_tx - .send(Message::open_failure( - channel.distant_id, - ReasonCode::from(error.kind()), - error.to_string(), - )) - .await - .context("couldn't send OPEN FAILURE message through mpsc channel")?; + let mut last_error = None; - anyhow::bail!("couldn't open TCP stream to {host}:{port}: {error}"); - } else { - anyhow::bail!("no addresses resolved for {host}:{port}"); + for socket_addr in socket_addrs { + match TcpStream::connect(socket_addr).await { + Ok(stream) => { + channel.target_ip = Some(socket_addr.ip()); + return Ok(Box::new(stream)); + } + Err(error) => { + debug!(?error, ?socket_addr, "TcpStream::connect failed"); + last_error = Some((socket_addr, error)); } } - _ => anyhow::bail!("unsupported scheme: {scheme}"), + } + + if let Some((failed_addr, error)) = last_error { + if let Some(callback) = traffic_callback { + let connect_and_disconnect_time = SystemTime::now(); + + callback(TrafficEvent { + outcome: EventOutcome::ConnectFailure, + protocol: TransportProtocol::Tcp, + target_host: channel.target_host.clone(), + target_ip: failed_addr.ip(), + target_port: failed_addr.port(), + connect_at: connect_and_disconnect_time, + disconnect_at: connect_and_disconnect_time, + active_duration: std::time::Duration::ZERO, + bytes_tx: 0, + bytes_rx: 0, + }); + } + + Err(StreamResolutionFailure { + reason_code: ReasonCode::from(error.kind()), + description: error.to_string(), + error: anyhow::Error::new(error).context(format!("couldn't open TCP stream to {host}:{port}")), + }) + } else { + Err(StreamResolutionFailure { + error: anyhow::anyhow!("no addresses resolved for {host}:{port}"), + reason_code: ReasonCode::GENERAL_FAILURE, + description: "no addresses resolved".to_owned(), + }) } } } diff --git a/crates/jmux-proxy/tests/target_connector.rs b/crates/jmux-proxy/tests/target_connector.rs new file mode 100644 index 000000000..9f7f36655 --- /dev/null +++ b/crates/jmux-proxy/tests/target_connector.rs @@ -0,0 +1,167 @@ +use std::time::Duration; + +use jmux_proto::{Bytes, BytesMut, DistantChannelId, Header, LocalChannelId, Message, ReasonCode}; +use jmux_proxy::{DestinationUrl, JmuxConfig, JmuxProxy, TrafficEvent}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::sync::mpsc; +use tokio::time::timeout; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); +const NO_EVENT_TIMEOUT: Duration = Duration::from_millis(100); + +async fn send_message(writer: &mut (impl AsyncWrite + Unpin), message: Message) { + let mut bytes = BytesMut::new(); + message.encode(&mut bytes).expect("encode JMUX message"); + writer.write_all(&bytes).await.expect("send JMUX message"); +} + +async fn receive_message(reader: &mut (impl AsyncRead + Unpin)) -> Message { + timeout(TEST_TIMEOUT, async { + let mut header = [0; Header::SIZE]; + reader.read_exact(&mut header).await.expect("read JMUX header"); + let message_size = usize::from(u16::from_be_bytes([header[1], header[2]])); + let mut body = vec![0; message_size - Header::SIZE]; + reader.read_exact(&mut body).await.expect("read JMUX body"); + + let mut bytes = BytesMut::with_capacity(message_size); + bytes.extend_from_slice(&header); + bytes.extend_from_slice(&body); + Message::decode(bytes.freeze()).expect("decode JMUX message") + }) + .await + .expect("JMUX response timed out") +} + +async fn assert_no_traffic_event(receiver: &mut mpsc::UnboundedReceiver) { + match timeout(NO_EVENT_TIMEOUT, receiver.recv()).await { + Err(_) => {} + Ok(Some(event)) => panic!("unexpected traffic event: {event:?}"), + Ok(None) => panic!("traffic event callback dropped"), + } +} + +#[tokio::test] +async fn override_stream_carries_channel_data() { + let (proxy_stream, peer_stream) = tokio::io::duplex(8192); + let (proxy_reader, proxy_writer) = tokio::io::split(proxy_stream); + let (mut peer_reader, mut peer_writer) = tokio::io::split(peer_stream); + let (traffic_tx, mut traffic_rx) = mpsc::unbounded_channel(); + let proxy = JmuxProxy::new(Box::new(proxy_reader), Box::new(proxy_writer)) + .with_config(JmuxConfig::permissive()) + .with_outgoing_traffic_event_callback(move |event| { + traffic_tx.send(event).expect("capture traffic event"); + }) + .with_target_connector_override(|_| async move { + let (target_stream, mut target_peer) = tokio::io::duplex(64); + tokio::spawn(async move { + let mut payload = [0; 4]; + target_peer.read_exact(&mut payload).await.expect("read target data"); + let mut eof_probe = [0u8; 1]; + assert_eq!(target_peer.read(&mut eof_probe).await.expect("read target EOF"), 0); + target_peer.write_all(&payload).await.expect("echo target data"); + }); + Ok(Some(target_stream)) + }); + let _proxy_task = tokio::spawn(proxy.run()); + + send_message( + &mut peer_writer, + Message::open( + LocalChannelId::from(7), + 4096, + DestinationUrl::new("tcp", "agent.example", 443), + ), + ) + .await; + + let Message::OpenSuccess(open_success) = receive_message(&mut peer_reader).await else { + panic!("expected OPEN SUCCESS"); + }; + let local_id = DistantChannelId::from(open_success.sender_channel_id); + + send_message(&mut peer_writer, Message::data(local_id, Bytes::from_static(b"ping"))).await; + send_message(&mut peer_writer, Message::eof(local_id)).await; + let Message::Data(data) = receive_message(&mut peer_reader).await else { + panic!("expected CHANNEL DATA"); + }; + assert_eq!(data.recipient_channel_id, 7); + assert_eq!(data.transfer_data, b"ping"[..]); + + let Message::Close(close) = receive_message(&mut peer_reader).await else { + panic!("expected CHANNEL CLOSE"); + }; + assert_eq!(close.recipient_channel_id, 7); + send_message(&mut peer_writer, Message::close(local_id)).await; + assert_no_traffic_event(&mut traffic_rx).await; +} + +#[tokio::test] +async fn resolution_failures_free_id_and_keep_direct_fallback() { + let (proxy_stream, peer_stream) = tokio::io::duplex(8192); + let (proxy_reader, proxy_writer) = tokio::io::split(proxy_stream); + let (mut peer_reader, mut peer_writer) = tokio::io::split(peer_stream); + let (traffic_tx, mut traffic_rx) = mpsc::unbounded_channel(); + let proxy = JmuxProxy::new(Box::new(proxy_reader), Box::new(proxy_writer)) + .with_config(JmuxConfig::permissive()) + .with_outgoing_traffic_event_callback(move |event| { + traffic_tx.send(event).expect("capture traffic event"); + }) + .with_target_connector_override(|destination| async move { + if destination.host() == "fail.example" { + anyhow::bail!("agent error"); + } + Ok(None::) + }); + let _proxy_task = tokio::spawn(proxy.run()); + + send_message( + &mut peer_writer, + Message::open( + LocalChannelId::from(11), + 4096, + DestinationUrl::new("tcp", "fail.example", 443), + ), + ) + .await; + + let Message::OpenFailure(open_failure) = receive_message(&mut peer_reader).await else { + panic!("expected OPEN FAILURE"); + }; + assert_eq!(open_failure.reason_code, ReasonCode::GENERAL_FAILURE); + assert_eq!(open_failure.description, "target connection failed"); + assert_no_traffic_event(&mut traffic_rx).await; + + send_message( + &mut peer_writer, + Message::open( + LocalChannelId::from(12), + 4096, + DestinationUrl::new("tcp", "127.0.0.1", 0), + ), + ) + .await; + assert!(matches!( + receive_message(&mut peer_reader).await, + Message::OpenFailure(_) + )); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind direct target"); + let target_port = listener.local_addr().expect("read direct target address").port(); + + send_message( + &mut peer_writer, + Message::open( + LocalChannelId::from(13), + 4096, + DestinationUrl::new("tcp", "127.0.0.1", target_port), + ), + ) + .await; + let Message::OpenSuccess(open_success) = receive_message(&mut peer_reader).await else { + panic!("expected OPEN SUCCESS"); + }; + // The failed channel released ID 0 for reuse. + assert_eq!(open_success.sender_channel_id, 0); +} diff --git a/devolutions-agent/src/config.rs b/devolutions-agent/src/config.rs index 81755248a..fb04d203b 100644 --- a/devolutions-agent/src/config.rs +++ b/devolutions-agent/src/config.rs @@ -1,11 +1,14 @@ use std::fs::File; use std::io::BufReader; -use std::net::SocketAddr; +use std::net::{Ipv6Addr, SocketAddr}; +use std::num::NonZeroU16; use std::sync::Arc; +use agent_tunnel_proto::DomainName; use anyhow::{Context as _, bail}; use camino::{Utf8Path, Utf8PathBuf}; use devolutions_agent_shared::{default_schedule_window_start, get_data_dir}; +use ipnetwork::Ipv4Network; use serde::{Deserialize, Serialize}; use tap::prelude::*; use url::Url; @@ -27,19 +30,22 @@ pub struct Conf { pub debug: dto::DebugConf, } -/// Validated tunnel configuration — required fields are guaranteed present. +/// Validated tunnel configuration. /// -/// Constructed from `dto::TunnelConf` via `TryFrom`. If the tunnel is disabled -/// or not yet enrolled, the `enabled` field is `false` and path fields are empty -/// (but the struct is always constructible). +/// Required fields and field values are validated when this is constructed from [`dto::TunnelConf`]. #[derive(Debug, Clone)] -pub struct TunnelConf { - pub enabled: bool, - pub gateway_endpoint: String, +pub enum TunnelConf { + Disabled, + Enabled(Box), +} + +#[derive(Debug, Clone)] +pub struct EnabledTunnelConf { + pub gateway_endpoint: GatewayEndpoint, pub client_cert_path: Utf8PathBuf, pub client_key_path: Utf8PathBuf, pub gateway_ca_cert_path: Utf8PathBuf, - pub advertise_subnets: Vec, + pub advertise_subnets: Vec, pub advertise_domains: Vec, pub auto_detect_domain: bool, pub heartbeat_interval_secs: u64, @@ -47,59 +53,165 @@ pub struct TunnelConf { pub server_spki_sha256: Option, } -impl TryFrom for TunnelConf { - type Error = anyhow::Error; +#[derive(Debug, Clone)] +pub struct GatewayEndpoint { + host: String, + port: NonZeroU16, +} - fn try_from(conf: dto::TunnelConf) -> anyhow::Result { - if !conf.enabled { - // Disabled tunnel — return a placeholder with defaults. +impl GatewayEndpoint { + pub fn host(&self) -> &str { + &self.host + } + + pub fn port(&self) -> u16 { + self.port.get() + } +} + +impl std::str::FromStr for GatewayEndpoint { + type Err = anyhow::Error; + + fn from_str(endpoint: &str) -> anyhow::Result { + anyhow::ensure!(!endpoint.is_empty(), "value is required when Tunnel.Enabled is true"); + anyhow::ensure!( + endpoint.trim() == endpoint, + "value must not contain leading or trailing whitespace" + ); + + if let Ok(socket_addr) = endpoint.parse::() { + let port = NonZeroU16::new(socket_addr.port()).context("port must be greater than zero")?; return Ok(Self { - enabled: false, - gateway_endpoint: String::new(), - client_cert_path: Utf8PathBuf::new(), - client_key_path: Utf8PathBuf::new(), - gateway_ca_cert_path: Utf8PathBuf::new(), - advertise_subnets: Vec::new(), - advertise_domains: Vec::new(), - auto_detect_domain: true, - heartbeat_interval_secs: 60, - route_advertise_interval_secs: 30, - server_spki_sha256: None, + host: socket_addr.ip().to_string(), + port, }); } - // Enabled tunnel — all required fields must be present. - let client_cert_path = conf - .client_cert_path - .context("tunnel enabled but client_cert_path not configured")?; - let client_key_path = conf - .client_key_path - .context("tunnel enabled but client_key_path not configured")?; - let gateway_ca_cert_path = conf - .gateway_ca_cert_path - .context("tunnel enabled but gateway_ca_cert_path not configured")?; + let (hostname, port) = endpoint + .rsplit_once(':') + .context("expected an endpoint in host:port format")?; + anyhow::ensure!(!hostname.is_empty(), "hostname must not be empty"); + anyhow::ensure!( + hostname.parse::().is_ok() + || !hostname.chars().any(|character| matches!(character, ':' | '[' | ']')), + "IPv6 addresses must use bracketed host:port notation" + ); + + rustls_pki_types::ServerName::try_from(hostname.to_owned()) + .map_err(|_| anyhow::anyhow!("invalid hostname `{hostname}`"))?; + + let port = port + .parse::() + .with_context(|| format!("invalid port `{port}`"))?; + + Ok(Self { + host: hostname.to_owned(), + port, + }) + } +} + +impl TunnelConf { + pub fn is_enabled(&self) -> bool { + matches!(self, Self::Enabled(_)) + } + + pub(crate) fn as_enabled(&self) -> Option<&EnabledTunnelConf> { + match self { + Self::Disabled => None, + Self::Enabled(conf) => Some(conf), + } + } + + pub(crate) fn from_dto(conf: dto::TunnelConf) -> anyhow::Result { + if !conf.enabled { + return Ok(Self::Disabled); + } + EnabledTunnelConf::from_dto(conf).map(Box::new).map(Self::Enabled) + } +} + +impl EnabledTunnelConf { + fn from_dto(conf: dto::TunnelConf) -> anyhow::Result { + let gateway_endpoint = conf + .gateway_endpoint + .parse() + .context("invalid Tunnel.GatewayEndpoint")?; + + let client_cert_path = required_tunnel_path(conf.client_cert_path).context("invalid Tunnel.ClientCertPath")?; + let client_key_path = required_tunnel_path(conf.client_key_path).context("invalid Tunnel.ClientKeyPath")?; + let gateway_ca_cert_path = + required_tunnel_path(conf.gateway_ca_cert_path).context("invalid Tunnel.GatewayCaCertPath")?; + + let advertise_subnets = conf + .advertise_subnets + .into_iter() + .enumerate() + .map(|(index, subnet)| { + subnet + .parse::() + .with_context(|| format!("invalid Tunnel.AdvertiseSubnets[{index}] value `{subnet}`")) + }) + .collect::>>()?; + + for (index, route) in conf.advertise_domains.iter().enumerate() { + anyhow::ensure!( + DomainName::is_valid_route(route), + "invalid Tunnel.AdvertiseDomains[{index}] value `{route}`" + ); + } + + let heartbeat_interval_secs = conf.heartbeat_interval_secs.unwrap_or(60); + anyhow::ensure!( + heartbeat_interval_secs > 0, + "invalid Tunnel.HeartbeatIntervalSecs: value must be greater than zero" + ); + + let route_advertise_interval_secs = conf.route_advertise_interval_secs.unwrap_or(30); anyhow::ensure!( - !conf.gateway_endpoint.is_empty(), - "tunnel enabled but gateway_endpoint is empty" + route_advertise_interval_secs > 0, + "invalid Tunnel.RouteAdvertiseIntervalSecs: value must be greater than zero" + ); + anyhow::ensure!( + heartbeat_interval_secs.min(route_advertise_interval_secs) + <= agent_tunnel_proto::AGENT_OFFLINE_TIMEOUT_SECS / 3, + "invalid Tunnel.HeartbeatIntervalSecs and Tunnel.RouteAdvertiseIntervalSecs: \ + at least one value must be at most {} seconds", + agent_tunnel_proto::AGENT_OFFLINE_TIMEOUT_SECS / 3 ); + let server_spki_sha256 = conf + .server_spki_sha256 + .map(|hash| { + anyhow::ensure!( + hash.len() == 64 && hash.bytes().all(|byte| byte.is_ascii_hexdigit()), + "invalid Tunnel.ServerSpkiSha256: expected 64 hexadecimal characters" + ); + Ok(hash.to_ascii_lowercase()) + }) + .transpose()?; Ok(Self { - enabled: true, - gateway_endpoint: conf.gateway_endpoint, + gateway_endpoint, client_cert_path, client_key_path, gateway_ca_cert_path, - advertise_subnets: conf.advertise_subnets, + advertise_subnets, advertise_domains: conf.advertise_domains, auto_detect_domain: conf.auto_detect_domain, - heartbeat_interval_secs: conf.heartbeat_interval_secs.unwrap_or(60), - route_advertise_interval_secs: conf.route_advertise_interval_secs.unwrap_or(30), - server_spki_sha256: conf.server_spki_sha256, + heartbeat_interval_secs, + route_advertise_interval_secs, + server_spki_sha256, }) } } +fn required_tunnel_path(path: Option) -> anyhow::Result { + let path = path.context("value is required when Tunnel.Enabled is true")?; + anyhow::ensure!(!path.as_str().trim().is_empty(), "path must not be empty"); + Ok(path) +} + /// Validated PSU agent configuration. /// /// Constructed from `dto::PsuConf` via `TryFrom for Option`. @@ -178,7 +290,7 @@ impl Conf { .tunnel .clone() .unwrap_or_default() - .pipe(TunnelConf::try_from) + .pipe(TunnelConf::from_dto) .context("invalid tunnel config")?, proxy: conf_file.proxy.clone().unwrap_or_default(), debug: conf_file.debug.clone().unwrap_or_default(), @@ -954,6 +1066,112 @@ pub fn handle_cli(command: &str) -> Result<(), anyhow::Error> { mod tests { use super::*; + fn valid_tunnel_json() -> serde_json::Value { + serde_json::json!({ + "Enabled": true, + "GatewayEndpoint": "[::1]:4433", + "ClientCertPath": "client.crt", + "ClientKeyPath": "client.key", + "GatewayCaCertPath": "gateway-ca.crt", + "HeartbeatIntervalSecs": 60, + "ServerSpkiSha256": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" + }) + } + + fn load_tunnel_json(tunnel: serde_json::Value) -> anyhow::Result { + let conf_file = serde_json::from_value(serde_json::json!({ "Tunnel": tunnel })) + .context("deserialize test configuration")?; + Conf::from_conf_file(&conf_file) + } + + #[test] + fn tunnel_config_normalizes_spki_and_hostname() { + let conf = load_tunnel_json(valid_tunnel_json()).expect("load valid tunnel configuration"); + let tunnel = conf.tunnel.as_enabled().expect("tunnel enabled"); + + assert_eq!(tunnel.gateway_endpoint.host(), "::1"); + assert_eq!( + tunnel.server_spki_sha256.as_deref(), + Some("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") + ); + } + + #[test] + fn tunnel_config_accepts_legacy_unbracketed_ipv6_endpoint() { + let mut tunnel = valid_tunnel_json(); + tunnel["GatewayEndpoint"] = serde_json::json!("::1:4433"); + + let conf = load_tunnel_json(tunnel).expect("load legacy tunnel endpoint"); + let tunnel = conf.tunnel.as_enabled().expect("tunnel enabled"); + + assert_eq!(tunnel.gateway_endpoint.host(), "::1"); + assert_eq!(tunnel.gateway_endpoint.port(), 4433); + } + + #[test] + fn tunnel_config_errors_identify_invalid_fields() { + let cases = [ + ( + "GatewayEndpoint", + serde_json::json!("gateway.example.com"), + "invalid Tunnel.GatewayEndpoint", + ), + ("ClientKeyPath", serde_json::json!(""), "invalid Tunnel.ClientKeyPath"), + ( + "AdvertiseDomains", + serde_json::json!(["invalid.*.example.com"]), + "invalid Tunnel.AdvertiseDomains[0]", + ), + ( + "HeartbeatIntervalSecs", + serde_json::json!(0), + "invalid Tunnel.HeartbeatIntervalSecs", + ), + ( + "ServerSpkiSha256", + serde_json::json!("not-a-sha256-hash"), + "invalid Tunnel.ServerSpkiSha256", + ), + ]; + + for (field, value, expected) in cases { + let mut tunnel = valid_tunnel_json(); + tunnel[field] = value; + + let error = match load_tunnel_json(tunnel) { + Ok(_) => panic!("invalid {field} should fail loading"), + Err(error) => error, + }; + let error = format!("{error:#}"); + assert!(error.contains(expected), "expected `{expected}` in `{error}`"); + } + } + + #[test] + fn disabled_tunnel_skips_validation() { + let conf = load_tunnel_json(serde_json::json!({ + "Enabled": false, + "GatewayEndpoint": "invalid" + })) + .expect("load disabled tunnel configuration"); + + assert!(matches!(conf.tunnel, TunnelConf::Disabled)); + } + + #[test] + fn tunnel_config_requires_liveness_margin() { + let mut tunnel = valid_tunnel_json(); + let invalid_interval = agent_tunnel_proto::AGENT_OFFLINE_TIMEOUT_SECS / 3 + 1; + tunnel["HeartbeatIntervalSecs"] = serde_json::json!(invalid_interval); + tunnel["RouteAdvertiseIntervalSecs"] = serde_json::json!(invalid_interval); + + let error = load_tunnel_json(tunnel).expect_err("stale liveness intervals should fail loading"); + + assert!( + format!("{error:#}").contains("invalid Tunnel.HeartbeatIntervalSecs and Tunnel.RouteAdvertiseIntervalSecs") + ); + } + #[test] fn psu_config_deserializes() { let conf_file: dto::ConfFile = serde_json::from_value(serde_json::json!({ diff --git a/devolutions-agent/src/service.rs b/devolutions-agent/src/service.rs index 631d793b4..6739ea391 100644 --- a/devolutions-agent/src/service.rs +++ b/devolutions-agent/src/service.rs @@ -258,7 +258,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { tasks.register(RemoteDesktopTask::new(conf_handle.clone())); } - if conf.tunnel.enabled { + if conf.tunnel.is_enabled() { tasks.register(TunnelTask::new(conf_handle.clone())); } diff --git a/devolutions-agent/src/tunnel.rs b/devolutions-agent/src/tunnel.rs index e19c69d65..57c36f9d5 100644 --- a/devolutions-agent/src/tunnel.rs +++ b/devolutions-agent/src/tunnel.rs @@ -16,8 +16,9 @@ use async_trait::async_trait; use devolutions_gateway_task::{ShutdownSignal, Task}; use ipnetwork::Ipv4Network; use sha2::Digest as _; +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt as _}; -use crate::config::ConfHandle; +use crate::config::{ConfHandle, EnabledTunnelConf}; use crate::tunnel_helpers::{Target, connect_to_target, resolve_target}; // --------------------------------------------------------------------------- @@ -251,31 +252,18 @@ async fn run_single_connection( shutdown_signal: &mut ShutdownSignal, ) -> anyhow::Result { let agent_conf = conf_handle.get_conf(); - let tunnel_conf = &agent_conf.tunnel; + let tunnel_conf = agent_conf.tunnel.as_enabled().context("agent tunnel is not enabled")?; let cert_path = &tunnel_conf.client_cert_path; let key_path = &tunnel_conf.client_key_path; let ca_path = &tunnel_conf.gateway_ca_cert_path; - let advertise_subnets: Vec = tunnel_conf - .advertise_subnets - .iter() - .map(|subnet| subnet.parse()) - .collect::, _>>() - .context("failed to parse advertise_subnets")?; + let advertise_subnets = tunnel_conf.advertise_subnets.clone(); if advertise_subnets.is_empty() { warn!("No subnets configured to advertise"); } - if let Some(route) = tunnel_conf - .advertise_domains - .iter() - .find(|route| !DomainName::is_valid_route(route)) - { - bail!("invalid advertise domain route: {route}"); - } - let detected_domain = if tunnel_conf.auto_detect_domain && tunnel_conf.advertise_domains.is_empty() { crate::domain_detect::detect_domain() } else { @@ -383,9 +371,7 @@ async fn run_single_connection( /// Build the mTLS client config, resolve the gateway endpoint, and perform the /// QUIC handshake, returning the live endpoint and connection. -async fn connect_to_gateway( - tunnel_conf: &crate::config::TunnelConf, -) -> anyhow::Result<(quinn::Endpoint, quinn::Connection)> { +async fn connect_to_gateway(tunnel_conf: &EnabledTunnelConf) -> anyhow::Result<(quinn::Endpoint, quinn::Connection)> { // Ensure rustls crypto provider is installed (ring). let _ = rustls::crypto::ring::default_provider().install_default(); @@ -454,16 +440,14 @@ async fn connect_to_gateway( // -- DNS resolve -- // Extract hostname for TLS server name validation. - let (gateway_hostname, _) = tunnel_conf - .gateway_endpoint - .rsplit_once(':') - .context("gateway_endpoint missing port separator")?; + let gateway_hostname = tunnel_conf.gateway_endpoint.host(); - let gateway_addr = tokio::net::lookup_host(&tunnel_conf.gateway_endpoint) - .await - .context("failed to resolve gateway endpoint")? - .next() - .context("no addresses resolved for gateway endpoint")?; + let gateway_addr = + tokio::net::lookup_host((tunnel_conf.gateway_endpoint.host(), tunnel_conf.gateway_endpoint.port())) + .await + .context("failed to resolve gateway endpoint")? + .next() + .context("no addresses resolved for gateway endpoint")?; info!(gateway_addr = %gateway_addr, %gateway_hostname, "Connecting to gateway"); @@ -504,9 +488,7 @@ async fn connect_to_gateway( /// wait for the gateway to reply with a Version Negotiation packet. That reply proves UDP/4433 /// reaches the gateway while creating zero connection state on it, and needs no client cert. pub async fn probe_connectivity(tunnel_conf: &crate::config::TunnelConf, timeout: Duration) -> anyhow::Result<()> { - if !tunnel_conf.enabled { - bail!("agent tunnel is not enabled"); - } + let tunnel_conf = tunnel_conf.as_enabled().context("agent tunnel is not enabled")?; // The whole probe — DNS resolution, socket setup, and the retransmit loop — is bounded by // `timeout`, so a stalled resolver or a black-holed path can't hang past it. @@ -516,12 +498,13 @@ pub async fn probe_connectivity(tunnel_conf: &crate::config::TunnelConf, timeout } } -async fn reach_gateway(tunnel_conf: &crate::config::TunnelConf) -> anyhow::Result<()> { - let gateway_addr = tokio::net::lookup_host(&tunnel_conf.gateway_endpoint) - .await - .context("failed to resolve gateway endpoint")? - .next() - .context("no addresses resolved for gateway endpoint")?; +async fn reach_gateway(tunnel_conf: &EnabledTunnelConf) -> anyhow::Result<()> { + let gateway_addr = + tokio::net::lookup_host((tunnel_conf.gateway_endpoint.host(), tunnel_conf.gateway_endpoint.port())) + .await + .context("failed to resolve gateway endpoint")? + .next() + .context("no addresses resolved for gateway endpoint")?; // Match the local bind family to the resolved gateway address (see connect_to_gateway). let bind_addr: SocketAddr = if gateway_addr.is_ipv4() { @@ -598,8 +581,8 @@ async fn try_renew_certificate( ca_path: &camino::Utf8Path, ) -> anyhow::Result> where - S: tokio::io::AsyncWrite + Unpin, - R: tokio::io::AsyncRead + Unpin, + S: AsyncWrite + Unpin, + R: AsyncRead + Unpin, { const RENEWAL_THRESHOLD_DAYS: u32 = 15; const RENEWAL_TIMEOUT: Duration = Duration::from_secs(30); @@ -675,7 +658,7 @@ where // Control stream reader // --------------------------------------------------------------------------- -async fn run_control_reader(mut ctrl: FramedRecv) { +async fn run_control_reader(mut ctrl: FramedRecv) { let _ = async move { loop { let message: ControlMessage = ctrl.recv().await.context("recv control message")?; @@ -717,6 +700,35 @@ async fn run_control_reader(mut ctrl: FramedRec /// black-holed target outlives its deadline and it never hears why we failed. const CONNECT_DEADLINE: Duration = Duration::from_secs(20); +async fn proxy_session_traffic( + (mut tunnel_send, mut tunnel_recv): (impl AsyncWrite + Unpin, impl AsyncRead + Unpin), + (mut target_read, mut target_write): (impl AsyncRead + Unpin, impl AsyncWrite + Unpin), +) -> anyhow::Result<()> { + let tunnel_to_target = async { + tokio::io::copy(&mut tunnel_recv, &mut target_write) + .await + .context("proxy QUIC to TCP")?; + if let Err(error) = target_write.shutdown().await { + debug!(%error, "TCP write shutdown failed"); + } + Ok::<_, anyhow::Error>(()) + }; + let target_to_tunnel = async { + tokio::io::copy(&mut target_read, &mut tunnel_send) + .await + .context("proxy TCP to QUIC")?; + if let Err(error) = tunnel_send.shutdown().await { + debug!(%error, "QUIC send shutdown failed"); + } + Ok::<_, anyhow::Error>(()) + }; + + // Keep both directions alive so one-sided EOF propagates as a half-close without cancelling reverse traffic. + let (tunnel_to_target, target_to_tunnel) = tokio::join!(tunnel_to_target, target_to_tunnel); + tunnel_to_target?; + target_to_tunnel +} + async fn run_session_proxy( advertise_subnets: Vec, advertise_domains: Vec, @@ -785,18 +797,7 @@ async fn run_session_proxy( let (mut send, mut recv) = session.into_inner(); let (mut tcp_read, mut tcp_write) = tcp_stream.into_split(); - - // Use join! (not select!) to wait for BOTH directions to finish. - // select! would cancel in-flight data when one direction closes first. - let (r1, r2) = tokio::join!( - tokio::io::copy(&mut recv, &mut tcp_write), - tokio::io::copy(&mut tcp_read, &mut send), - ); - r1.inspect_err(|e| debug!(%e, "QUIC->TCP copy ended"))?; - r2.inspect_err(|e| debug!(%e, "TCP->QUIC copy ended"))?; - - // Gracefully finish the QUIC send stream (signals EOF to peer). - let _ = send.finish(); + proxy_session_traffic((&mut send, &mut recv), (&mut tcp_read, &mut tcp_write)).await?; Ok(()) } @@ -807,30 +808,105 @@ async fn run_session_proxy( #[cfg(test)] mod tests { use camino::Utf8PathBuf; + use tokio::io::AsyncReadExt as _; use super::*; - use crate::config::TunnelConf; + use crate::config::{TunnelConf, dto}; - fn tunnel_conf_template() -> TunnelConf { - TunnelConf { + async fn tcp_pair() -> (tokio::net::TcpStream, tokio::net::TcpStream) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind TCP listener"); + let client = tokio::net::TcpStream::connect(listener.local_addr().expect("read listener address")) + .await + .expect("connect TCP client"); + let (server, _) = listener.accept().await.expect("accept TCP client"); + (client, server) + } + + async fn spawn_relay() -> ( + tokio::io::DuplexStream, + tokio::net::TcpStream, + tokio::task::JoinHandle>, + ) { + let (tunnel, gateway) = tokio::io::duplex(64); + let (tunnel_recv, tunnel_send) = tokio::io::split(tunnel); + let (target, peer) = tcp_pair().await; + let (target_read, target_write) = target.into_split(); + let relay = tokio::spawn(proxy_session_traffic( + (tunnel_send, tunnel_recv), + (target_read, target_write), + )); + (gateway, peer, relay) + } + + #[tokio::test] + async fn tunnel_eof_half_closes_target_and_preserves_response() { + tokio::time::timeout(Duration::from_secs(5), async { + let (mut gateway, mut peer, relay) = spawn_relay().await; + + gateway.write_all(b"request").await.expect("write tunnel request"); + gateway.shutdown().await.expect("finish tunnel request"); + + let mut request = [0; 7]; + peer.read_exact(&mut request).await.expect("read target request"); + assert_eq!(&request, b"request"); + let mut eof_probe = [0u8; 1]; + assert_eq!(peer.read(&mut eof_probe).await.expect("read target EOF"), 0); + + peer.write_all(b"response").await.expect("write target response"); + peer.shutdown().await.expect("finish target response"); + + let mut response = [0; 8]; + gateway.read_exact(&mut response).await.expect("read tunnel response"); + assert_eq!(&response, b"response"); + relay.await.expect("relay task panicked").expect("relay traffic"); + }) + .await + .expect("tunnel EOF test timed out"); + } + + #[tokio::test] + async fn target_eof_finishes_tunnel_send_and_preserves_request() { + tokio::time::timeout(Duration::from_secs(5), async { + let (mut gateway, mut peer, relay) = spawn_relay().await; + + peer.write_all(b"response").await.expect("write target response"); + peer.shutdown().await.expect("finish target response"); + + let mut response = [0; 8]; + gateway.read_exact(&mut response).await.expect("read tunnel response"); + assert_eq!(&response, b"response"); + let mut eof_probe = [0u8; 1]; + assert_eq!(gateway.read(&mut eof_probe).await.expect("read tunnel EOF"), 0); + + gateway.write_all(b"request").await.expect("write tunnel request"); + gateway.shutdown().await.expect("finish tunnel request"); + + let mut request = [0; 7]; + peer.read_exact(&mut request).await.expect("read target request"); + assert_eq!(&request, b"request"); + relay.await.expect("relay task panicked").expect("relay traffic"); + }) + .await + .expect("target EOF test timed out"); + } + + fn tunnel_conf(endpoint: impl Into) -> TunnelConf { + let conf = dto::TunnelConf { enabled: true, - gateway_endpoint: String::new(), - client_cert_path: Utf8PathBuf::new(), - client_key_path: Utf8PathBuf::new(), - gateway_ca_cert_path: Utf8PathBuf::new(), - advertise_subnets: Vec::new(), - advertise_domains: Vec::new(), - auto_detect_domain: false, - heartbeat_interval_secs: 15, - route_advertise_interval_secs: 60, - server_spki_sha256: None, - } + gateway_endpoint: endpoint.into(), + client_cert_path: Some(Utf8PathBuf::from("client.crt")), + client_key_path: Some(Utf8PathBuf::from("client.key")), + gateway_ca_cert_path: Some(Utf8PathBuf::from("gateway-ca.crt")), + ..dto::TunnelConf::default() + }; + TunnelConf::from_dto(conf).expect("validate tunnel configuration") } #[tokio::test] async fn probe_fails_fast_when_tunnel_disabled() { - let mut conf = tunnel_conf_template(); - conf.enabled = false; + let conf = TunnelConf::from_dto(dto::TunnelConf::default()).expect("validate disabled tunnel configuration"); let error = probe_connectivity(&conf, Duration::from_millis(200)) .await @@ -852,8 +928,7 @@ mod tests { .expect("bind blackhole socket"); let blackhole_addr = blackhole.local_addr().expect("blackhole addr"); - let mut conf = tunnel_conf_template(); - conf.gateway_endpoint = blackhole_addr.to_string(); + let conf = tunnel_conf(blackhole_addr.to_string()); let started = std::time::Instant::now(); let result = probe_connectivity(&conf, Duration::from_millis(300)).await; diff --git a/devolutions-gateway/src/api/jmux.rs b/devolutions-gateway/src/api/jmux.rs index b9199c01f..e40bcf3e6 100644 --- a/devolutions-gateway/src/api/jmux.rs +++ b/devolutions-gateway/src/api/jmux.rs @@ -22,6 +22,7 @@ pub async fn handler( shutdown_signal, conf_handle, traffic_audit_handle, + agent_tunnel_handle, .. }): State, JmuxToken(claims): JmuxToken, @@ -35,6 +36,7 @@ pub async fn handler( sessions, subscriber_tx, traffic_audit_handle, + agent_tunnel_handle, claims, source_addr, Duration::from_secs(conf_handle.get_conf().debug.ws_keep_alive_interval), @@ -54,6 +56,7 @@ async fn handle_socket( sessions: SessionMessageSender, subscriber_tx: SubscriberSender, traffic_audit_handle: TrafficAuditHandle, + agent_tunnel_handle: Option>, claims: JmuxTokenClaims, source_addr: SocketAddr, keep_alive_interval: Duration, @@ -65,9 +68,16 @@ async fn handle_socket( ); let session_id = claims.jet_aid; - let result = crate::jmux::handle(stream, claims, sessions, subscriber_tx, traffic_audit_handle) - .instrument(info_span!("jmux", client = %source_addr, %session_id)) - .await; + let result = crate::jmux::handle( + stream, + claims, + sessions, + subscriber_tx, + traffic_audit_handle, + agent_tunnel_handle, + ) + .instrument(info_span!("jmux", client = %source_addr, %session_id)) + .await; if let Err(error) = result { close_handle.server_error("JMUX failure".to_owned()).await; diff --git a/devolutions-gateway/src/jmux.rs b/devolutions-gateway/src/jmux.rs index 500ec0963..a2c8ea2fd 100644 --- a/devolutions-gateway/src/jmux.rs +++ b/devolutions-gateway/src/jmux.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use anyhow::Context as _; use devolutions_gateway_task::ChildTask; -use jmux_proxy::{FilteringRule, JmuxConfig, JmuxProxy}; +use jmux_proxy::{DestinationUrl, FilteringRule, JmuxConfig, JmuxProxy}; use tap::prelude::*; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::sync::Notify; @@ -10,8 +10,10 @@ use transport::{ErasedRead, ErasedWrite}; use crate::session::{ConnectionModeDetails, SessionInfo, SessionMessageSender}; use crate::subscriber::SubscriberSender; +use crate::target_addr::TargetAddr; use crate::token::{JmuxTokenClaims, RecordingPolicy}; use crate::traffic_audit::TrafficAuditHandle; +use crate::upstream::route_target_from_target_addr; pub async fn handle( stream: impl AsyncRead + AsyncWrite + Send + 'static, @@ -19,6 +21,7 @@ pub async fn handle( sessions: SessionMessageSender, subscriber_tx: SubscriberSender, traffic_audit_handle: TrafficAuditHandle, + agent_tunnel_handle: Option>, ) -> anyhow::Result<()> { match claims.jet_rec { RecordingPolicy::None | RecordingPolicy::Stream => (), @@ -105,10 +108,43 @@ pub async fn handle( }); }; - let proxy_fut = JmuxProxy::new(reader, writer) + let mut proxy = JmuxProxy::new(reader, writer) .with_config(config) - .with_outgoing_traffic_event_callback(traffic_event_callback) - .run(); + .with_outgoing_traffic_event_callback(traffic_event_callback); + + if let Some(agent_tunnel_handle) = agent_tunnel_handle { + proxy = proxy.with_target_connector_override(move |destination_url: DestinationUrl| { + let agent_tunnel_handle = Arc::clone(&agent_tunnel_handle); + + async move { + let target = TargetAddr::from_components( + destination_url.scheme(), + destination_url.host(), + destination_url.port(), + ) + .context("invalid JMUX target")?; + let route_target = route_target_from_target_addr(&target); + + let routed = agent_tunnel::routing::try_route( + Some(agent_tunnel_handle.as_ref()), + // TODO: Pass `jet_agent_id` after JMUX consumers start issuing it. + None, + &route_target, + session_id, + target.as_addr(), + ) + .await?; + + let Some((stream, _agent)) = routed else { + return Ok(None); + }; + + Ok(Some(stream)) + } + }); + } + + let proxy_fut = proxy.run(); let proxy_handle = ChildTask::spawn(proxy_fut); let join_fut = proxy_handle.join(); tokio::pin!(join_fut); diff --git a/testsuite/Cargo.toml b/testsuite/Cargo.toml index 1cc45390a..88e163f60 100644 --- a/testsuite/Cargo.toml +++ b/testsuite/Cargo.toml @@ -40,6 +40,7 @@ devolutions-gateway-task = { path = "../crates/devolutions-gateway-task" } devolutions-gateway = { path = "../devolutions-gateway" } futures-util = "0.3" ipnetwork = "0.20" +jmux-proto = { path = "../crates/jmux-proto" } libsql = { version = "0.9", default-features = false, features = ["core"] } mcp-proxy.path = "../crates/mcp-proxy" network-scanner = { path = "../crates/network-scanner", features = ["test-utils"] } diff --git a/testsuite/tests/agent_tunnel/integration.rs b/testsuite/tests/agent_tunnel/integration.rs index b36433689..1389324e0 100644 --- a/testsuite/tests/agent_tunnel/integration.rs +++ b/testsuite/tests/agent_tunnel/integration.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Duration; use agent_tunnel::AgentTunnelHandle; @@ -5,10 +6,17 @@ use agent_tunnel::registry::AgentRegistry; use agent_tunnel_proto::{ CertRenewalResult, ConnectResponse, ControlMessage, ControlStream, DomainAdvertisement, DomainName, }; +use devolutions_gateway::recording::recording_message_channel; +use devolutions_gateway::session::SessionManagerTask; +use devolutions_gateway::subscriber::subscriber_channel; use devolutions_gateway::target_addr::TargetAddr; +use devolutions_gateway::token::{ApplicationProtocol, JmuxTokenClaims, RecordingPolicy, SessionTtl}; +use devolutions_gateway::traffic_audit::TrafficAuditHandle; use devolutions_gateway::upstream::{ConnectedUpstream, UpstreamLeg, connect_upstream}; +use devolutions_gateway_task::{ShutdownHandle, Task}; +use jmux_proto::{Bytes, BytesMut, DistantChannelId, Header, LocalChannelId, Message, ReasonCode}; use nonempty::NonEmpty; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use uuid::Uuid; @@ -21,6 +29,32 @@ fn target(host: &str, port: u16) -> TargetAddr { TargetAddr::from_components("tcp", host, port).expect("build target address") } +async fn send_jmux_message(writer: &mut (impl AsyncWrite + Unpin), message: Message) { + let mut bytes = BytesMut::new(); + message.encode(&mut bytes).expect("encode JMUX message"); + writer.write_all(&bytes).await.expect("send JMUX message"); +} + +async fn receive_jmux_message(reader: &mut (impl AsyncRead + Unpin)) -> Message { + tokio::time::timeout(Duration::from_secs(5), async { + let mut header = [0; Header::SIZE]; + reader.read_exact(&mut header).await.expect("read JMUX header"); + let header = Header::decode(Bytes::copy_from_slice(&header)).expect("decode JMUX header"); + let body_size = usize::from(header.size) + .checked_sub(Header::SIZE) + .expect("JMUX message size smaller than header"); + let mut body = vec![0; body_size]; + reader.read_exact(&mut body).await.expect("read JMUX body"); + + let mut bytes = BytesMut::with_capacity(usize::from(header.size)); + header.encode(&mut bytes); + bytes.extend_from_slice(&body); + Message::decode(bytes.freeze()).expect("decode JMUX message") + }) + .await + .expect("JMUX response timed out") +} + async fn advertise_domain( connection: &quinn::Connection, registry: &AgentRegistry, @@ -311,6 +345,134 @@ async fn gateway_connect_upstream_does_not_bypass_failed_agent_routes() { listener.shutdown().await; } +#[tokio::test] +async fn gateway_jmux_uses_agent_route_without_direct_fallback() { + let listener = bind_test_listener().await; + let (agent_id, connection) = listener.connect_agent("jmux-agent").await; + let _ctrl = advertise_routes( + &connection, + listener.handle.registry(), + agent_id, + 1, + vec!["127.0.0.0/8".parse().expect("parse test subnet")], + vec![], + ) + .await; + let direct_listener = TcpListener::bind("127.0.0.1:0").await.expect("bind direct target"); + let target_port = direct_listener.local_addr().expect("read direct target address").port(); + let target = target("127.0.0.1", target_port); + let session_id = Uuid::new_v4(); + + let (recordings, _recording_rx) = recording_message_channel(); + let session_manager = SessionManagerTask::init(recordings); + let sessions = session_manager.handle(); + let (session_shutdown, session_shutdown_signal) = ShutdownHandle::new(); + let session_task = tokio::spawn(session_manager.run(session_shutdown_signal)); + let (subscriber_tx, _subscriber_rx) = subscriber_channel(); + let (traffic_audit_handle, _traffic_audit_rx) = TrafficAuditHandle::new(); + let claims = JmuxTokenClaims { + jet_aid: session_id, + hosts: NonEmpty::new(target.clone()), + jet_ap: ApplicationProtocol::unknown(), + jet_rec: RecordingPolicy::None, + jet_ttl: SessionTtl::Unlimited, + exp: i64::MAX, + jti: Uuid::new_v4(), + }; + let (proxy_stream, mut peer_stream) = tokio::io::duplex(8192); + let proxy_task = tokio::spawn(devolutions_gateway::jmux::handle( + proxy_stream, + claims, + sessions, + subscriber_tx, + traffic_audit_handle, + Some(Arc::new(listener.handle.clone())), + )); + + send_jmux_message( + &mut peer_stream, + Message::open( + LocalChannelId::from(20), + 4096, + jmux_proto::DestinationUrl::new("tcp", "127.0.0.1", target_port), + ), + ) + .await; + let mut routed_session = tokio::time::timeout( + Duration::from_secs(5), + accept_session_request(&connection, session_id, target.as_addr()), + ) + .await + .expect("routed JMUX request timed out"); + routed_session + .send_response(&ConnectResponse::success()) + .await + .expect("accept routed JMUX request"); + let Message::OpenSuccess(success) = receive_jmux_message(&mut peer_stream).await else { + panic!("expected OPEN SUCCESS"); + }; + assert_eq!(success.recipient_channel_id, 20); + + let local_id = DistantChannelId::from(success.sender_channel_id); + send_jmux_message(&mut peer_stream, Message::data(local_id, Bytes::from_static(b"ping"))).await; + let (mut routed_send, mut routed_recv) = routed_session.into_inner(); + let mut request = [0; 4]; + routed_recv + .read_exact(&mut request) + .await + .expect("read routed JMUX payload"); + assert_eq!(&request, b"ping"); + routed_send + .write_all(b"pong") + .await + .expect("write routed JMUX response"); + let Message::Data(response) = receive_jmux_message(&mut peer_stream).await else { + panic!("expected CHANNEL DATA"); + }; + assert_eq!(response.recipient_channel_id, 20); + assert_eq!(response.transfer_data, b"pong"[..]); + + send_jmux_message( + &mut peer_stream, + Message::open( + LocalChannelId::from(21), + 4096, + jmux_proto::DestinationUrl::new("tcp", "127.0.0.1", target_port), + ), + ) + .await; + let mut failed_session = tokio::time::timeout( + Duration::from_secs(5), + accept_session_request(&connection, session_id, target.as_addr()), + ) + .await + .expect("failed routed JMUX request timed out"); + failed_session + .send_response(&ConnectResponse::error("connection refused")) + .await + .expect("reject routed JMUX request"); + let Message::OpenFailure(failure) = receive_jmux_message(&mut peer_stream).await else { + panic!("expected OPEN FAILURE"); + }; + assert_eq!(failure.recipient_channel_id, 21); + assert_eq!(failure.reason_code, ReasonCode::GENERAL_FAILURE); + assert!( + tokio::time::timeout(Duration::from_millis(100), direct_listener.accept()) + .await + .is_err(), + "matched Agent route must not fall back to direct TCP" + ); + + proxy_task.abort(); + session_shutdown.signal(); + session_task + .await + .expect("session manager task panicked") + .expect("session manager shutdown"); + connection.close(0u32.into(), b"test done"); + listener.shutdown().await; +} + #[tokio::test] async fn gateway_listener_rejects_certificate_renewal_key_rotation() { let listener = bind_test_listener().await;