Skip to content

feat(dgw,agent,jetsocat): route JMUX channels through agents #1974

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
irvingouj@Devolutions (irvingoujAtDevolution) wants to merge 7 commits into
base: master
Choose a base branch
from
Open

feat(dgw,agent,jetsocat): route JMUX channels through agents #1974

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c816627
feat(dgw): route JMUX channels through agents
irvingoujAtDevolution Sep 7, 2026
8995646
fix(dgw,agent): preserve routed target addresses
irvingoujAtDevolution Sep 7, 2026
54c106c
fix(dgw,jetsocat): release failed JMUX channel IDs
irvingoujAtDevolution Sep 7, 2026
d954814
docs(dgw): document routed target address
irvingoujAtDevolution Sep 7, 2026
c10d66a
docs(dgw): clarify agent routing result
irvingoujAtDevolution Sep 7, 2026
6572260
fix(dgw,agent,jetsocat): simplify JMUX routing
irvingoujAtDevolution Sep 8, 2026
201f03c
fix(agent): preserve tunneled TCP half-close
irvingoujAtDevolution Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 109 additions & 10 deletions crates/jmux-proxy/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -53,6 +55,26 @@ pub type ApiResponseReceiver = oneshot::Receiver<JmuxApiResponse>;
pub type ApiRequestSender = mpsc::Sender<JmuxApiRequest>;
pub type ApiRequestReceiver = mpsc::Receiver<JmuxApiRequest>;

trait TargetStream: AsyncRead + AsyncWrite + Unpin + Send {}

impl<T> TargetStream for T where T: AsyncRead + AsyncWrite + Unpin + Send {}

type ErasedTargetStream = Box<dyn TargetStream>;
type TargetConnectorFuture = Pin<Box<dyn Future<Output = anyhow::Result<Option<ConnectedTarget>>> + Send>>;
type TargetConnector = Arc<dyn Fn(DestinationUrl) -> TargetConnectorFuture + Send + Sync>;

pub struct ConnectedTarget {
stream: ErasedTargetStream,
}

impl ConnectedTarget {
pub fn new(stream: impl AsyncRead + AsyncWrite + Unpin + Send + 'static) -> Self {
Self {
stream: Box::new(stream),
}
}
}

