Add native remote port forwarding

src/bin/dosh-server.rs

@@ -23,7 +23,7 @@ use std::net::SocketAddr;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::{TcpStream, UdpSocket};
+use tokio::net::{TcpListener, TcpStream, UdpSocket};
 use tokio::sync::mpsc;
 
 #[derive(Debug, Parser)]
@@ -164,6 +164,7 @@ struct ServerState {
     pty_tx: mpsc::UnboundedSender<PtyOutput>,
     sessions: HashMap<String, Session>,
     pending_native: HashMap<[u8; 16], PendingNativeAuth>,
+    next_server_stream_id: u64,
 }
 
 struct Session {
@@ -219,6 +220,7 @@ impl ServerState {
             pty_tx,
             sessions: HashMap::new(),
             pending_native: HashMap::new(),
+            next_server_stream_id: 1u64 << 63,
         }
     }
 
@@ -271,6 +273,8 @@ async fn handle_packet(
         | PacketKind::Ping
         | PacketKind::Ack
         | PacketKind::StreamOpen
+        | PacketKind::StreamOpenOk
+        | PacketKind::StreamOpenReject
         | PacketKind::StreamData
         | PacketKind::StreamClose
         | PacketKind::StreamEof
@@ -285,6 +289,8 @@ async fn handle_packet(
         PacketKind::Ack => handle_ack(state, &packet).await,
         PacketKind::Detach => handle_detach(state, &packet).await,
         PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await,
+        PacketKind::StreamOpenOk => handle_stream_open_ok(state, peer, &packet).await,
+        PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await,
         PacketKind::StreamData => handle_stream_data(state, peer, &packet).await,
         PacketKind::StreamClose => handle_stream_close(state, &packet).await,
         _ => Ok(()),
@@ -406,6 +412,7 @@ async fn handle_native_user_auth(
         String,
         u64,
         Vec<u8>,
+        Vec<ForwardingRequest>,
     )> = {
         let mut locked = state.lock().expect("server state poisoned");
         verify_native_user_auth_from_config(
@@ -486,6 +493,7 @@ async fn handle_native_user_auth(
                 mode,
                 output_seq,
                 snapshot,
+                req.auth.requested_forwardings.clone(),
             ))
         })
     };
@@ -499,6 +507,7 @@ async fn handle_native_user_auth(
         mode,
         output_seq,
         snapshot,
+        requested_forwardings,
     ) = match attached {
         Ok(attached) => attached,
         Err(err) => {
@@ -506,6 +515,18 @@ async fn handle_native_user_auth(
                 .await;
         }
     };
+    if let Err(err) = start_remote_forwards(
+        state,
+        socket,
+        client_id,
+        requested_forwardings.clone(),
+        pending.client.requested_session.clone(),
+    )
+    .await
+    {
+        remove_client(state, client_id);
+        return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string()).await;
+    }
 
     let ok = NativeAuthOk {
         client_id,
@@ -909,6 +930,13 @@ async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packe
     Ok(())
 }
 
+fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
+    let mut locked = state.lock().expect("server state poisoned");
+    for session in locked.sessions.values_mut() {
+        session.clients.remove(&client_id);
+    }
+}
+
 async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
     let (key, _) = find_client_key(state, &packet.header.conn_id)?;
     let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
@@ -1073,6 +1101,57 @@ async fn handle_stream_data(
     Ok(())
 }
 
