Skip to content

Commit 1723645

Browse files
brianmaybrianmay
committed
Add UDP support framework
1 parent aa6ad0e commit 1723645

File tree

8 files changed

+507
-132
lines changed

8 files changed

+507
-132
lines changed

‎.vscode/settings.json

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,13 @@
1010
"getaddrinfo",
1111
"getsockopt",
1212
"INET",
13+
"IPPROTO",
14+
"IPPROTO_IP",
15+
"libc",
1316
"lport",
1417
"noncommercially",
1518
"PREROUTING",
19+
"RECVORIGDSTADDR",
1620
"reprs",
1721
"rustfmt",
1822
"setsocketopt",

‎src/client.rs

Lines changed: 68 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
1-
use std::net::IpAddr;
1+
use std::io::IoSliceMut;
2+
use std::net::{IpAddr, UdpSocket};
3+
use std::os::unix::prelude::AsRawFd;
24
use std::sync::Arc;
35
use std::time::Duration;
46
use std::{error::Error, fmt::Display, net::SocketAddr};
57

68
use fast_socks5::client::Socks5Stream;
79

10+
use nix::sys::socket::{recvmsg, MsgFlags, RecvMsg};
811
use tokio::io::copy_bidirectional;
912
use tokio::net::{TcpListener, TcpStream};
1013
use tokio::select;
@@ -17,14 +20,14 @@ use crate::command::CommandError;
1720
use crate::firewall::{
1821
Firewall, FirewallConfig, FirewallError, FirewallListenerConfig, FirewallSubnetConfig,
1922
};
20-
use crate::network::Subnets;
23+
use crate::network::{ListenerAddr, Subnets};
2124
use crate::options::FirewallType;
2225