#[derive(Debug)]
pub enum JmuxApiRequest {
OpenChannel {
Expand Down Expand Up @@ -84,6 +106,7 @@ pub struct JmuxProxy {
jmux_reader: Box<dyn AsyncRead + Unpin + Send>,
jmux_writer: Box<dyn AsyncWrite + Unpin + Send>,
traffic_callback: Option<TrafficCallback>,
target_connector: Option<TargetConnector>,
}

impl JmuxProxy {
Expand All @@ -98,6 +121,7 @@ impl JmuxProxy {
jmux_reader,
jmux_writer,
traffic_callback: None,
target_connector: None,
}
}

Expand All @@ -113,6 +137,19 @@ impl JmuxProxy {
self
}

/// Tries a custom target connection before falling back to direct TCP.
///
/// Return `Ok(None)` when the target should use the default direct connection.
#[must_use]
pub fn with_target_connector<C, F>(mut self, connector: C) -> Self
where
C: Fn(DestinationUrl) -> F + Send + Sync + 'static,
F: Future<Output = anyhow::Result<Option<ConnectedTarget>>> + Send + 'static,
Comment thread
irvingoujAtDevolution marked this conversation as resolved.
{
self.target_connector = Some(Arc::new(move |destination_url| Box::pin(connector(destination_url))));
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
Expand Down Expand Up @@ -186,6 +223,7 @@ async fn run_proxy_impl(proxy: JmuxProxy, span: Span) -> anyhow::Result<()> {
jmux_reader,
jmux_writer,
traffic_callback,
target_connector,
} = proxy;

let (msg_to_send_tx, msg_to_send_rx) = mpsc::channel::<Message>(JMUX_MESSAGE_MPSC_CHANNEL_SIZE);
Expand All @@ -206,6 +244,7 @@ async fn run_proxy_impl(proxy: JmuxProxy, span: Span) -> anyhow::Result<()> {
msg_to_send_tx,
api_request_rx,
traffic_callback,
target_connector,
parent_span: span,
}
.spawn();
Expand Down Expand Up @@ -255,7 +294,7 @@ struct JmuxChannelCtx {
// Traffic audit metadata
target_host: String,
/// Target server resolved address IP
target_ip: Option<std::net::IpAddr>,
target_ip: Option<IpAddr>,
/// Target server port
target_port: u16,
/// Time the connection with target peer was established at
Expand Down Expand Up @@ -344,7 +383,6 @@ type DataReceiver = mpsc::Receiver<Bytes>;
type DataSender = mpsc::Sender<Bytes>;
type InternalMessageSender = mpsc::Sender<InternalMessage>;

#[derive(Debug)]
enum InternalMessage {
Eof {
id: LocalChannelId,
Expand All @@ -353,7 +391,11 @@ enum InternalMessage {
// Boxing reduces enum size from 224 bytes to ~16 bytes
// (clippy::large_enum_variant)
channel: Box<JmuxChannelCtx>,
stream: TcpStream,
stream: ErasedTargetStream,
},
TargetConnectionFailed {
id: LocalChannelId,
distant_id: DistantChannelId,
},
AbnormalTermination {
id: LocalChannelId,
Expand Down Expand Up @@ -427,6 +469,7 @@ struct JmuxSchedulerTask<T: AsyncRead + Unpin + Send + 'static> {
msg_to_send_tx: MessageSender,
api_request_rx: ApiRequestReceiver,
traffic_callback: Option<TrafficCallback>,
target_connector: Option<TargetConnector>,
parent_span: Span,
}

Expand All @@ -448,6 +491,7 @@ async fn scheduler_task_impl<T: AsyncRead + Unpin + Send + 'static>(task: JmuxSc
msg_to_send_tx,
mut api_request_rx,
traffic_callback,
target_connector,
parent_span,
} = task;

Expand Down Expand Up @@ -501,7 +545,8 @@ async fn scheduler_task_impl<T: AsyncRead + Unpin + Send + 'static>(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,
Expand Down Expand Up @@ -626,7 +671,7 @@ async fn scheduler_task_impl<T: AsyncRead + Unpin + Send + 'static>(task: JmuxSc
debug!("Channel accepted");
});

let (reader, writer) = stream.into_split();
let (reader, writer) = tokio::io::split(stream);

DataWriterTask {
writer,
Expand All @@ -652,6 +697,17 @@ async fn scheduler_task_impl<T: AsyncRead + Unpin + Send + 'static>(task: JmuxSc
.spawn(channel_span)
.detach();
}
InternalMessage::TargetConnectionFailed { id, distant_id } => {
jmux_ctx.id_allocator.free(id);
msg_to_send_tx
.send(Message::open_failure(
distant_id,
ReasonCode::GENERAL_FAILURE,
"target connection failed",
))
.await
.context("couldn't send OPEN FAILURE message through mpsc channel")?;
}
}
}
msg = jmux_stream.next() => {
Expand Down Expand Up @@ -760,6 +816,7 @@ async fn scheduler_task_impl<T: AsyncRead + Unpin + Send + 'static>(task: JmuxSc
internal_msg_tx: internal_msg_tx.clone(),
msg_to_send_tx: msg_to_send_tx.clone(),
traffic_callback: traffic_callback.clone(),
target_connector: target_connector.clone(),
}
.spawn()
.detach();
Expand Down Expand Up @@ -958,7 +1015,7 @@ async fn scheduler_task_impl<T: AsyncRead + Unpin + Send + 'static>(task: JmuxSc
// ---------------------- //

struct DataReaderTask {
reader: OwnedReadHalf,
reader: tokio::io::ReadHalf<ErasedTargetStream>,
local_id: LocalChannelId,
distant_id: DistantChannelId,
window_size_updated: Arc<Notify>,
Expand Down Expand Up @@ -1087,7 +1144,7 @@ impl DataReaderTask {
// ---------------------- //

struct DataWriterTask {
writer: OwnedWriteHalf,
writer: tokio::io::WriteHalf<ErasedTargetStream>,
data_rx: DataReceiver,
/// Tracks bytes written into the stream.
bytes_tx: Arc<AtomicU64>,
Expand Down Expand Up @@ -1123,6 +1180,8 @@ impl DataWriterTask {

bytes_tx.fetch_add(data.len() as u64, Ordering::SeqCst);
}

let _ = writer.shutdown().await;
}
.instrument(span),
);
Expand All @@ -1139,6 +1198,7 @@ struct StreamResolverTask {
internal_msg_tx: InternalMessageSender,
msg_to_send_tx: MessageSender,
traffic_callback: Option<TrafficCallback>,
target_connector: Option<TargetConnector>,
}

impl StreamResolverTask {
Expand All @@ -1164,6 +1224,7 @@ impl StreamResolverTask {
internal_msg_tx,
msg_to_send_tx,
traffic_callback,
target_connector,
} = self;

let scheme = destination_url.scheme();
Expand All @@ -1172,6 +1233,42 @@ impl StreamResolverTask {

match scheme {
"tcp" => {
if let Some(connector) = target_connector {
match connector(destination_url.clone()).await {
Ok(Some(ConnectedTarget { stream })) => {
channel.connect_at = SystemTime::now();

internal_msg_tx
.send(InternalMessage::StreamResolved {
channel: Box::new(channel),
stream,
})
.await
.map_err(|_| {
anyhow::anyhow!("couldn't send back resolved stream through internal mpsc channel")
})?;

return Ok(());
}
Ok(None) => {}
Err(error) => {
internal_msg_tx
.send(InternalMessage::TargetConnectionFailed {
id: channel.local_id,
distant_id: channel.distant_id,
})
.await
.map_err(|_| {
anyhow::anyhow!(
"couldn't report target connection failure through internal mpsc channel"
)
})?;

return Err(error.context(format!("couldn't connect to {host}:{port}")));
}
}
}

// Perform DNS resolution first to get concrete IP addresses.
let socket_addrs = match tokio::net::lookup_host((host, port)).await {
Ok(addrs) => addrs,
Expand Down Expand Up @@ -1203,10 +1300,12 @@ impl StreamResolverTask {
internal_msg_tx
.send(InternalMessage::StreamResolved {
channel: Box::new(channel),
stream,
stream: Box::new(stream),
})
.await
.context("couldn't send back resolved stream through internal mpsc channel")?;
.map_err(|_| {
anyhow::anyhow!("couldn't send back resolved stream through internal mpsc channel")
})?;

return Ok(());
}
Expand Down
Loading
Loading