refine SocksAddressType

This commit is contained in:
ssrlive 2023-04-10 18:26:32 +08:00 committed by B. Blechschmidt
parent fd48be5feb
commit c0cff1da58

View file

@ -6,6 +6,7 @@ use crate::tun2proxy::{
use crate::Credentials; use crate::Credentials;
use smoltcp::wire::IpProtocol; use smoltcp::wire::IpProtocol;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::convert::TryFrom;
use std::net::{IpAddr, SocketAddr}; use std::net::{IpAddr, SocketAddr};
use std::rc::Rc; use std::rc::Rc;
@ -22,13 +23,31 @@ enum SocksState {
} }
#[repr(u8)] #[repr(u8)]
#[derive(Copy, Clone)] #[derive(Copy, Clone, PartialEq, Debug)]
enum SocksAddressType { enum SocksAddressType {
Ipv4 = 1, Ipv4 = 1,
DomainName = 3, DomainName = 3,
Ipv6 = 4, Ipv6 = 4,
} }
impl TryFrom<u8> for SocksAddressType {
type Error = Error;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
1 => Ok(SocksAddressType::Ipv4),
3 => Ok(SocksAddressType::DomainName),
4 => Ok(SocksAddressType::Ipv6),
_ => Err(format!("Unknown address type: {}", value).into()),
}
}
}
impl From<SocksAddressType> for u8 {
fn from(value: SocksAddressType) -> Self {
value as u8
}
}
#[derive(Copy, Clone)] #[derive(Copy, Clone)]
pub enum SocksVersion { pub enum SocksVersion {
V4 = 4, V4 = 4,
@ -82,11 +101,11 @@ impl SocksConnection {
let mut result = Self { let mut result = Self {
connection: connection.clone(), connection: connection.clone(),
state: SocksState::ServerHello, state: SocksState::ServerHello,
client_inbuf: Default::default(), client_inbuf: VecDeque::default(),
server_inbuf: Default::default(), server_inbuf: VecDeque::default(),
client_outbuf: Default::default(), client_outbuf: VecDeque::default(),
server_outbuf: Default::default(), server_outbuf: VecDeque::default(),
data_buf: Default::default(), data_buf: VecDeque::default(),
manager, manager,
version, version,
}; };
@ -235,29 +254,18 @@ impl SocksConnection {
return Err("SOCKS5 connection unsuccessful.".into()); return Err("SOCKS5 connection unsuccessful.".into());
} }
if atyp != SocksAddressType::Ipv4 as u8 let message_length = match SocksAddressType::try_from(atyp)? {
&& atyp != SocksAddressType::Ipv6 as u8 SocksAddressType::DomainName => {
&& atyp != SocksAddressType::DomainName as u8 if self.server_inbuf.len() < 5 {
{
return Err("SOCKS5 server replied with unrecognized address type.".into());
}
if atyp == SocksAddressType::DomainName as u8 && self.server_inbuf.len() < 5 {
return Ok(()); return Ok(());
} }
if self.server_inbuf.len() < 7 + (self.server_inbuf[4] as usize) {
if atyp == SocksAddressType::DomainName as u8
&& self.server_inbuf.len() < 7 + (self.server_inbuf[4] as usize)
{
return Ok(()); return Ok(());
} }
let message_length = if atyp == SocksAddressType::Ipv4 as u8 {
10
} else if atyp == SocksAddressType::Ipv6 as u8 {
22
} else {
7 + (self.server_inbuf[4] as usize) 7 + (self.server_inbuf[4] as usize)
}
SocksAddressType::Ipv4 => 10,
SocksAddressType::Ipv6 => 22,
}; };
self.server_inbuf.drain(0..message_length); self.server_inbuf.drain(0..message_length);
@ -277,7 +285,7 @@ impl SocksConnection {
} else { } else {
SocksAddressType::Ipv6 SocksAddressType::Ipv6
}; };
self.server_outbuf.extend(&[cmd as u8]); self.server_outbuf.extend(&[u8::from(cmd)]);
match dst_ip { match dst_ip {
IpAddr::V4(ip) => self.server_outbuf.extend(ip.octets().as_ref()), IpAddr::V4(ip) => self.server_outbuf.extend(ip.octets().as_ref()),
IpAddr::V6(ip) => self.server_outbuf.extend(ip.octets().as_ref()), IpAddr::V6(ip) => self.server_outbuf.extend(ip.octets().as_ref()),
@ -285,7 +293,7 @@ impl SocksConnection {
} }
DestinationHost::Hostname(host) => { DestinationHost::Hostname(host) => {
self.server_outbuf self.server_outbuf
.extend(&[SocksAddressType::DomainName as u8, host.len() as u8]); .extend(&[u8::from(SocksAddressType::DomainName), host.len() as u8]);
self.server_outbuf.extend(host.as_bytes()); self.server_outbuf.extend(host.as_bytes());
} }
} }