2326
pub struct Config {
2427
pub includes: Subnets,
2528
pub excludes: Subnets,
2629
pub remote: Option<String>,
27-
pub listen: Vec<SocketAddr>,
30+
pub listen: Vec<ListenerAddr>,
2831
pub socks_addr: SocketAddr,
2932
pub firewall: FirewallType,
3033
}
@@ -225,13 +228,13 @@ fn get_firewall_config(config: &Config) -> FirewallConfig {
225228
.map(|addr| match addr.ip() {
226229
IpAddr::V4(_) => FirewallListenerConfig::Ipv4(FirewallSubnetConfig {
227230
enable: true,
228-
port: addr.port(),
231+
listener: addr.clone(),
229232
includes: config.includes.ipv4(),
230233
excludes: config.excludes.ipv4(),
231234
}),
232235
IpAddr::V6(_) => FirewallListenerConfig::Ipv6(FirewallSubnetConfig {
233236
enable: true,
234-
port: addr.port(),
237+
listener: addr.clone(),
235238
includes: config.includes.ipv6(),
236239
excludes: config.excludes.ipv6(),
237240
}),
@@ -307,32 +310,74 @@ async fn run_client(
307310
let listen = config.listen.clone();
308311

309312
let firewall: Arc<dyn Firewall + Send + Sync> = Arc::from(firewall);
310-
for addr in listen {
311-
let firewall = Arc::clone(&firewall);
312-
let listener = TcpListener::bind(addr).await?;
313-
firewall.setup_tcp_listener(&listener)?;
314-
315-
let _handle = tokio::spawn(async move {
316-
firewall.setup_tcp_listener(&listener).unwrap();
317-
318-
loop {
319-
let firewall = Arc::clone(&firewall);
320-
let (socket, _) = listener.accept().await.unwrap();
321-
tokio::spawn(async move {
322-
handle_tcp_client(socket, addr, socks_addr, firewall).await;
323-
});
324-
}
325-
});
313+
for l_addr in listen {
314+
println!("----> {}", l_addr);
315+
match l_addr.protocol {
316+
crate::network::Protocol::Tcp => listen_tcp(&firewall, l_addr, socks_addr).await?,
317+
crate::network::Protocol::Udp => listen_udp(&firewall, l_addr, socks_addr).await?,
318+
}
326319
}
327320

328321
loop {
329322
sleep(Duration::from_secs(60)).await;
330323
}
331324
}
332325

326+
async fn listen_tcp(
327+
firewall: &Arc<dyn Firewall + Send + Sync>,
328+
l_addr: ListenerAddr,
329+
socks_addr: SocketAddr,
330+
) -> Result<(), ClientError> {
331+
let firewall = Arc::clone(firewall);
332+
let listener = TcpListener::bind(l_addr.addr).await?;
333+
firewall.setup_tcp_listener(&listener)?;
334+
335+
let _handle = tokio::spawn(async move {
336+
loop {
337+
let firewall = Arc::clone(&firewall);
338+
let (socket, _) = listener.accept().await.unwrap();
339+
let l_addr = l_addr.clone();
340+
tokio::spawn(async move {
341+
handle_tcp_client(socket, &l_addr, socks_addr, firewall).await;
342+
});
343+
}
344+
});
345+
Ok(())
346+
}
347+
348+
async fn listen_udp(
349+
firewall: &Arc<dyn Firewall + Send + Sync>,
350+
l_addr: ListenerAddr,
351+
_socks_addr: SocketAddr,
352+
) -> Result<(), ClientError> {
353+
let firewall = Arc::clone(firewall);
354+
tokio::task::spawn_blocking(move || {
355+
// let firewall = Arc::clone(firewall);
356+
let local = UdpSocket::bind(l_addr.addr)?;
357+
local.set_nonblocking(false)?;
358+
firewall.setup_udp_socket(&local)?;
359+
std::thread::spawn(move || loop {
360+
let mut buf = vec![0u8; 1024];
361+
let ioslice = IoSliceMut::new(&mut buf);
362+
let mut cmsg_buffer = vec![0u8; 24];
363+
println!("{l_addr} UDP waiting");
364+
let cmsg: RecvMsg<nix::sys::socket::SockAddr> = recvmsg(
365+
local.as_raw_fd(),
366+
&mut [ioslice],
367+
Some(&mut cmsg_buffer),
368+
MsgFlags::empty(),
369+
)
370+
.unwrap();
371+
println!("{l_addr} UDP {cmsg:?}");
372+
});
373+
Ok(())
374+
})
375+
.await?
376+
}
377+
333378
async fn handle_tcp_client(
334379
socket: TcpStream,
335-
addr: SocketAddr,
380+
l_addr: &ListenerAddr,
336381
socks_addr: SocketAddr,
337382
firewall: Arc<dyn Firewall + Send + Sync>,
338383
) {
@@ -341,7 +386,7 @@ async fn handle_tcp_client(
341386
log::debug!("new connection from: {}", local_addr);
342387

343388
let remote_addr = firewall.get_dst_addr(&local).unwrap();
344-
log::info!("{addr} got connection from {local_addr} to {remote_addr}");
389+
log::info!("{l_addr} got connection from {local_addr} to {remote_addr}");
345390

346391
let (addr_str, port) = {
347392
let addr = remote_addr.ip().to_string();

‎src/firewall.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use std::{
22
error::Error,
33
fmt::Display,
4-
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
4+
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket},
55
os::unix::prelude::AsRawFd,
66
};
77

@@ -16,8 +16,8 @@ use tokio::net::{TcpListener, TcpStream};
1616

1717
use crate::{
1818
commands::Commands,
19-
network::SubnetsV4,
2019
network::{Family, SubnetsFamily, SubnetsV6},
20+
network::{ListenerAddr, SubnetsV4},
2121
};
2222

2323
pub mod nat;
@@ -81,6 +81,10 @@ pub trait Firewall {
8181
Ok(())
8282
}
8383

84+
fn setup_udp_socket(&self, _l: &UdpSocket) -> Result<(), FirewallError> {
85+
Ok(())
86+
}
87+
8488
fn get_dst_addr(&self, s: &TcpStream) -> Result<SocketAddr, FirewallError> {
8589
get_dst_addr_sockopt(s)
8690
}
@@ -91,7 +95,7 @@ pub trait Firewall {
9195

9296
pub struct FirewallSubnetConfig<T: SubnetsFamily> {
9397
pub enable: bool,
94-
pub port: u16,
98+
pub listener: ListenerAddr,
9599
pub includes: T,
96100
pub excludes: T,
97101
}

‎src/firewall/nat.rs

Lines changed: 40 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use crate::network::Ports;
2+
use crate::network::Protocol;
23
use crate::network::SubnetFamily;
34
use crate::network::SubnetsFamily;
45

@@ -18,7 +19,13 @@ impl NatFirewall {
1819
fconfig: &FirewallSubnetConfig<T>,
1920
commands: &mut Commands,
2021
) -> Result<(), FirewallError> {
21-
let port = fconfig.port.to_string();
22+
if !matches!(fconfig.listener.protocol, Protocol::Tcp) {
23+
return Err(FirewallError {
24+
message: "Only TCP is supported for NAT".to_string(),
25+
});
26+
}
27+
28+
let port = fconfig.listener.port().to_string();
2229
let chain = format!("sshuttle-{}", port);
2330
let family = fconfig.family();
2431

@@ -95,8 +102,14 @@ impl NatFirewall {
95102
fconfig: &FirewallSubnetConfig<T>,
96103
commands: &mut Commands,
97104
) -> Result<(), FirewallError> {
98-
let port = fconfig.port.to_string();
99-
let chain = format!("sshuttle-{}", fconfig.port);
105+
if !matches!(fconfig.listener.protocol, Protocol::Tcp) {
106+
return Err(FirewallError {
107+
message: "Only TCP is supported for NAT".to_string(),
108+
});
109+
}
110+
111+
let port = fconfig.listener.port().to_string();
112+
let chain = format!("sshuttle-{}", port);
100113
let family = fconfig.family();
101114

102115
macro_rules! ipt {
@@ -167,7 +180,7 @@ impl Firewall for NatFirewall {
167180
mod tests {
168181
use crate::{
169182
command::CommandLine,
170-
network::{SubnetsV4, SubnetsV6},
183+
network::{ListenerAddr, SubnetsV4, SubnetsV6},
171184
};
172185

173186
use super::*;
@@ -177,7 +190,10 @@ mod tests {
177190
let firewall = NatFirewall::new();
178191
let ipv4_family = FirewallSubnetConfig {
179192
enable: true,
180-
port: 1025,
193+
listener: ListenerAddr {
194+
protocol: Protocol::Tcp,
195+
addr: "127.0.0.1:1024".parse().unwrap(),
196+
},
181197
includes: "1.2.3.0/24:8000-9000".parse::<SubnetsV4>().unwrap(),
182198
excludes: "1.2.3.66:8080".parse::<SubnetsV4>().unwrap(),
183199
};
@@ -187,13 +203,13 @@ mod tests {
187203
};
188204

189205
let expected_ipv4: [&str; 7] = [
190-
"iptables -w -t nat -N sshuttle-1025",
191-
"iptables -w -t nat -F sshuttle-1025",
192-
"iptables -w -t nat -I OUTPUT 1 -j sshuttle-1025",
193-
"iptables -w -t nat -I PREROUTING 1 -j sshuttle-1025",
194-
"iptables -w -t nat -A sshuttle-1025 -j RETURN -m addrtype --dst-type LOCAL",
195-
"iptables -w -t nat -A sshuttle-1025 -j RETURN --dest 1.2.3.66/32 -p tcp --dport 8080",
196-
"iptables -w -t nat -A sshuttle-1025 -j REDIRECT --dest 1.2.3.0/24 -p tcp --dport 8000:9000 --to-ports 1025",
206+
"iptables -w -t nat -N sshuttle-1024",
207+
"iptables -w -t nat -F sshuttle-1024",
208+
"iptables -w -t nat -I OUTPUT 1 -j sshuttle-1024",
209+
"iptables -w -t nat -I PREROUTING 1 -j sshuttle-1024",
210+
"iptables -w -t nat -A sshuttle-1024 -j RETURN -m addrtype --dst-type LOCAL",
211+
"iptables -w -t nat -A sshuttle-1024 -j RETURN --dest 1.2.3.66/32 -p tcp --dport 8080",
212+
"iptables -w -t nat -A sshuttle-1024 -j REDIRECT --dest 1.2.3.0/24 -p tcp --dport 8000:9000 --to-ports 1024",
197213
];
198214

199215
let mut commands = Commands::default();
@@ -213,7 +229,10 @@ mod tests {
213229
let firewall = NatFirewall::new();
214230
let ipv6_family = FirewallSubnetConfig {
215231
enable: true,
216-
port: 1024,
232+
listener: ListenerAddr {
233+
protocol: Protocol::Tcp,
234+
addr: "127.0.0.1:1024".parse().unwrap(),
235+
},
217236
includes: "2404:6800:4004:80c::/64".parse::<SubnetsV6>().unwrap(),
218237
excludes: "[2404:6800:4004:80c::101f]:80".parse().unwrap(),
219238
};
@@ -248,7 +267,10 @@ mod tests {
248267
let firewall = NatFirewall::new();
249268
let ipv4_family = FirewallSubnetConfig {
250269
enable: true,
251-
port: 1024,
270+
listener: ListenerAddr {
271+
protocol: Protocol::Tcp,
272+
addr: "127.0.0.1:1024".parse().unwrap(),
273+
},
252274
includes: "1.2.3.0/32".parse::<SubnetsV4>().unwrap(),
253275
excludes: "1.2.3.66".parse::<SubnetsV4>().unwrap(),
254276
};
@@ -281,7 +303,10 @@ mod tests {
281303
let firewall = NatFirewall::new();
282304
let ipv6_family = FirewallSubnetConfig {
283305
enable: true,
284-
port: 1024,
306+
listener: ListenerAddr {
307+
protocol: Protocol::Tcp,
308+
addr: "127.0.0.1:1024".parse().unwrap(),
309+
},
285310
includes: "2404:6800:4004:80c::/64".parse::<SubnetsV6>().unwrap(),
286311
excludes: "[2404:6800:4004:80c::101f]:80".parse().unwrap(),
287312
};

0 commit comments

Comments
 (0)