mirror of
https://github.com/dndx/phantun.git
synced 2025-11-03 11:25:35 +08:00
chore(*) refactor Phantun and fake-tcp to their own crates
This commit is contained in:
Datong Sun
424
fake-tcp/src/lib.rs
Normal file
424
fake-tcp/src/lib.rs
Normal file
@@ -0,0 +1,424 @@
|
||||
pub mod packet;
|
||||
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use log::{info, trace};
|
||||
use packet::*;
|
||||
use pnet::packet::{tcp, Packet};
|
||||
use rand::prelude::*;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fmt;
|
||||
use std::net::{Ipv4Addr, SocketAddrV4};
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::sync::mpsc::{self, Receiver, Sender};
|
||||
use tokio::sync::watch;
|
||||
use tokio::sync::Mutex as AsyncMutex;
|
||||
use tokio::{io, time};
|
||||
use tokio_tun::Tun;
|
||||
|
||||
const TIMEOUT: time::Duration = time::Duration::from_secs(1);
|
||||
const RETRIES: usize = 6;
|
||||
const MPSC_BUFFER_LEN: usize = 512;
|
||||
|
||||
#[derive(Debug, Hash, Eq, PartialEq)]
|
||||
pub struct AddrTuple {
|
||||
local_addr: SocketAddrV4,
|
||||
remote_addr: SocketAddrV4,
|
||||
}
|
||||
|
||||
impl AddrTuple {
|
||||
fn new(local_addr: SocketAddrV4, remote_addr: SocketAddrV4) -> AddrTuple {
|
||||
AddrTuple {
|
||||
local_addr,
|
||||
remote_addr,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Shared {
|
||||
tuples: RwLock<HashMap<AddrTuple, Arc<Sender<Bytes>>>>,
|
||||
listening: RwLock<HashSet<u16>>,
|
||||
outgoing: Sender<Bytes>,
|
||||
ready: Sender<Socket>,
|
||||
}
|
||||
|
||||
pub struct Stack {
|
||||
shared: Arc<Shared>,
|
||||
local_ip: Ipv4Addr,
|
||||
ready: Receiver<Socket>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum State {
|
||||
Idle,
|
||||
SynSent,
|
||||
SynReceived,
|
||||
Established,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Mode {
|
||||
Client,
|
||||
Server,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Socket {
|
||||
mode: Mode,
|
||||
shared: Arc<Shared>,
|
||||
incoming: AsyncMutex<Receiver<Bytes>>,
|
||||
local_addr: SocketAddrV4,
|
||||
remote_addr: SocketAddrV4,
|
||||
seq: AtomicU32,
|
||||
ack: AtomicU32,
|
||||
state: State,
|
||||
closing_tx: watch::Sender<()>,
|
||||
closing_rx: watch::Receiver<()>,
|
||||
}
|
||||
|
||||
impl Socket {
|
||||
fn new(
|
||||
mode: Mode,
|
||||
shared: Arc<Shared>,
|
||||
local_addr: SocketAddrV4,
|
||||
remote_addr: SocketAddrV4,
|
||||
ack: Option<u32>,
|
||||
state: State,
|
||||
) -> (Socket, Sender<Bytes>) {
|
||||
let (incoming_tx, incoming_rx) = mpsc::channel(MPSC_BUFFER_LEN);
|
||||
let (closing_tx, closing_rx) = watch::channel(());
|
||||
|
||||
(
|
||||
Socket {
|
||||
mode,
|
||||
shared,
|
||||
incoming: AsyncMutex::new(incoming_rx),
|
||||
local_addr,
|
||||
remote_addr,
|
||||
seq: AtomicU32::new(0),
|
||||
ack: AtomicU32::new(ack.unwrap_or(0)),
|
||||
state,
|
||||
closing_tx,
|
||||
closing_rx,
|
||||
},
|
||||
incoming_tx,
|
||||
)
|
||||
}
|
||||
|
||||
fn build_tcp_packet(&self, flags: u16, payload: Option<&[u8]>) -> Bytes {
|
||||
build_tcp_packet(
|
||||
self.local_addr,
|
||||
self.remote_addr,
|
||||
self.seq.load(Ordering::Relaxed),
|
||||
self.ack.load(Ordering::Relaxed),
|
||||
flags,
|
||||
payload,
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn send(&self, payload: &[u8]) -> Option<()> {
|
||||
let mut closing = self.closing_rx.clone();
|
||||
|
||||
match self.state {
|
||||
State::Established => {
|
||||
let buf = self.build_tcp_packet(tcp::TcpFlags::ACK, Some(payload));
|
||||
self.seq.fetch_add(buf.len() as u32, Ordering::Relaxed);
|
||||
|
||||
tokio::select! {
|
||||
res = self.shared.outgoing.send(buf) => {
|
||||
res.unwrap();
|
||||
Some(())
|
||||
},
|
||||
_ = closing.changed() => {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn recv(&self, buf: &mut [u8]) -> Option<usize> {
|
||||
let mut closing = self.closing_rx.clone();
|
||||
|
||||
match self.state {
|
||||
State::Established => {
|
||||
let mut incoming = self.incoming.lock().await;
|
||||
tokio::select! {
|
||||
Some(raw_buf) = incoming.recv() => {
|
||||
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&raw_buf);
|
||||
|
||||
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
|
||||
info!("Connection {} reset by peer", self);
|
||||
self.close();
|
||||
return None;
|
||||
}
|
||||
|
||||
let payload = tcp_packet.payload();
|
||||
|
||||
self.ack
|
||||
.store(tcp_packet.get_sequence().wrapping_add(1), Ordering::Relaxed);
|
||||
|
||||
buf[..payload.len()].copy_from_slice(payload);
|
||||
|
||||
Some(payload.len())
|
||||
},
|
||||
_ = closing.changed() => {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn close(&self) {
|
||||
self.closing_tx.send(()).unwrap();
|
||||
}
|
||||
|
||||
async fn accept(mut self) {
|
||||
for _ in 0..RETRIES {
|
||||
match self.state {
|
||||
State::Idle => {
|
||||
let buf = self.build_tcp_packet(tcp::TcpFlags::SYN | tcp::TcpFlags::ACK, None);
|
||||
// ACK set by constructor
|
||||
self.shared.outgoing.send(buf).await.unwrap();
|
||||
self.state = State::SynReceived;
|
||||
info!("Sent SYN + ACK to client");
|
||||
}
|
||||
State::SynReceived => {
|
||||
let res = time::timeout(TIMEOUT, self.incoming.lock().await.recv()).await;
|
||||
if let Ok(buf) = res {
|
||||
let buf = buf.unwrap();
|
||||
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&buf);
|
||||
|
||||
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
if tcp_packet.get_flags() == tcp::TcpFlags::ACK
|
||||
&& tcp_packet.get_acknowledgement()
|
||||
== self.seq.load(Ordering::Relaxed) + 1
|
||||
{
|
||||
// found our ACK
|
||||
self.seq.fetch_add(1, Ordering::Relaxed);
|
||||
self.state = State::Established;
|
||||
|
||||
info!("Connection from {:?} established", self.remote_addr);
|
||||
let ready = self.shared.ready.clone();
|
||||
ready.send(self).await.unwrap();
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
info!("Waiting for client ACK timed out");
|
||||
self.state = State::Idle;
|
||||
}
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect(&mut self) -> Option<()> {
|
||||
for _ in 0..RETRIES {
|
||||
match self.state {
|
||||
State::Idle => {
|
||||
let buf = self.build_tcp_packet(tcp::TcpFlags::SYN, None);
|
||||
self.shared.outgoing.send(buf).await.unwrap();
|
||||
self.state = State::SynSent;
|
||||
info!("Sent SYN to server");
|
||||
}
|
||||
State::SynSent => {
|
||||
match time::timeout(TIMEOUT, self.incoming.lock().await.recv()).await {
|
||||
Ok(buf) => {
|
||||
let buf = buf.unwrap();
|
||||
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&buf);
|
||||
|
||||
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
if tcp_packet.get_flags() == tcp::TcpFlags::SYN | tcp::TcpFlags::ACK
|
||||
&& tcp_packet.get_acknowledgement()
|
||||
== self.seq.load(Ordering::Relaxed) + 1
|
||||
{
|
||||
// found our SYN + ACK
|
||||
self.seq.fetch_add(1, Ordering::Relaxed);
|
||||
self.ack
|
||||
.store(tcp_packet.get_sequence() + 1, Ordering::Relaxed);
|
||||
|
||||
// send ACK to finish handshake
|
||||
let buf = self.build_tcp_packet(tcp::TcpFlags::ACK, None);
|
||||
self.shared.outgoing.send(buf).await.unwrap();
|
||||
|
||||
self.state = State::Established;
|
||||
|
||||
info!("Connection to {:?} established", self.remote_addr);
|
||||
return Some(());
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
info!("Waiting for SYN + ACK timed out");
|
||||
self.state = State::Idle;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Socket {
|
||||
fn drop(&mut self) {
|
||||
// dissociates ourself from the dispatch map
|
||||
assert!(self
|
||||
.shared
|
||||
.tuples
|
||||
.write()
|
||||
.unwrap()
|
||||
.remove(&AddrTuple::new(self.local_addr, self.remote_addr))
|
||||
.is_some());
|
||||
|
||||
let buf = self.build_tcp_packet(tcp::TcpFlags::RST, None);
|
||||
self.shared.outgoing.try_send(buf).unwrap();
|
||||
self.close();
|
||||
info!("Fake TCP connection to {} closed", self);
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for Socket {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"(Fake TCP connection from {} to {})",
|
||||
self.local_addr, self.remote_addr
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Stack {
|
||||
pub fn new(tun: Tun) -> Stack {
|
||||
let (outgoing_tx, outgoing_rx) = mpsc::channel(MPSC_BUFFER_LEN);
|
||||
let (ready_tx, ready_rx) = mpsc::channel(MPSC_BUFFER_LEN);
|
||||
let shared = Arc::new(Shared {
|
||||
tuples: RwLock::new(HashMap::new()),
|
||||
outgoing: outgoing_tx,
|
||||
listening: RwLock::new(HashSet::new()),
|
||||
ready: ready_tx,
|
||||
});
|
||||
let local_ip = tun.destination().unwrap();
|
||||
|
||||
tokio::spawn(Stack::dispatch(tun, outgoing_rx, shared.clone()));
|
||||
Stack {
|
||||
shared,
|
||||
local_ip,
|
||||
ready: ready_rx,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn listen(&mut self, port: u16) {
|
||||
assert!(self.shared.listening.write().unwrap().insert(port));
|
||||
}
|
||||
|
||||
pub async fn accept(&mut self) -> Socket {
|
||||
self.ready.recv().await.unwrap()
|
||||
}
|
||||
|
||||
pub async fn connect(&mut self, addr: SocketAddrV4) -> 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 tuple = AddrTuple::new(local_addr, addr);
|
||||
let (mut sock, incoming) = Socket::new(
|
||||
Mode::Client,
|
||||
self.shared.clone(),
|
||||
local_addr,
|
||||
addr,
|
||||
None,
|
||||
State::Idle,
|
||||
);
|
||||
|
||||
{
|
||||
let mut tuples = self.shared.tuples.write().unwrap();
|
||||
assert!(tuples.insert(tuple, Arc::new(incoming.clone())).is_none());
|
||||
}
|
||||
|
||||
sock.connect().await.map(|_| sock)
|
||||
}
|
||||
|
||||
async fn dispatch(tun: Tun, mut outgoing: Receiver<Bytes>, shared: Arc<Shared>) {
|
||||
let (mut tun_r, mut tun_w) = io::split(tun);
|
||||
|
||||
loop {
|
||||
let mut buf = BytesMut::with_capacity(MAX_PACKET_LEN);
|
||||
|
||||
tokio::select! {
|
||||
buf = outgoing.recv() => {
|
||||
let buf = buf.unwrap();
|
||||
tun_w.write_all(&buf).await.unwrap();
|
||||
},
|
||||
s = tun_r.read_buf(&mut buf) => {
|
||||
s.unwrap();
|
||||
let buf = buf.freeze();
|
||||
if buf[0] >> 4 != 4 {
|
||||
// not an IPv4 packet
|
||||
continue;
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
let sender;
|
||||
{
|
||||
let tuples = shared.tuples.read().unwrap();
|
||||
sender = tuples.get(&tuple).cloned();
|
||||
}
|
||||
|
||||
if let Some(c) = sender {
|
||||
c.send(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(Mode::Server, shared.clone(), local_addr, remote_addr, Some(tcp_packet.get_sequence() + 1), State::Idle);
|
||||
assert!(shared.tuples.write().unwrap().insert(tuple, Arc::new(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() + 1,
|
||||
tcp::TcpFlags::RST,
|
||||
None,
|
||||
);
|
||||
shared.outgoing.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(),
|
||||
0,
|
||||
tcp::TcpFlags::RST,
|
||||
None,
|
||||
);
|
||||
shared.outgoing.try_send(buf).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
70
fake-tcp/src/packet.rs
Normal file
70
fake-tcp/src/packet.rs
Normal file
@@ -0,0 +1,70 @@
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use pnet::packet::{ip, ipv4, tcp};
|
||||
use std::convert::TryInto;
|
||||
use std::net::SocketAddrV4;
|
||||
|
||||
const IPV4_HEADER_LEN: usize = 20;
|
||||
const TCP_HEADER_LEN: usize = 20;
|
||||
pub const MAX_PACKET_LEN: usize = 1500;
|
||||
|
||||
pub fn build_tcp_packet(
|
||||
local_addr: SocketAddrV4,
|
||||
remote_addr: SocketAddrV4,
|
||||
seq: u32,
|
||||
ack: u32,
|
||||
flags: u16,
|
||||
payload: Option<&[u8]>,
|
||||
) -> Bytes {
|
||||
let wscale = (flags & tcp::TcpFlags::SYN) != 0;
|
||||
let tcp_total_len = TCP_HEADER_LEN + if wscale {4} else {0} // nop + wscale
|
||||
+ payload.map_or(0, |payload| payload.len());
|
||||
let total_len = IPV4_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 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(32);
|
||||
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);
|
||||
v4.set_checksum(ipv4::checksum(&v4.to_immutable()));
|
||||
|
||||
let mut tcp = tcp::MutableTcpPacket::new(&mut tcp_buf).unwrap();
|
||||
tcp.set_window(0xffff);
|
||||
tcp.set_source(local_addr.port());
|
||||
tcp.set_destination(remote_addr.port());
|
||||
tcp.set_sequence(seq);
|
||||
tcp.set_acknowledgement(ack);
|
||||
tcp.set_flags(flags);
|
||||
tcp.set_data_offset(TCP_HEADER_LEN as u8 / 4 + if wscale { 1 } else { 0 });
|
||||
if wscale {
|
||||
let wscale = tcp::TcpOption::wscale(14);
|
||||
tcp.set_options(&[tcp::TcpOption::nop(), wscale]);
|
||||
}
|
||||
|
||||
if let Some(payload) = payload {
|
||||
tcp.set_payload(payload);
|
||||
}
|
||||
|
||||
let checksum = tcp::ipv4_checksum(&tcp.to_immutable(), local_addr.ip(), remote_addr.ip());
|
||||
tcp.set_checksum(checksum);
|
||||
|
||||
v4_buf.unsplit(tcp_buf);
|
||||
|
||||
v4_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();
|
||||
|
||||
(v4, tcp)
|
||||
}
|
||||
Reference in New Issue
Block a user