Add native local port forwarding

src/bin/dosh-server.rs

@@ -7,22 +7,23 @@ use dosh::auth::{
 use dosh::config::{ServerConfig, load_server_config};
 use dosh::crypto;
 use dosh::native::{
-    NativeAuthOk, NativeServerHello, derive_native_session_key, generate_native_ephemeral,
-    host_public_key, host_public_key_line, load_or_create_host_key, sign_server_hello,
-    verify_native_user_auth_from_config,
+    ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello, derive_native_session_key,
+    generate_native_ephemeral, host_public_key, host_public_key_line, load_or_create_host_key,
+    sign_server_hello, verify_native_user_auth_from_config,
 };
 use dosh::protocol::{
     self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
     NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind,
-    ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope,
-    TicketAttachOkEnvelope,
+    ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen,
+    StreamOpenOk, StreamOpenReject, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session};
 use std::collections::{HashMap, VecDeque};
 use std::net::SocketAddr;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
-use tokio::net::UdpSocket;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpStream, UdpSocket};
 use tokio::sync::mpsc;
 
 #[derive(Debug, Parser)]
@@ -186,6 +187,8 @@ struct ClientState {
     last_seen: Instant,
     pending: VecDeque<PendingFrame>,
     last_screen: Option<vt100::Screen>,
+    allowed_forwardings: Vec<ForwardingRequest>,
+    stream_writers: HashMap<u64, mpsc::UnboundedSender<Vec<u8>>>,
 }
 
 #[derive(Clone)]
@@ -263,7 +266,15 @@ async fn handle_packet(
             handle_ticket_attach(state, socket, peer, packet.body).await
         }
         PacketKind::ResumeRequest => handle_resume(state, socket, peer, &packet).await,
-        PacketKind::Input | PacketKind::Resize | PacketKind::Ping | PacketKind::Ack
+        PacketKind::Input
+        | PacketKind::Resize
+        | PacketKind::Ping
+        | PacketKind::Ack
+        | PacketKind::StreamOpen
+        | PacketKind::StreamData
+        | PacketKind::StreamClose
+        | PacketKind::StreamEof
+        | PacketKind::StreamWindowAdjust
             if find_client_key(state, &packet.header.conn_id).is_err() =>
         {
             send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").await
@@ -273,6 +284,9 @@ async fn handle_packet(
         PacketKind::Ping => handle_ping(state, socket, peer, &packet).await,
         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::StreamData => handle_stream_data(state, peer, &packet).await,
+        PacketKind::StreamClose => handle_stream_close(state, &packet).await,
         _ => Ok(()),
     }
 }
@@ -459,6 +473,8 @@ async fn handle_native_user_auth(
                     last_seen: Instant::now(),
                     pending: VecDeque::new(),
                     last_screen: Some(screen),
+                    allowed_forwardings: req.auth.requested_forwardings.clone(),
+                    stream_writers: HashMap::new(),
                 },
             );
             Ok((
@@ -562,6 +578,8 @@ async fn handle_bootstrap_attach(
                 last_seen: Instant::now(),
                 pending: VecDeque::new(),
                 last_screen: Some(screen),
+                allowed_forwardings: Vec::new(),
+                stream_writers: HashMap::new(),
             },
         );
         (
@@ -670,6 +688,8 @@ async fn handle_ticket_attach(
                 last_seen: Instant::now(),
                 pending: VecDeque::new(),
                 last_screen: Some(screen),
+                allowed_forwardings: Vec::new(),
+                stream_writers: HashMap::new(),
             },
         );
         (client_id, output_seq, snapshot)
@@ -913,6 +933,273 @@ async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet)
     Ok(())
 }
 
