内网穿透技术-- rust

内网穿透技术：很简单。

---

技术实现： A请求B服务，B服务在自己家里部署。 （A是不能直接连接B的） 此时，需要有一个角色C， 部署在公网上。B请求C ， 然后维持一个会话。 A请求C，C在缓存的会话中，找到B，转发流量即可。

• 啥协议都可以，这里使用的tcp
• 需要维护心跳，还可以加一些其他的机制， 比如， 安全啊， 流量方面的考虑。

---

中继服务器代码

use std::net::{TcpListener, TcpStream};
use std::thread;
use std::io::{self, Read, Write};
use std::time::{Duration, Instant};
use std::sync::{Arc, Mutex};
use std::collections::HashMap;

const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10);
const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(30);

type SharedState = Arc<Mutex<HashMap<String, (TcpStream, Instant)>>>;

fn forward(mut from: TcpStream, mut to: TcpStream) -> io::Result<()> {
    let mut buffer = [0; 4096];
    loop {
        let bytes_read = from.read(&mut buffer)?;
        if bytes_read == 0 {
            return Ok(());
        }
        to.write_all(&buffer[..bytes_read])?;
    }
}

fn handle_client(client_stream: TcpStream, target_address: String, state: SharedState) -> io::Result<()> {
    let map = state.lock().unwrap();
    if let Some((server_stream, _)) = map.get(&target_address) {
        let client_stream_clone = client_stream.try_clone()?;
        let server_stream_clone = server_stream.try_clone()?;

        let client_to_server = thread::spawn(move || {
            forward(client_stream_clone, server_stream_clone).unwrap_or_else(|err| {
                eprintln!("Client to Server error: {:?}", err);
            });
        });

        let server_to_client = thread::spawn(move || {
            forward(server_stream.try_clone().unwrap(), client_stream).unwrap_or_else(|err| {
                eprintln!("Server to Client error: {:?}", err);
            });
        });

        client_to_server.join().unwrap();
        server_to_client.join().unwrap();
    } else {
        eprintln!("Target address not available: {}", target_address);
    }
    Ok(())
}

fn heartbeat_checker(state: SharedState) {
    let interval = std::time::Duration::from_secs(5);
    loop {
        thread::sleep(interval);
        let mut map = state.lock().unwrap();
        let now = Instant::now();
        map.retain(|addr, (stream, last_heartbeat)| {
            if now.duration_since(*last_heartbeat) > HEARTBEAT_TIMEOUT {
                eprintln!("Removing inactive client: {}", addr);
                false
            } else {
                true
            }
        });
    }
}

fn main() -> io::Result<()> {
    let listener = TcpListener::bind("0.0.0.0:4000")?;
    let state = Arc::new(Mutex::new(HashMap::new()));
    let state_clone = Arc::clone(&state);

    // 心跳检测线程
    thread::spawn(move || {
        heartbeat_checker(state_clone);
    });

    println!("Relay server listening on port 4000");

    for stream in listener.incoming() {
        match stream {
            Ok(mut client_stream) => {
                let mut buf = [0; 256];
                let size = client_stream.read(&mut buf)?;
                let initial_msg = String::from_utf8_lossy(&buf[..size]);

                let parts: Vec<&str> = initial_msg.trim().split(' ').collect();
                if parts.len() == 2 {
                    let command = parts[0];
                    let address = parts[1].to_string();

                    if command == "REGISTER" {
                        let mut map = state.lock().unwrap();
                        map.insert(address.clone(), (client_stream.try_clone().unwrap(), Instant::now()));
                        println!("Registered service from: {}", address);

                        // 处理心跳消息
                        let state = Arc::clone(&state);
                        thread::spawn(move || {
                            loop {
                                let mut heartbeat_buf = [0; 10];
                                match client_stream.read(&mut heartbeat_buf) {
                                    Ok(size) if size > 0 => {
                                        let message = String::from_utf8_lossy(&heartbeat_buf[..size]);
                                        if message.trim() == "HEARTBEAT" {
                                            let mut map = state.lock().unwrap();
                                            if let Some((_, last_heartbeat)) = map.get_mut(&address) {
                                                *last_heartbeat = Instant::now();
                                            }
                                        }
                                    }
                                    Ok(_) => break,
                                    Err(err) => {
                                        eprintln!("Heartbeat error: {:?}", err);
                                        break;
                                    }
                                }
                            }
                        });
                    } else if command == "CONNECT" {
                        let state = Arc::clone(&state);
                        thread::spawn(move || {
                            handle_client(client_stream, address, state).unwrap_or_else(|err| {
                                eprintln!("Failed to handle client: {:?}", err);
                            });
                        });
                    } else {
                        eprintln!("Unknown command received: {}", command);
                    }
                } else {
                    eprintln!("Invalid initial message received");
                }
            }
            Err(e) => eprintln!("Connection failed: {}", e),
        }
    }

    Ok(())
}

---

use std::net::{TcpStream, TcpListener};
use std::thread;
use std::io::{self, Read, Write};
use std::sync::{Arc, Mutex};
use std::time::Duration;

const RELAY_SERVER: &str = "public_server_ip:4000"; // 替换为你的中继服务器的公共 IP 地址
const LOCAL_SERVICE: &str = "127.0.0.1:3000"; // 内网服务地址
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10);

fn forward(mut from: TcpStream, mut to: TcpStream) -> io::Result<()> {
    let mut buffer = [0; 4096];
    loop {
        let bytes_read = from.read(&mut buffer)?;
        if bytes_read == 0 {
            return Ok(());
        }
        to.write_all(&buffer[..bytes_read])?;
    }
}

fn main() -> io::Result<()> {
    // 连接到 Relay Server
    let relay_stream = TcpStream::connect(RELAY_SERVER)?;
    let relay_stream = Arc::new(Mutex::new(relay_stream));

    // 注册内网服务地址
    {
        let mut relay = relay_stream.lock().unwrap();
        relay.write_all(format!("REGISTER {}", LOCAL_SERVICE).as_bytes())?;
    }

    // 启动心跳线程
    let relay_clone = Arc::clone(&relay_stream);
    thread::spawn(move || {
        loop {
            thread::sleep(HEARTBEAT_INTERVAL);
            let mut relay = relay_clone.lock().unwrap();
            if let Err(e) = relay.write_all(b"HEARTBEAT") {
                eprintln!("Failed to send heartbeat: {}", e);
                break;
            }
        }
    });

    // 建立本地监听器，等待中继服务器的连接
    let listener = TcpListener::bind(LOCAL_SERVICE)?;
    for stream in listener.incoming() {
        match stream {
            Ok(local_stream) => {
                let relay_clone = Arc::clone(&relay_stream);
                let local_clone = local_stream.try_clone()?;

                let relay_to_local = thread::spawn(move || {
                    forward(relay_clone.lock().unwrap().try_clone().unwrap(), local_clone).unwrap_or_else(|err| {
                        eprintln!("Relay to Local error: {:?}", err);
                    });
                });

                let local_to_relay = thread::spawn(move || {
                    forward(local_stream, relay_clone.lock().unwrap().try_clone().unwrap()).unwrap_or_else(|err| {
                        eprintln!("Local to Relay error: {:?}", err);
                    });
                });

                relay_to_local.join().unwrap();
                local_to_relay.join().unwrap();
            }
            Err(e) => eprintln!("Connection failed: {}", e),
        }
    }

    Ok(())
}