feat(*) IPv6 support

This commit is contained in:
Datong Sun
2022-04-16 01:41:19 -07:00
parent 74183071f1
commit 85555f2a34
6 changed files with 351 additions and 139 deletions

View File

@@ -49,7 +49,7 @@ use pnet::packet::{tcp, Packet};
use rand::prelude::*;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::net::{Ipv4Addr, SocketAddrV4};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, RwLock};
use tokio::sync::broadcast;
@@ -65,12 +65,12 @@ const MAX_UNACKED_LEN: u32 = 128 * 1024 * 1024; // 128MB
#[derive(Hash, Eq, PartialEq, Clone, Debug)]
struct AddrTuple {
local_addr: SocketAddrV4,
remote_addr: SocketAddrV4,
local_addr: SocketAddr,
remote_addr: SocketAddr,
}
impl AddrTuple {
fn new(local_addr: SocketAddrV4, remote_addr: SocketAddrV4) -> AddrTuple {
fn new(local_addr: SocketAddr, remote_addr: SocketAddr) -> AddrTuple {
AddrTuple {
local_addr,
remote_addr,
@@ -89,6 +89,7 @@ struct Shared {
pub struct Stack {
shared: Arc<Shared>,
local_ip: Ipv4Addr,
local_ip6: Option<Ipv6Addr>,
ready: mpsc::Receiver<Socket>,
}
@@ -103,8 +104,8 @@ pub struct Socket {
shared: Arc<Shared>,
tun: Arc<Tun>,
incoming: flume::Receiver<Bytes>,
local_addr: SocketAddrV4,
remote_addr: SocketAddrV4,
local_addr: SocketAddr,
remote_addr: SocketAddr,
seq: AtomicU32,
ack: AtomicU32,
last_ack: AtomicU32,
@@ -122,8 +123,8 @@ impl Socket {
fn new(
shared: Arc<Shared>,
tun: Arc<Tun>,
local_addr: SocketAddrV4,
remote_addr: SocketAddrV4,
local_addr: SocketAddr,
remote_addr: SocketAddr,
ack: Option<u32>,
state: State,
) -> (Socket, flume::Sender<Bytes>) {
@@ -188,7 +189,7 @@ impl Socket {
match self.state {
State::Established => {
self.incoming.recv_async().await.ok().and_then(|raw_buf| {
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&raw_buf);
let (_v4_packet, tcp_packet) = parse_ip_packet(&raw_buf).unwrap();
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
info!("Connection {} reset by peer", self);
@@ -233,7 +234,7 @@ impl Socket {
let res = time::timeout(TIMEOUT, self.incoming.recv_async()).await;
if let Ok(buf) = res {
let buf = buf.unwrap();
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&buf);
let (_v4_packet, tcp_packet) = parse_ip_packet(&buf).unwrap();
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
return;
@@ -277,7 +278,7 @@ impl Socket {
match time::timeout(TIMEOUT, self.incoming.recv_async()).await {
Ok(buf) => {
let buf = buf.unwrap();
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&buf);
let (_v4_packet, tcp_packet) = parse_ip_packet(&buf).unwrap();
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
return None;
@@ -358,7 +359,7 @@ impl Stack {
/// When more than one [`Tun`](tokio_tun::Tun) object is passed in, same amount
/// of reader will be spawned later. This allows user to utilize the performance
/// benefit of Multiqueue Tun support on machines with SMP.
pub fn new(tun: Vec<Tun>) -> Stack {
pub fn new(tun: Vec<Tun>, local_ip: Ipv4Addr, local_ip6: Option<Ipv6Addr>) -> Stack {
let tun: Vec<Arc<Tun>> = tun.into_iter().map(Arc::new).collect();
let (ready_tx, ready_rx) = mpsc::channel(MPSC_BUFFER_LEN);
let (tuples_purge_tx, _tuples_purge_rx) = broadcast::channel(16);
@@ -369,7 +370,6 @@ impl Stack {
ready: ready_tx,
tuples_purge: tuples_purge_tx.clone(),
});
let local_ip = tun[0].destination().unwrap();
for t in tun {
tokio::spawn(Stack::reader_task(
@@ -382,6 +382,7 @@ impl Stack {
Stack {
shared,
local_ip,
local_ip6,
ready: ready_rx,
}
}
@@ -398,10 +399,17 @@ impl Stack {
/// Connects to the remote end. `None` returned means
/// the connection attempt failed.
pub async fn connect(&mut self, addr: SocketAddrV4) -> Option<Socket> {
pub async fn connect(&mut self, addr: SocketAddr) -> Option<Socket> {
let mut rng = SmallRng::from_entropy();
let local_port: u16 = rng.gen_range(1024..65535);
let local_addr = SocketAddrV4::new(self.local_ip, local_port);
let local_addr = SocketAddr::new(
if addr.is_ipv4() {
IpAddr::V4(self.local_ip)
} else {
IpAddr::V6(self.local_ip6.expect("IPv6 local address undefined"))
},
local_port,
);
let tuple = AddrTuple::new(local_addr, addr);
let (mut sock, incoming) = Socket::new(
self.shared.clone(),
@@ -437,90 +445,90 @@ impl Stack {
buf.truncate(size);
let buf = buf.freeze();
if buf[0] >> 4 != 4 {
// not an IPv4 packet
continue;
}
match parse_ip_packet(&buf) {
Some((ip_packet, tcp_packet)) => {
let local_addr =
SocketAddr::new(ip_packet.get_destination(), tcp_packet.get_destination());
let remote_addr = SocketAddr::new(ip_packet.get_source(), tcp_packet.get_source());
let (ip_packet, tcp_packet) = parse_ipv4_packet(&buf);
let local_addr =
SocketAddrV4::new(ip_packet.get_destination(), tcp_packet.get_destination());
let remote_addr = SocketAddrV4::new(ip_packet.get_source(), tcp_packet.get_source());
let tuple = AddrTuple::new(local_addr, remote_addr);
if let Some(c) = tuples.get(&tuple) {
if c.send_async(buf).await.is_err() {
trace!("Cache hit, but receiver already closed, dropping packet");
}
let tuple = AddrTuple::new(local_addr, remote_addr);
if let Some(c) = tuples.get(&tuple) {
if c.send_async(buf).await.is_err() {
trace!("Cache hit, but receiver already closed, dropping packet");
continue;
// If not Ok, receiver has been closed and just fall through to the slow
// path below
} else {
trace!("Cache miss, checking the shared tuples table for connection");
let sender = {
let tuples = shared.tuples.read().unwrap();
tuples.get(&tuple).cloned()
};
if let Some(c) = sender {
trace!("Storing connection information into local tuples");
tuples.insert(tuple, c.clone());
c.send_async(buf).await.unwrap();
continue;
}
}
if tcp_packet.get_flags() == tcp::TcpFlags::SYN
&& shared
.listening
.read()
.unwrap()
.contains(&tcp_packet.get_destination())
{
// SYN seen on listening socket
if tcp_packet.get_sequence() == 0 {
let (sock, incoming) = Socket::new(
shared.clone(),
tun.clone(),
local_addr,
remote_addr,
Some(tcp_packet.get_sequence() + 1),
State::Idle,
);
assert!(shared
.tuples
.write()
.unwrap()
.insert(tuple, incoming)
.is_none());
tokio::spawn(sock.accept());
} else {
trace!("Bad TCP SYN packet from {}, sending RST", remote_addr);
let buf = build_tcp_packet(
local_addr,
remote_addr,
0,
tcp_packet.get_sequence() + tcp_packet.payload().len() as u32 + 1, // +1 because of SYN flag set
tcp::TcpFlags::RST | tcp::TcpFlags::ACK,
None,
);
shared.tun[0].try_send(&buf).unwrap();
}
} else if (tcp_packet.get_flags() & tcp::TcpFlags::RST) == 0 {
info!("Unknown TCP packet from {}, sending RST", remote_addr);
let buf = build_tcp_packet(
local_addr,
remote_addr,
tcp_packet.get_acknowledgement(),
tcp_packet.get_sequence() + tcp_packet.payload().len() as u32,
tcp::TcpFlags::RST | tcp::TcpFlags::ACK,
None,
);
shared.tun[0].try_send(&buf).unwrap();
}
}
continue;
// If not Ok, receiver has been closed and just fall through to the slow
// path below
} else {
trace!("Cache miss, checking the shared tuples table for connection");
let sender = {
let tuples = shared.tuples.read().unwrap();
tuples.get(&tuple).cloned()
};
if let Some(c) = sender {
trace!("Storing connection information into local tuples");
tuples.insert(tuple, c.clone());
c.send_async(buf).await.unwrap();
None => {
continue;
}
}
if tcp_packet.get_flags() == tcp::TcpFlags::SYN
&& shared
.listening
.read()
.unwrap()
.contains(&tcp_packet.get_destination())
{
// SYN seen on listening socket
if tcp_packet.get_sequence() == 0 {
let (sock, incoming) = Socket::new(
shared.clone(),
tun.clone(),
local_addr,
remote_addr,
Some(tcp_packet.get_sequence() + 1),
State::Idle,
);
assert!(shared
.tuples
.write()
.unwrap()
.insert(tuple, incoming)
.is_none());
tokio::spawn(sock.accept());
} else {
trace!("Bad TCP SYN packet from {}, sending RST", remote_addr);
let buf = build_tcp_packet(
local_addr,
remote_addr,
0,
tcp_packet.get_sequence() + tcp_packet.payload().len() as u32 + 1, // +1 because of SYN flag set
tcp::TcpFlags::RST | tcp::TcpFlags::ACK,
None,
);
shared.tun[0].try_send(&buf).unwrap();
}
} else if (tcp_packet.get_flags() & tcp::TcpFlags::RST) == 0 {
info!("Unknown TCP packet from {}, sending RST", remote_addr);
let buf = build_tcp_packet(
local_addr,
remote_addr,
tcp_packet.get_acknowledgement(),
tcp_packet.get_sequence() + tcp_packet.payload().len() as u32,
tcp::TcpFlags::RST | tcp::TcpFlags::ACK,
None,
);
shared.tun[0].try_send(&buf).unwrap();
}
},
tuple = tuples_purge.recv() => {
let tuple = tuple.unwrap();

View File

@@ -1,45 +1,85 @@
use bytes::{Bytes, BytesMut};
use internet_checksum::Checksum;
use pnet::packet::Packet;
use pnet::packet::{ip, ipv4, tcp};
use pnet::packet::{ip, ipv4, ipv6, tcp};
use std::convert::TryInto;
use std::net::SocketAddrV4;
use std::net::{IpAddr, SocketAddr};
const IPV4_HEADER_LEN: usize = 20;
const IPV6_HEADER_LEN: usize = 40;
const TCP_HEADER_LEN: usize = 20;
pub const MAX_PACKET_LEN: usize = 1500;
pub enum IPPacket<'p> {
V4(ipv4::Ipv4Packet<'p>),
V6(ipv6::Ipv6Packet<'p>),
}
impl<'a> IPPacket<'a> {
pub fn get_source(&self) -> IpAddr {
match self {
IPPacket::V4(p) => IpAddr::V4(p.get_source()),
IPPacket::V6(p) => IpAddr::V6(p.get_source()),
}
}
pub fn get_destination(&self) -> IpAddr {
match self {
IPPacket::V4(p) => IpAddr::V4(p.get_destination()),
IPPacket::V6(p) => IpAddr::V6(p.get_destination()),
}
}
}
pub fn build_tcp_packet(
local_addr: SocketAddrV4,
remote_addr: SocketAddrV4,
local_addr: SocketAddr,
remote_addr: SocketAddr,
seq: u32,
ack: u32,
flags: u16,
payload: Option<&[u8]>,
) -> Bytes {
let ip_header_len = match local_addr {
SocketAddr::V4(_) => IPV4_HEADER_LEN,
SocketAddr::V6(_) => IPV6_HEADER_LEN,
};
let wscale = (flags & tcp::TcpFlags::SYN) != 0;
let tcp_header_len = TCP_HEADER_LEN + if wscale { 4 } else { 0 }; // nop + wscale
let tcp_total_len = tcp_header_len + payload.map_or(0, |payload| payload.len());
let total_len = IPV4_HEADER_LEN + tcp_total_len;
let total_len = ip_header_len + tcp_total_len;
let mut buf = BytesMut::with_capacity(total_len);
buf.resize(total_len, 0);
let mut v4_buf = buf.split_to(IPV4_HEADER_LEN);
let mut ip_buf = buf.split_to(ip_header_len);
let mut tcp_buf = buf.split_to(tcp_total_len);
assert_eq!(0, buf.len());
let mut v4 = ipv4::MutableIpv4Packet::new(&mut v4_buf).unwrap();
v4.set_version(4);
v4.set_header_length(IPV4_HEADER_LEN as u8 / 4);
v4.set_next_level_protocol(ip::IpNextHeaderProtocols::Tcp);
v4.set_ttl(64);
v4.set_source(*local_addr.ip());
v4.set_destination(*remote_addr.ip());
v4.set_total_length(total_len.try_into().unwrap());
v4.set_flags(ipv4::Ipv4Flags::DontFragment);
let mut cksm = Checksum::new();
cksm.add_bytes(v4.packet());
v4.set_checksum(u16::from_be_bytes(cksm.checksum()));
match (local_addr, remote_addr) {
(SocketAddr::V4(local), SocketAddr::V4(remote)) => {
let mut v4 = ipv4::MutableIpv4Packet::new(&mut ip_buf).unwrap();
v4.set_version(4);
v4.set_header_length(IPV4_HEADER_LEN as u8 / 4);
v4.set_next_level_protocol(ip::IpNextHeaderProtocols::Tcp);
v4.set_ttl(64);
v4.set_source(*local.ip());
v4.set_destination(*remote.ip());
v4.set_total_length(total_len.try_into().unwrap());
v4.set_flags(ipv4::Ipv4Flags::DontFragment);
let mut cksm = Checksum::new();
cksm.add_bytes(v4.packet());
v4.set_checksum(u16::from_be_bytes(cksm.checksum()));
}
(SocketAddr::V6(local), SocketAddr::V6(remote)) => {
let mut v6 = ipv6::MutableIpv6Packet::new(&mut ip_buf).unwrap();
v6.set_version(6);
v6.set_payload_length(tcp_total_len.try_into().unwrap());
v6.set_next_header(ip::IpNextHeaderProtocols::Tcp);
v6.set_hop_limit(64);
v6.set_source(*local.ip());
v6.set_destination(*remote.ip());
}
_ => unreachable!(),
};
let mut tcp = tcp::MutableTcpPacket::new(&mut tcp_buf).unwrap();
tcp.set_window(0xffff);
@@ -59,24 +99,55 @@ pub fn build_tcp_packet(
}
let mut cksm = Checksum::new();
cksm.add_bytes(&local_addr.ip().octets());
cksm.add_bytes(&remote_addr.ip().octets());
let ip::IpNextHeaderProtocol(tcp_protocol) = ip::IpNextHeaderProtocols::Tcp;
let mut pseudo = [0u8, tcp_protocol, 0, 0];
pseudo[2..].copy_from_slice(&(tcp_total_len as u16).to_be_bytes());
cksm.add_bytes(&pseudo);
match (local_addr, remote_addr) {
(SocketAddr::V4(local), SocketAddr::V4(remote)) => {
cksm.add_bytes(&local.ip().octets());
cksm.add_bytes(&remote.ip().octets());
let mut pseudo = [0u8, tcp_protocol, 0, 0];
pseudo[2..].copy_from_slice(&(tcp_total_len as u16).to_be_bytes());
cksm.add_bytes(&pseudo);
}
(SocketAddr::V6(local), SocketAddr::V6(remote)) => {
cksm.add_bytes(&local.ip().octets());
cksm.add_bytes(&remote.ip().octets());
let mut pseudo = [0u8, 0, 0, 0, 0, 0, 0, tcp_protocol];
pseudo[0..4].copy_from_slice(&(tcp_total_len as u32).to_be_bytes());
cksm.add_bytes(&pseudo);
}
_ => unreachable!(),
};
cksm.add_bytes(tcp.packet());
tcp.set_checksum(u16::from_be_bytes(cksm.checksum()));
v4_buf.unsplit(tcp_buf);
v4_buf.freeze()
ip_buf.unsplit(tcp_buf);
ip_buf.freeze()
}
pub fn parse_ipv4_packet(buf: &Bytes) -> (ipv4::Ipv4Packet, tcp::TcpPacket) {
let v4 = ipv4::Ipv4Packet::new(buf).unwrap();
let tcp = tcp::TcpPacket::new(&buf[IPV4_HEADER_LEN..]).unwrap();
pub fn parse_ip_packet(buf: &Bytes) -> Option<(IPPacket, tcp::TcpPacket)> {
if buf[0] >> 4 == 4 {
let v4 = ipv4::Ipv4Packet::new(buf).unwrap();
if v4.get_next_level_protocol() != ip::IpNextHeaderProtocols::Tcp {
return None;
}
(v4, tcp)
let tcp = tcp::TcpPacket::new(&buf[IPV4_HEADER_LEN..]).unwrap();
Some((IPPacket::V4(v4), tcp))
} else if buf[0] >> 4 == 6 {
let v6 = ipv6::Ipv6Packet::new(buf).unwrap();
if v6.get_next_header() != ip::IpNextHeaderProtocols::Tcp {
return None;
}
let tcp = tcp::TcpPacket::new(&buf[IPV6_HEADER_LEN..]).unwrap();
Some((IPPacket::V6(v6), tcp))
} else {
None
}
}
#[cfg(all(test, feature = "benchmark"))]