+async fn handle_stream_open(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    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 open: StreamOpen = protocol::from_body(&body)?;
+    let allowed = {
+        let mut locked = state.lock().expect("server state poisoned");
+        let config = locked.config.clone();
+        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();
+        stream_open_allowed(&config, client, &open)
+    };
+    if let Err(err) = allowed {
+        return send_stream_open_reject(
+            state,
+            socket,
+            packet.header.conn_id,
+            open.stream_id,
+            err.to_string(),
+        )
+        .await;
+    }
+
+    let target = format!("{}:{}", open.target_host, open.target_port);
+    let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
+        Ok(stream) => stream,
+        Err(err) => {
+            return send_stream_open_reject(
+                state,
+                socket,
+                packet.header.conn_id,
+                open.stream_id,
+                format!("connect {target}: {err}"),
+            )
+            .await;
+        }
+    };
+    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");
+        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"))?;
+        client.stream_writers.insert(open.stream_id, writer_tx);
+    }
+
+    tokio::spawn(async move {
+        while let Some(bytes) = writer_rx.recv().await {
+            if writer.write_all(&bytes).await.is_err() {
+                break;
+            }
+        }
+    });
+
+    send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;
+
+    let state = Arc::clone(state);
+    let socket = Arc::clone(socket);
+    let client_id = packet.header.conn_id;
+    let stream_id = open.stream_id;
+    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!("stream send error: {err:#}");
+                        break;
+                    }
+                }
+                Err(_) => break,
+            }
+        }
+        let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
+    });
+
+    Ok(())
+}
+
+async fn handle_stream_data(
+    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 data: StreamData = protocol::from_body(&body)?;
+    let writer = {
+        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.get(&data.stream_id).cloned()
+    };
+    if let Some(writer) = writer {
+        let _ = writer.send(data.bytes);
+    }
+    Ok(())
+}
+
+async fn handle_stream_close(
+    state: &Arc<Mutex<ServerState>>,
+    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 close: StreamClose = protocol::from_body(&body)?;
+    let mut locked = state.lock().expect("server state poisoned");
+    let Some(session) = locked.sessions.get_mut(&session_name) else {
+        return Ok(());
+    };
+    let Some(client) = session.clients.get_mut(&packet.header.conn_id) else {
+        return Ok(());
+    };
+    if !client.replay.accept(packet.header.seq) {
+        return Ok(());
+    }
+    client.last_seen = Instant::now();
+    client.stream_writers.remove(&close.stream_id);
+    Ok(())
+}
+
+fn stream_open_allowed(
+    config: &ServerConfig,
+    client: &ClientState,
+    open: &StreamOpen,
+) -> Result<()> {
+    anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled");
+    let allowed = client.allowed_forwardings.iter().any(|forward| {
+        forward.kind == ForwardingKind::Local
+            && forward.target_host.as_deref() == Some(open.target_host.as_str())
+            && forward.target_port == Some(open.target_port)
+    });
+    anyhow::ensure!(
+        allowed,
+        "stream target {}:{} was not requested during auth",
+        open.target_host,
+        open.target_port
+    );
+    Ok(())
+}
+
+async fn send_stream_open_ok(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamOpenOk { stream_id })?;
+    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
+}
+
+async fn send_stream_open_reject(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+    reason: String,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
+    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenReject, body).await
+}
+
+async fn send_stream_data_to_client(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+    bytes: Vec<u8>,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamData { stream_id, bytes })?;
+    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamData, body).await
+}
+
+async fn send_stream_close_to_client(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+) -> Result<()> {
+    {
+        let mut locked = state.lock().expect("server state poisoned");
+        for session in locked.sessions.values_mut() {
+            if let Some(client) = session.clients.get_mut(&client_id) {
+                client.stream_writers.remove(&stream_id);
+                break;
+            }
+        }
+    }
+    let body = protocol::to_body(&StreamClose { stream_id })?;
+    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await
+}
+
+async fn send_stream_packet_to_client(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    kind: PacketKind,
+    body: Vec<u8>,
+) -> Result<()> {
+    let send = {
+        let mut locked = state.lock().expect("server state poisoned");
+        let mut found = None;
+        for session in locked.sessions.values_mut() {
+            if let Some(client) = session.clients.get_mut(&client_id) {
+                client.send_seq += 1;
+                let packet = protocol::encode_encrypted(
+                    kind,
+                    client_id,
+                    client.send_seq,
+                    client.last_acked,
+                    &client.session_key,
+                    SERVER_TO_CLIENT,
+                    &body,
+                )?;
+                found = Some((client.endpoint, packet));
+                break;
+            }
+        }
+        found
+    };
+    if let Some((endpoint, packet)) = send {
+        socket.send_to(&packet, endpoint).await?;
+    }
+    Ok(())
+}
+
 async fn broadcast_output(
     state: &Arc<Mutex<ServerState>>,
     socket: &Arc<UdpSocket>,