70 lines
2.6 KiB
Rust
70 lines
2.6 KiB
Rust
use std::io;
|
|
use std::net::{SocketAddr, SocketAddrV4};
|
|
use std::sync::Arc;
|
|
use crossbeam_skiplist::SkipMap;
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
pub struct TcpProxy {
|
|
tcp_listener: TcpListener,
|
|
// todo 怎么过期
|
|
map: Arc<SkipMap<SocketAddrV4, (SocketAddrV4, SocketAddrV4)>>,
|
|
}
|
|
|
|
impl TcpProxy {
|
|
pub fn new(tcp_listener: TcpListener, map: Arc<SkipMap<SocketAddrV4, (SocketAddrV4, SocketAddrV4)>>) -> Self {
|
|
Self {
|
|
tcp_listener,
|
|
map,
|
|
}
|
|
}
|
|
pub async fn start(self) {
|
|
let tcp_listener = self.tcp_listener;
|
|
let map = self.map;
|
|
loop {
|
|
match tcp_listener.accept().await {
|
|
Ok((tcp_stream, sender_addr)) => {
|
|
match sender_addr {
|
|
SocketAddr::V4(sender_addr) => {
|
|
if let Some(entry) = map.get(&sender_addr) {
|
|
let (src_addr, dest_addr) = *entry.value();
|
|
let peer_tcp_stream = match TcpStream::connect(dest_addr).await {
|
|
Ok(peer_tcp_stream) => {peer_tcp_stream}
|
|
Err(e) => {
|
|
log::warn!("tcp代理异常:{:?},来源:{},目标：{}",e,src_addr,dest_addr);
|
|
continue;
|
|
}
|
|
};
|
|
tokio::spawn(async move {
|
|
match proxy(tcp_stream, peer_tcp_stream).await {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
log::warn!("tcp代理异常:{:?},来源:{},目标：{}",e,src_addr,dest_addr);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
SocketAddr::V6(_) => {}
|
|
}
|
|
|
|
}
|
|
Err(e) => {
|
|
log::warn!("tcp代理监听:{:?}",e);
|
|
}
|
|
}
|
|
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn proxy(mut client: TcpStream, mut server: TcpStream) -> io::Result<()> {
|
|
let (mut client_reader, mut client_writer) = client.split();
|
|
let (mut server_reader, mut server_writer) = server.split();
|
|
|
|
let client_to_server = tokio::io::copy(&mut client_reader, &mut server_writer);
|
|
let server_to_client = tokio::io::copy(&mut server_reader, &mut client_writer);
|
|
|
|
tokio::try_join!(client_to_server, server_to_client)?;
|
|
|
|
Ok(())
|
|
} |