porta-rs 0.1.1

Zero-trust CGNAT bypass via WireGuard tunnel
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

use anyhow::{Context, Result};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use std::net::SocketAddr;
use std::sync::Arc;
use parking_lot::RwLock;

use crate::config::PortMapping;

/// TCP connection statistics
#[derive(Debug, Default, Clone)]
pub struct TcpStats {
    pub connections_active: u64,
    pub connections_total: u64,
    pub bytes_sent: u64,
    pub bytes_received: u64,
}

/// TCP forwarder
pub struct TcpForwarder {
    stats: Arc<RwLock<TcpStats>>,
}

impl TcpForwarder {
    pub fn new() -> Self {
        Self {
            stats: Arc::new(RwLock::new(TcpStats::default())),
        }
    }

    /// Start forwarding for a port mapping
    pub async fn start_forwarding(&self, mapping: &PortMapping) -> Result<()> {
        let listener = TcpListener::bind(format!("0.0.0.0:{}", mapping.remote_port))
            .await
            .context(format!("Failed to bind TCP port {}", mapping.remote_port))?;

        tracing::info!(
            "TCP forwarder listening on :{} -> localhost:{}",
            mapping.remote_port,
            mapping.local_port
        );

        let stats = self.stats.clone();
        let local_port = mapping.local_port;

        tokio::spawn(async move {
            loop {
                match listener.accept().await {
                    Ok((stream, addr)) => {
                        let stats = stats.clone();
                        tokio::spawn(async move {
                            if let Err(e) = Self::handle_connection(stream, addr, local_port, stats).await {
                                tracing::error!("TCP connection error: {}", e);
                            }
                        });
                    }
                    Err(e) => {
                        tracing::error!("Failed to accept TCP connection: {}", e);
                    }
                }
            }
        });

        Ok(())
    }

    /// Handle a single TCP connection
    async fn handle_connection(
        client_stream: TcpStream,
        client_addr: SocketAddr,
        local_port: u16,
        stats: Arc<RwLock<TcpStats>>,
    ) -> Result<()> {
        // Connect to local service
        let local_addr = format!("127.0.0.1:{}", local_port);
        let local_stream = TcpStream::connect(&local_addr)
            .await
            .context(format!("Failed to connect to {}", local_addr))?;

        tracing::debug!("TCP connection: {} -> :{} -> {}", client_addr, local_port, local_addr);

        // Update stats
        {
            let mut s = stats.write();
            s.connections_active += 1;
            s.connections_total += 1;
        }

        // Bidirectional copy
        let (mut client_read, mut client_write) = client_stream.into_split();
        let (mut local_read, mut local_write) = local_stream.into_split();

        let stats_send = stats.clone();
        let stats_recv = stats.clone();

        let send_task = tokio::spawn(async move {
            let mut buf = vec![0u8; 8192];
            let mut total = 0u64;
            loop {
                match client_read.read(&mut buf).await {
                    Ok(0) => break,
                    Ok(n) => {
                        if local_write.write_all(&buf[..n]).await.is_err() {
                            break;
                        }
                        total += n as u64;
                    }
                    Err(_) => break,
                }
            }
            let mut s = stats_send.write();
            s.bytes_sent += total;
            s.connections_active = s.connections_active.saturating_sub(1);
        });

        let recv_task = tokio::spawn(async move {
            let mut buf = vec![0u8; 8192];
            let mut total = 0u64;
            loop {
                match local_read.read(&mut buf).await {
                    Ok(0) => break,
                    Ok(n) => {
                        if client_write.write_all(&buf[..n]).await.is_err() {
                            break;
                        }
                        total += n as u64;
                    }
                    Err(_) => break,
                }
            }
            let mut s = stats_recv.write();
            s.bytes_received += total;
        });

        tokio::select! {
            _ = send_task => {},
            _ = recv_task => {},
        }

        Ok(())
    }

    /// Get current statistics
    pub fn stats(&self) -> TcpStats {
        self.stats.read().clone()
    }
}

impl Default for TcpForwarder {
    fn default() -> Self {
        Self::new()
    }
}