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