UDP associate

This commit is contained in:
ssrlive 2023-08-11 19:18:18 +08:00
parent 46ca342aba
commit da87fa8d5a
3 changed files with 34 additions and 6 deletions

View file

@ -382,6 +382,10 @@ impl TcpProxy for HttpConnection {
fn reset_connection(&self) -> bool {
self.state == HttpState::Reset
}
fn get_udp_associate(&self) -> Option<SocketAddr> {
None
}
}
pub(crate) struct HttpManager {

View file

@ -224,7 +224,9 @@ impl SocksProxyImpl {
return Err(format!("SOCKS connection failed: {}", response.reply).into());
}
if self.command == protocol::Command::UdpAssociate {
self.udp_associate = Some(SocketAddr::try_from(response.address)?);
self.udp_associate = Some(SocketAddr::try_from(&response.address)?);
assert!(self.data_buf.is_empty());
log::debug!("UDP associate: {}", response.address);
}
self.server_outbuf.append(&mut self.data_buf);
@ -326,6 +328,10 @@ impl TcpProxy for SocksProxyImpl {
fn reset_connection(&self) -> bool {
false
}
fn get_udp_associate(&self) -> Option<SocketAddr> {
self.udp_associate
}
}
pub(crate) struct SocksProxyManager {

View file

@ -7,7 +7,7 @@ use smoltcp::{
time::Instant,
wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket, UDP_HEADER_LEN},
};
use socks5_impl::protocol::{Address, UserKey};
use socks5_impl::protocol::{Address, StreamOperation, UdpHeader, UserKey};
use std::{
collections::{HashMap, HashSet},
convert::{From, TryFrom},
@ -184,6 +184,7 @@ pub(crate) trait TcpProxy {
fn connection_established(&self) -> bool;
fn have_data(&mut self, dir: Direction) -> bool;
fn reset_connection(&self) -> bool;
fn get_udp_associate(&self) -> Option<SocketAddr>;
}
pub(crate) trait UdpProxy {
@ -496,11 +497,28 @@ impl<'a> TunToProxy<'a> {
}
} else {
// Another UDP packet
let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info, true)?;
let state = self.create_new_tcp_connection_state(server_addr, dst, tcp_proxy_handler)?;
self.connection_map.insert(connection_info.clone(), state);
if !self.connection_map.contains_key(&connection_info) {
log::trace!("New UDP connection {} ({})", connection_info, dst);
let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info, true)?;
let state = self.create_new_tcp_connection_state(server_addr, dst, tcp_proxy_handler)?;
self.connection_map.insert(connection_info.clone(), state);
}
// TODO: Handle UDP packets
self.expect_smoltcp_send()?;
self.tunsocket_read_and_forward(&connection_info)?;
self.write_to_server(&connection_info)?;
let mut s5_udp_data = Vec::<u8>::new();
UdpHeader::new(0, connection_info.dst.clone()).write_to_stream(&mut s5_udp_data)?;
s5_udp_data.extend_from_slice(payload);
let state = self.connection_map.get(&connection_info).ok_or("udp associate state")?;
if let Some(udp_associate) = state.tcp_proxy_handler.get_udp_associate() {
log::debug!("UDP associate address: {}", udp_associate);
// Send packets via UDP associate...
} else {
// UDP associate tunnel not ready yet, we must cache the packet...
}
}
} else {
log::warn!("Unsupported protocol: {} ({})", connection_info, dst);