qssh/
port_forward.rs

1//! Port forwarding implementation for QSSH
2//! Supports local (-L), remote (-R), and dynamic (-D) forwarding
3
4use std::collections::HashMap;
5use std::net::SocketAddr;
6use std::sync::Arc;
7use tokio::net::{TcpListener, TcpStream};
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use tokio::sync::{mpsc, Mutex};
10use crate::{Result, QsshError};
11use crate::transport::{Transport, Message, ChannelMessage, ChannelType,
12    GlobalRequestMessage, GlobalRequestType};
13
14/// Port forwarding types
15#[derive(Debug, Clone)]
16pub enum ForwardType {
17    /// Local port forwarding (-L)
18    /// Listen locally, forward to remote
19    Local {
20        bind_addr: SocketAddr,
21        remote_host: String,
22        remote_port: u16,
23    },
24    /// Remote port forwarding (-R)
25    /// Listen on remote, forward to local
26    Remote {
27        remote_bind_addr: SocketAddr,
28        local_host: String,
29        local_port: u16,
30    },
31    /// Dynamic port forwarding (-D)
32    /// SOCKS proxy
33    Dynamic {
34        bind_addr: SocketAddr,
35    },
36}
37
38/// Type alias for the remote-forward mapping table.
39type ForwardMapping = HashMap<(String, u16), (String, u16)>;
40
41/// Registry mapping (bind_host, bind_port) on the server to (local_host, local_port) on the client.
42/// Used by the client to know where to connect when a ForwardedTcpip channel arrives.
43#[derive(Debug, Clone)]
44pub struct RemoteForwardRegistry {
45    mappings: Arc<Mutex<ForwardMapping>>,
46}
47
48impl Default for RemoteForwardRegistry {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl RemoteForwardRegistry {
55    pub fn new() -> Self {
56        Self {
57            mappings: Arc::new(Mutex::new(HashMap::new())),
58        }
59    }
60
61    /// Register a remote forward mapping
62    pub async fn insert(&self, bind_host: String, bind_port: u16, local_host: String, local_port: u16) {
63        let mut map = self.mappings.lock().await;
64        map.insert((bind_host, bind_port), (local_host, local_port));
65    }
66
67    /// Look up the local target for an incoming forwarded connection
68    pub async fn lookup(&self, connected_host: &str, connected_port: u16) -> Option<(String, u16)> {
69        let map = self.mappings.lock().await;
70        map.get(&(connected_host.to_string(), connected_port)).cloned()
71    }
72
73    /// Remove a mapping
74    pub async fn remove(&self, bind_host: &str, bind_port: u16) -> Option<(String, u16)> {
75        let mut map = self.mappings.lock().await;
76        map.remove(&(bind_host.to_string(), bind_port))
77    }
78}
79
80/// Routes channel data from the client's message loop to active forwarded connections.
81/// When a ForwardedTcpip channel is set up, a sender is registered here. The client's
82/// message loop dispatches `ChannelMessage::Data` to the matching sender.
83#[derive(Debug, Clone)]
84pub struct ForwardedChannelRouter {
85    senders: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
86}
87
88impl Default for ForwardedChannelRouter {
89    fn default() -> Self {
90        Self::new()
91    }
92}
93
94impl ForwardedChannelRouter {
95    pub fn new() -> Self {
96        Self {
97            senders: Arc::new(Mutex::new(HashMap::new())),
98        }
99    }
100
101    /// Register a sender for a channel
102    pub async fn register(&self, channel_id: u32, sender: mpsc::Sender<Vec<u8>>) {
103        let mut map = self.senders.lock().await;
104        map.insert(channel_id, sender);
105    }
106
107    /// Route data to a channel's handler. Returns true if routed, false if no handler.
108    pub async fn route_data(&self, channel_id: u32, data: Vec<u8>) -> bool {
109        let map = self.senders.lock().await;
110        if let Some(sender) = map.get(&channel_id) {
111            sender.send(data).await.is_ok()
112        } else {
113            false
114        }
115    }
116
117    /// Check if a channel has a registered handler
118    pub async fn has_channel(&self, channel_id: u32) -> bool {
119        let map = self.senders.lock().await;
120        map.contains_key(&channel_id)
121    }
122
123    /// Remove a channel's handler (called when the channel closes)
124    pub async fn remove(&self, channel_id: u32) {
125        let mut map = self.senders.lock().await;
126        map.remove(&channel_id);
127    }
128}
129
130/// Port forwarding manager
131pub struct PortForwardManager {
132    transport: Arc<Transport>,
133    forwards: Vec<ForwardType>,
134    remote_registry: RemoteForwardRegistry,
135    channel_router: ForwardedChannelRouter,
136}
137
138impl PortForwardManager {
139    pub fn new(transport: Arc<Transport>) -> Self {
140        Self {
141            transport,
142            forwards: Vec::new(),
143            remote_registry: RemoteForwardRegistry::new(),
144            channel_router: ForwardedChannelRouter::new(),
145        }
146    }
147
148    /// Get a clone of the remote forward registry (for sharing with the client message handler)
149    pub fn remote_registry(&self) -> RemoteForwardRegistry {
150        self.remote_registry.clone()
151    }
152
153    /// Get a clone of the forwarded channel router (for sharing with the client message handler)
154    pub fn channel_router(&self) -> ForwardedChannelRouter {
155        self.channel_router.clone()
156    }
157    
158    /// Parse forwarding spec (e.g., "8080:localhost:80")
159    pub fn parse_forward_spec(spec: &str, forward_type: &str) -> Result<ForwardType> {
160        let parts: Vec<&str> = spec.split(':').collect();
161        
162        match forward_type {
163            "local" => {
164                if parts.len() != 3 {
165                    return Err(QsshError::Config("Invalid local forward spec. Use: local_port:remote_host:remote_port".into()));
166                }
167                
168                let local_port: u16 = parts[0].parse()
169                    .map_err(|_| QsshError::Config("Invalid local port".into()))?;
170                let remote_host = parts[1].to_string();
171                let remote_port: u16 = parts[2].parse()
172                    .map_err(|_| QsshError::Config("Invalid remote port".into()))?;
173                
174                Ok(ForwardType::Local {
175                    bind_addr: ([127, 0, 0, 1], local_port).into(),
176                    remote_host,
177                    remote_port,
178                })
179            }
180            "remote" => {
181                if parts.len() != 3 {
182                    return Err(QsshError::Config("Invalid remote forward spec. Use: remote_port:local_host:local_port".into()));
183                }
184                
185                let remote_port: u16 = parts[0].parse()
186                    .map_err(|_| QsshError::Config("Invalid remote port".into()))?;
187                let local_host = parts[1].to_string();
188                let local_port: u16 = parts[2].parse()
189                    .map_err(|_| QsshError::Config("Invalid local port".into()))?;
190                
191                Ok(ForwardType::Remote {
192                    remote_bind_addr: ([0, 0, 0, 0], remote_port).into(),
193                    local_host,
194                    local_port,
195                })
196            }
197            "dynamic" => {
198                let port: u16 = spec.parse()
199                    .map_err(|_| QsshError::Config("Invalid SOCKS port".into()))?;
200                
201                Ok(ForwardType::Dynamic {
202                    bind_addr: ([127, 0, 0, 1], port).into(),
203                })
204            }
205            _ => Err(QsshError::Config("Unknown forward type".into()))
206        }
207    }
208    
209    /// Start local port forwarding
210    pub async fn start_local_forward(
211        &self,
212        bind_addr: SocketAddr,
213        remote_host: String,
214        remote_port: u16,
215    ) -> Result<()> {
216        let listener = TcpListener::bind(bind_addr).await?;
217        let transport = self.transport.clone();
218        let channel_router = self.channel_router.clone();
219
220        log::info!("Local port forwarding: {} -> {}:{}", bind_addr, remote_host, remote_port);
221
222        tokio::spawn(async move {
223            loop {
224                match listener.accept().await {
225                    Ok((stream, peer_addr)) => {
226                        log::debug!("Accepted connection from {} for forwarding", peer_addr);
227
228                        let transport = transport.clone();
229                        let remote_host = remote_host.clone();
230                        let channel_router = channel_router.clone();
231
232                        tokio::spawn(async move {
233                            if let Err(e) = handle_local_forward(
234                                stream,
235                                transport,
236                                remote_host,
237                                remote_port,
238                                channel_router,
239                            ).await {
240                                log::error!("Forward error: {}", e);
241                            }
242                        });
243                    }
244                    Err(e) => {
245                        log::error!("Accept error: {}", e);
246                    }
247                }
248            }
249        });
250
251        Ok(())
252    }
253    
254    /// Start all configured forwards
255    pub async fn start_all(&mut self) -> Result<()> {
256        for forward in self.forwards.clone() {
257            match forward {
258                ForwardType::Local { bind_addr, remote_host, remote_port } => {
259                    self.start_local_forward(bind_addr, remote_host, remote_port).await?;
260                }
261                ForwardType::Remote { remote_bind_addr, local_host, local_port } => {
262                    self.start_remote_forward(remote_bind_addr, local_host, local_port).await?;
263                }
264                ForwardType::Dynamic { bind_addr } => {
265                    self.start_socks_proxy(bind_addr).await?;
266                }
267            }
268        }
269        Ok(())
270    }
271
272    /// Start remote port forwarding (-R)
273    ///
274    /// Sends a GlobalRequest(TcpipForward) to the server asking it to listen on
275    /// `remote_bind_addr`. When the server accepts connections on that port, it
276    /// opens ForwardedTcpip channels back to us; the client-side handler connects
277    /// those to `local_host:local_port`.
278    pub async fn start_remote_forward(
279        &self,
280        remote_bind_addr: SocketAddr,
281        local_host: String,
282        local_port: u16,
283    ) -> Result<()> {
284        let bind_host = remote_bind_addr.ip().to_string();
285        let bind_port = remote_bind_addr.port();
286
287        log::info!("Requesting remote forward: {}:{} (server) -> {}:{} (local)",
288            bind_host, bind_port, local_host, local_port);
289
290        // Send TcpipForward global request
291        let request = Message::GlobalRequest(GlobalRequestMessage {
292            request_type: GlobalRequestType::TcpipForward {
293                bind_host: bind_host.clone(),
294                bind_port,
295            },
296            want_reply: true,
297        });
298        self.transport.send_message(&request).await?;
299
300        // Wait for success or failure
301        let reply = self.transport.receive_message::<Message>().await?;
302        match reply {
303            Message::GlobalRequestSuccess(success) => {
304                let actual_port = success.bound_port;
305                log::info!("Remote forward established: server listening on {}:{}",
306                    bind_host, actual_port);
307
308                // Register mapping so handle_forwarded_channel knows where to connect
309                self.remote_registry.insert(
310                    bind_host, actual_port, local_host, local_port,
311                ).await;
312
313                Ok(())
314            }
315            Message::GlobalRequestFailure => {
316                Err(QsshError::Protocol(format!(
317                    "Server refused remote forward on {}:{}", bind_host, bind_port
318                )))
319            }
320            other => {
321                Err(QsshError::Protocol(format!(
322                    "Unexpected reply to TcpipForward request: {:?}", other
323                )))
324            }
325        }
326    }
327    
328    /// Start SOCKS proxy for dynamic forwarding
329    pub async fn start_socks_proxy(&self, bind_addr: SocketAddr) -> Result<()> {
330        let listener = TcpListener::bind(bind_addr).await?;
331        let transport = self.transport.clone();
332        let channel_router = self.channel_router.clone();
333
334        log::info!("SOCKS proxy listening on {}", bind_addr);
335
336        tokio::spawn(async move {
337            loop {
338                match listener.accept().await {
339                    Ok((stream, peer_addr)) => {
340                        log::debug!("SOCKS connection from {}", peer_addr);
341
342                        let transport = transport.clone();
343                        let channel_router = channel_router.clone();
344                        tokio::spawn(async move {
345                            if let Err(e) = handle_socks_connection(stream, transport, channel_router).await {
346                                log::error!("SOCKS error: {}", e);
347                            }
348                        });
349                    }
350                    Err(e) => {
351                        log::error!("SOCKS accept error: {}", e);
352                    }
353                }
354            }
355        });
356
357        Ok(())
358    }
359    
360    /// Add a forward configuration
361    pub fn add_forward(&mut self, forward: ForwardType) {
362        self.forwards.push(forward);
363    }
364}
365
366/// Handle a local forward connection
367async fn handle_local_forward(
368    local_stream: TcpStream,
369    transport: Arc<Transport>,
370    remote_host: String,
371    remote_port: u16,
372    channel_router: ForwardedChannelRouter,
373) -> Result<()> {
374    let channel_id = rand::random::<u32>() % 65536;
375
376    // Register with channel router to receive server responses
377    let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(256);
378    channel_router.register(channel_id, data_tx).await;
379
380    // Open DirectTcpip channel — includes target in the channel type
381    let open_msg = Message::Channel(ChannelMessage::Open {
382        channel_id,
383        channel_type: ChannelType::DirectTcpip {
384            host: remote_host.clone(),
385            port: remote_port,
386            originator_host: "127.0.0.1".to_string(),
387            originator_port: local_stream.local_addr()
388                .map(|a| a.port()).unwrap_or(0),
389        },
390        window_size: 1024 * 1024,
391        max_packet_size: 32768,
392    });
393    transport.send_message(&open_msg).await?;
394
395    // Wait for Accept from server via router (empty vec = accept signal)
396    log::debug!("Local forward channel {} waiting for Accept", channel_id);
397    match tokio::time::timeout(std::time::Duration::from_secs(10), data_rx.recv()).await {
398        Ok(Some(data)) if data.is_empty() => {
399            log::debug!("Local forward channel {} accepted", channel_id);
400        }
401        Ok(Some(_data)) => {
402            // Got data before accept — server may have started sending immediately
403            log::debug!("Local forward channel {} got data before explicit accept", channel_id);
404            // Re-queue this data by just proceeding — the write_half will consume it below
405            // Actually, just start the bridge — we'll handle data inline
406        }
407        Ok(None) => {
408            log::error!("Local forward channel {} router closed before accept", channel_id);
409            channel_router.remove(channel_id).await;
410            return Err(QsshError::Protocol("Channel closed before accept".into()));
411        }
412        Err(_) => {
413            log::error!("Local forward channel {} timed out waiting for accept", channel_id);
414            channel_router.remove(channel_id).await;
415            return Err(QsshError::Protocol("Timeout waiting for channel accept".into()));
416        }
417    }
418
419    // Bridge local TCP <-> channel data bidirectionally
420    let (mut read_half, mut write_half) = local_stream.into_split();
421
422    // Local -> Remote: read from TCP, send as channel data
423    let transport_out = transport.clone();
424    let local_to_remote = tokio::spawn(async move {
425        let mut buffer = vec![0u8; 8192];
426        loop {
427            match read_half.read(&mut buffer).await {
428                Ok(0) => break,
429                Ok(n) => {
430                    let msg = Message::Channel(ChannelMessage::Data {
431                        channel_id,
432                        data: buffer[..n].to_vec(),
433                    });
434                    if transport_out.send_message(&msg).await.is_err() {
435                        break;
436                    }
437                }
438                Err(_) => break,
439            }
440        }
441    });
442
443    // Remote -> Local: receive from router, write to TCP
444    let remote_to_local = tokio::spawn(async move {
445        while let Some(data) = data_rx.recv().await {
446            if data.is_empty() { continue; } // skip any stray accept signals
447            if write_half.write_all(&data).await.is_err() {
448                break;
449            }
450        }
451    });
452
453    // Wait for either direction to finish
454    tokio::select! {
455        _ = local_to_remote => {}
456        _ = remote_to_local => {}
457    }
458
459    // Clean up
460    channel_router.remove(channel_id).await;
461    let eof_msg = Message::Channel(ChannelMessage::Eof { channel_id });
462    let _ = transport.send_message(&eof_msg).await;
463
464    Ok(())
465}
466
467/// Handle an incoming ForwardedTcpip channel from the server (remote forward, client side).
468///
469/// The server opened a channel because someone connected to the remotely-forwarded port.
470/// We look up the registry to find the local target, connect to it, accept the channel,
471/// register the channel in the router so the message loop can feed data, and bridge
472/// data bidirectionally.
473pub async fn handle_forwarded_channel(
474    channel_id: u32,
475    connected_host: String,
476    connected_port: u16,
477    transport: Arc<Transport>,
478    registry: RemoteForwardRegistry,
479    router: ForwardedChannelRouter,
480) -> Result<()> {
481    // Look up local target
482    let (local_host, local_port) = registry
483        .lookup(&connected_host, connected_port)
484        .await
485        .ok_or_else(|| QsshError::Protocol(format!(
486            "No remote forward registered for {}:{}", connected_host, connected_port
487        )))?;
488
489    log::info!("Forwarded channel {} for {}:{} -> connecting to {}:{}",
490        channel_id, connected_host, connected_port, local_host, local_port);
491
492    // Connect to local target
493    let local_addr = format!("{}:{}", local_host, local_port);
494    let local_stream = TcpStream::connect(&local_addr).await
495        .map_err(|e| QsshError::Connection(format!(
496            "Failed to connect to local target {}: {}", local_addr, e
497        )))?;
498
499    // Accept the channel
500    let accept = Message::Channel(ChannelMessage::Accept {
501        channel_id,
502        sender_channel: channel_id,
503        window_size: 1024 * 1024,
504        max_packet_size: 32768,
505    });
506    transport.send_message(&accept).await?;
507
508    // Bridge: local TCP stream <-> channel data (bidirectional)
509    let (mut read_half, mut write_half) = local_stream.into_split();
510
511    // Register an mpsc sender so the client's message loop can feed channel data to us
512    let (tx, mut rx) = mpsc::channel::<Vec<u8>>(256);
513    router.register(channel_id, tx).await;
514
515    // Local TCP -> channel (read from local, send as channel data to server)
516    let transport_send = transport.clone();
517    let local_to_channel = tokio::spawn(async move {
518        let mut buffer = vec![0u8; 8192];
519        loop {
520            match read_half.read(&mut buffer).await {
521                Ok(0) => break,
522                Ok(n) => {
523                    let data_msg = Message::Channel(ChannelMessage::Data {
524                        channel_id,
525                        data: buffer[..n].to_vec(),
526                    });
527                    if transport_send.send_message(&data_msg).await.is_err() {
528                        break;
529                    }
530                }
531                Err(_) => break,
532            }
533        }
534    });
535
536    // Channel -> local TCP (receive data from router's mpsc, write to local TCP)
537    let channel_to_local = tokio::spawn(async move {
538        while let Some(data) = rx.recv().await {
539            if write_half.write_all(&data).await.is_err() {
540                break;
541            }
542        }
543    });
544
545    // Wait for either direction to finish
546    tokio::select! {
547        _ = local_to_channel => {}
548        _ = channel_to_local => {}
549    }
550
551    // Cleanup: remove from router, send EOF
552    router.remove(channel_id).await;
553    let eof_msg = Message::Channel(ChannelMessage::Eof { channel_id });
554    let _ = transport.send_message(&eof_msg).await;
555
556    Ok(())
557}
558
559/// Handle SOCKS5 connection
560async fn handle_socks_connection(
561    mut stream: TcpStream,
562    transport: Arc<Transport>,
563    channel_router: ForwardedChannelRouter,
564) -> Result<()> {
565    // SOCKS5 handshake
566    let mut buffer = vec![0u8; 1024];
567    
568    // Read version and methods
569    let n = stream.read(&mut buffer).await?;
570    if n < 3 || buffer[0] != 0x05 {
571        return Err(QsshError::Protocol("Invalid SOCKS5 handshake".into()));
572    }
573    
574    // Send no auth required
575    stream.write_all(&[0x05, 0x00]).await?;
576    
577    // Read connect request
578    let n = stream.read(&mut buffer).await?;
579    if n < 10 || buffer[0] != 0x05 || buffer[1] != 0x01 {
580        return Err(QsshError::Protocol("Invalid SOCKS5 connect request".into()));
581    }
582    
583    // Parse destination
584    let addr_type = buffer[3];
585    let (dest_host, dest_port) = match addr_type {
586        0x01 => {
587            // IPv4
588            let addr = format!("{}.{}.{}.{}", buffer[4], buffer[5], buffer[6], buffer[7]);
589            let port = u16::from_be_bytes([buffer[8], buffer[9]]);
590            (addr, port)
591        }
592        0x03 => {
593            // Domain name
594            let len = buffer[4] as usize;
595            let domain = String::from_utf8_lossy(&buffer[5..5+len]).to_string();
596            let port = u16::from_be_bytes([buffer[5+len], buffer[6+len]]);
597            (domain, port)
598        }
599        _ => return Err(QsshError::Protocol("Unsupported SOCKS5 address type".into())),
600    };
601    
602    // Send success response
603    stream.write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
604    
605    // Forward the connection
606    handle_local_forward(stream, transport, dest_host, dest_port, channel_router).await
607}
608
609#[cfg(test)]
610mod tests {
611    use super::*;
612
613    #[test]
614    fn test_parse_local_forward() {
615        let forward = PortForwardManager::parse_forward_spec("8080:localhost:80", "local").unwrap();
616
617        match forward {
618            ForwardType::Local { bind_addr, remote_host, remote_port } => {
619                assert_eq!(bind_addr.port(), 8080);
620                assert_eq!(remote_host, "localhost");
621                assert_eq!(remote_port, 80);
622            }
623            _ => panic!("Wrong forward type"),
624        }
625    }
626
627    #[test]
628    fn test_parse_remote_forward() {
629        let forward = PortForwardManager::parse_forward_spec("9090:localhost:8080", "remote").unwrap();
630
631        match forward {
632            ForwardType::Remote { remote_bind_addr, local_host, local_port } => {
633                assert_eq!(remote_bind_addr.port(), 9090);
634                assert_eq!(remote_bind_addr.ip().to_string(), "0.0.0.0");
635                assert_eq!(local_host, "localhost");
636                assert_eq!(local_port, 8080);
637            }
638            _ => panic!("Wrong forward type"),
639        }
640    }
641
642    #[test]
643    fn test_parse_remote_forward_different_ports() {
644        let forward = PortForwardManager::parse_forward_spec("443:127.0.0.1:3000", "remote").unwrap();
645
646        match forward {
647            ForwardType::Remote { remote_bind_addr, local_host, local_port } => {
648                assert_eq!(remote_bind_addr.port(), 443);
649                assert_eq!(local_host, "127.0.0.1");
650                assert_eq!(local_port, 3000);
651            }
652            _ => panic!("Wrong forward type"),
653        }
654    }
655
656    #[test]
657    fn test_parse_remote_forward_invalid() {
658        // Missing local port
659        assert!(PortForwardManager::parse_forward_spec("9090:localhost", "remote").is_err());
660        // Not a number
661        assert!(PortForwardManager::parse_forward_spec("abc:localhost:8080", "remote").is_err());
662    }
663
664    #[test]
665    fn test_parse_dynamic_forward() {
666        let forward = PortForwardManager::parse_forward_spec("1080", "dynamic").unwrap();
667
668        match forward {
669            ForwardType::Dynamic { bind_addr } => {
670                assert_eq!(bind_addr.port(), 1080);
671            }
672            _ => panic!("Wrong forward type"),
673        }
674    }
675
676    #[tokio::test]
677    async fn test_remote_forward_registry_insert_lookup() {
678        let registry = RemoteForwardRegistry::new();
679
680        // Insert a mapping
681        registry.insert("0.0.0.0".into(), 9090, "localhost".into(), 8080).await;
682
683        // Lookup should succeed
684        let result = registry.lookup("0.0.0.0", 9090).await;
685        assert_eq!(result, Some(("localhost".to_string(), 8080)));
686
687        // Lookup for non-existent should return None
688        let result = registry.lookup("0.0.0.0", 9999).await;
689        assert_eq!(result, None);
690    }
691
692    #[tokio::test]
693    async fn test_remote_forward_registry_remove() {
694        let registry = RemoteForwardRegistry::new();
695
696        registry.insert("0.0.0.0".into(), 9090, "localhost".into(), 8080).await;
697
698        // Remove should return the mapping
699        let removed = registry.remove("0.0.0.0", 9090).await;
700        assert_eq!(removed, Some(("localhost".to_string(), 8080)));
701
702        // Lookup should now fail
703        assert_eq!(registry.lookup("0.0.0.0", 9090).await, None);
704
705        // Remove again should return None
706        assert_eq!(registry.remove("0.0.0.0", 9090).await, None);
707    }
708
709    #[tokio::test]
710    async fn test_remote_forward_registry_multiple_entries() {
711        let registry = RemoteForwardRegistry::new();
712
713        registry.insert("0.0.0.0".into(), 9090, "localhost".into(), 8080).await;
714        registry.insert("0.0.0.0".into(), 9091, "localhost".into(), 3000).await;
715        registry.insert("127.0.0.1".into(), 443, "10.0.0.1".into(), 443).await;
716
717        assert_eq!(registry.lookup("0.0.0.0", 9090).await, Some(("localhost".to_string(), 8080)));
718        assert_eq!(registry.lookup("0.0.0.0", 9091).await, Some(("localhost".to_string(), 3000)));
719        assert_eq!(registry.lookup("127.0.0.1", 443).await, Some(("10.0.0.1".to_string(), 443)));
720    }
721
722    #[tokio::test]
723    async fn test_forwarded_channel_router() {
724        let router = ForwardedChannelRouter::new();
725
726        // No channel registered
727        assert!(!router.has_channel(1).await);
728        assert!(!router.route_data(1, vec![1, 2, 3]).await);
729
730        // Register a channel
731        let (tx, mut rx) = tokio::sync::mpsc::channel(16);
732        router.register(1, tx).await;
733
734        assert!(router.has_channel(1).await);
735
736        // Route data
737        assert!(router.route_data(1, vec![42, 43]).await);
738        let received = rx.recv().await.unwrap();
739        assert_eq!(received, vec![42, 43]);
740
741        // Remove channel
742        router.remove(1).await;
743        assert!(!router.has_channel(1).await);
744    }
745
746    #[tokio::test]
747    async fn test_forwarded_channel_router_multiple_channels() {
748        let router = ForwardedChannelRouter::new();
749
750        let (tx1, mut rx1) = tokio::sync::mpsc::channel(16);
751        let (tx2, mut rx2) = tokio::sync::mpsc::channel(16);
752
753        router.register(100, tx1).await;
754        router.register(200, tx2).await;
755
756        // Route to channel 100
757        assert!(router.route_data(100, vec![1]).await);
758        assert_eq!(rx1.recv().await.unwrap(), vec![1]);
759
760        // Route to channel 200
761        assert!(router.route_data(200, vec![2]).await);
762        assert_eq!(rx2.recv().await.unwrap(), vec![2]);
763
764        // Channel 300 doesn't exist
765        assert!(!router.route_data(300, vec![3]).await);
766    }
767
768    #[tokio::test]
769    async fn test_router_accept_signal() {
770        // Empty vec is used as the "channel accepted" signal
771        let router = ForwardedChannelRouter::new();
772        let (tx, mut rx) = tokio::sync::mpsc::channel(16);
773        router.register(42, tx).await;
774
775        // Send accept signal (empty vec)
776        assert!(router.route_data(42, Vec::new()).await);
777        let signal = rx.recv().await.unwrap();
778        assert!(signal.is_empty(), "Accept signal should be empty vec");
779
780        // Then send actual data
781        assert!(router.route_data(42, vec![0xDE, 0xAD]).await);
782        let data = rx.recv().await.unwrap();
783        assert_eq!(data, vec![0xDE, 0xAD]);
784    }
785
786    #[tokio::test]
787    async fn test_router_remove_drops_sender() {
788        let router = ForwardedChannelRouter::new();
789        let (tx, mut rx) = tokio::sync::mpsc::channel(16);
790        router.register(99, tx).await;
791
792        router.remove(99).await;
793
794        // After remove, routing should fail
795        assert!(!router.route_data(99, vec![1]).await);
796        // Receiver should get None (sender dropped)
797        assert!(rx.recv().await.is_none());
798    }
799}
qssh/port_forward.rs

qssh/
port_forward.rs