+async fn handle_stream_open_ok(
+    state: &Arc<Mutex<ServerState>>,
+    peer: SocketAddr,
+    packet: &protocol::Packet,
+) -> Result<()> {
+    let (key, session_name) = find_client_key(state, &packet.header.conn_id)?;
+    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
+    let _ok: StreamOpenOk = protocol::from_body(&body)?;
+    let mut locked = state.lock().expect("server state poisoned");
+    let session = locked
+        .sessions
+        .get_mut(&session_name)
+        .ok_or_else(|| anyhow!("unknown session"))?;
+    let client = session
+        .clients
+        .get_mut(&packet.header.conn_id)
+        .ok_or_else(|| anyhow!("unknown client"))?;
+    if !client.replay.accept(packet.header.seq) {
+        return Ok(());
+    }
+    client.endpoint = peer;
+    client.last_seen = Instant::now();
+    Ok(())
+}
+
+async fn handle_stream_open_reject(
+    state: &Arc<Mutex<ServerState>>,
+    peer: SocketAddr,
+    packet: &protocol::Packet,
+) -> Result<()> {
+    let (key, session_name) = find_client_key(state, &packet.header.conn_id)?;
+    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
+    let reject: StreamOpenReject = protocol::from_body(&body)?;
+    let mut locked = state.lock().expect("server state poisoned");
+    let session = locked
+        .sessions
+        .get_mut(&session_name)
+        .ok_or_else(|| anyhow!("unknown session"))?;
+    let client = session
+        .clients
+        .get_mut(&packet.header.conn_id)
+        .ok_or_else(|| anyhow!("unknown client"))?;
+    if !client.replay.accept(packet.header.seq) {
+        return Ok(());
+    }
+    client.endpoint = peer;
+    client.last_seen = Instant::now();
+    client.stream_writers.remove(&reject.stream_id);
+    Ok(())
+}
+
 async fn handle_stream_close(
     state: &Arc<Mutex<ServerState>>,
     packet: &protocol::Packet,
@@ -1095,6 +1174,134 @@ async fn handle_stream_close(
     Ok(())
 }
 
+async fn start_remote_forwards(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    requested_forwardings: Vec<ForwardingRequest>,
+    session_name: String,
+) -> Result<()> {
+    let remote_forwards = requested_forwardings
+        .into_iter()
+        .filter(|forward| forward.kind == ForwardingKind::Remote)
+        .collect::<Vec<_>>();
+    if remote_forwards.is_empty() {
+        return Ok(());
+    }
+    let config = {
+        let locked = state.lock().expect("server state poisoned");
+        locked.config.clone()
+    };
+    anyhow::ensure!(
+        config.allow_remote_forwarding,
+        "remote forwarding disabled on server"
+    );
+    for forward in remote_forwards {
+        let bind_host = forward
+            .bind_host
+            .clone()
+            .unwrap_or_else(|| "127.0.0.1".to_string());
+        anyhow::ensure!(
+            is_loopback_bind(&bind_host),
+            "remote forwarding may only bind loopback in this server build"
+        );
+        let listen_port = forward.listen_port;
+        let target_host = forward
+            .target_host
+            .clone()
+            .ok_or_else(|| anyhow!("remote forward target host is required"))?;
+        let target_port = forward
+            .target_port
+            .ok_or_else(|| anyhow!("remote forward target port is required"))?;
+        let bind = format!("{bind_host}:{listen_port}");
+        let listener = TcpListener::bind(&bind)
+            .await
+            .with_context(|| format!("bind remote forward {bind}"))?;
+        let state = Arc::clone(state);
+        let socket = Arc::clone(socket);
+        let session_name = session_name.clone();
+        tokio::spawn(async move {
+            loop {
+                let Ok((stream, _)) = listener.accept().await else {
+                    break;
+                };
+                let stream_id = allocate_server_stream_id(&state);
+                let (mut reader, mut writer) = stream.into_split();
+                let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>();
+                {
+                    let mut locked = state.lock().expect("server state poisoned");
+                    if let Some(session) = locked.sessions.get_mut(&session_name)
+                        && let Some(client) = session.clients.get_mut(&client_id)
+                    {
+                        client.stream_writers.insert(stream_id, writer_tx);
+                    } else {
+                        break;
+                    }
+                }
+                tokio::spawn(async move {
+                    while let Some(bytes) = writer_rx.recv().await {
+                        if writer.write_all(&bytes).await.is_err() {
+                            break;
+                        }
+                    }
+                });
+                if let Err(err) = send_stream_open_to_client(
+                    &state,
+                    &socket,
+                    client_id,
+                    stream_id,
+                    target_host.clone(),
+                    target_port,
+                )
+                .await
+                {
+                    eprintln!("remote forward open error: {err:#}");
+                    break;
+                }
+                let state = Arc::clone(&state);
+                let socket = Arc::clone(&socket);
+                tokio::spawn(async move {
+                    let mut buf = [0u8; 16 * 1024];
+                    loop {
+                        match reader.read(&mut buf).await {
+                            Ok(0) => break,
+                            Ok(n) => {
+                                if let Err(err) = send_stream_data_to_client(
+                                    &state,
+                                    &socket,
+                                    client_id,
+                                    stream_id,
+                                    buf[..n].to_vec(),
+                                )
+                                .await
+                                {
+                                    eprintln!("remote forward data error: {err:#}");
+                                    break;
+                                }
+                            }
+                            Err(_) => break,
+                        }
+                    }
+                    let _ =
+                        send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
+                });
+            }
+        });
+    }
+    Ok(())
+}
+
+fn allocate_server_stream_id(state: &Arc<Mutex<ServerState>>) -> u64 {
+    let mut locked = state.lock().expect("server state poisoned");
+    let stream_id = locked.next_server_stream_id;
+    locked.next_server_stream_id = locked.next_server_stream_id.wrapping_add(1).max(1u64 << 63);
+    stream_id
+}
+
+fn is_loopback_bind(host: &str) -> bool {
+    matches!(host, "127.0.0.1" | "localhost" | "::1")
+}
+
 fn stream_open_allowed(
     config: &ServerConfig,
     client: &ClientState,
@@ -1125,6 +1332,22 @@ async fn send_stream_open_ok(
     send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
 }
 
+async fn send_stream_open_to_client(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+    target_host: String,
+    target_port: u16,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamOpen {
+        stream_id,
+        target_host,
+        target_port,
+    })?;
+    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
+}
+
 async fn send_stream_open_reject(
     state: &Arc<Mutex<ServerState>>,
     socket: &Arc<UdpSocket>,