snowstorm/net/endpoint/

mod.rs

/**
 * Snowstorm SDK - Net Endpoint
 *
 * © Snowstorm Inc 2025+
 */
use etherparse::{IcmpEchoHeader, Icmpv4Header, IpNumber};
use ipnet::{Ipv4Net, Ipv6Net};
#[allow(unused_imports)]
use ipstack::{
    stream::{IpStackStream, IpStackTcpStream, IpStackUdpStream, IpStackUnknownTransport},
    IpStackConfig,
};
use log::{debug, warn};
use shadowsocks::{
    config::ServerType, context::Context, crypto::CipherKind, net::TcpStream, ProxyClientStream,
    ServerConfig,
};
use std::{
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    time::Duration,
};

use tokio::sync::watch;
use tokio::{
    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
    select,
};
use tokio_util::{compat::FuturesAsyncReadCompatExt, sync::CancellationToken};

#[cfg(not(target_os = "windows"))]
use std::os::unix::io::RawFd;

#[cfg(not(target_os = "windows"))]
use tunio::{
    traits::{DriverT, InterfaceT},
    DefaultAsyncInterface, DefaultDriver,
};

#[cfg(feature = "snowstorm_client_udp")]
use {
    shadowsocks::{config::Mode, relay::Address, ProxySocket},
    tokio::io::AsyncReadExt,
};

const MTU: u16 = 1500;

#[cfg(not(target_os = "windows"))]
pub async fn new_client(
    raw_fd: RawFd,
    token: CancellationToken,
    sserver_port: watch::Receiver<u16>,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut tun_driver = DefaultDriver::new()?;
    let tun_config = DefaultAsyncInterface::config_builder()
        .name("utun321".to_string()) // OSX naming convention.
        .platform(|mut b| b.raw_fd(Some(raw_fd)).build())
        .unwrap()
        .build()?;

    // Create the tun device.
    debug!("tun");
    let tun_dev = DefaultAsyncInterface::new_up(&mut tun_driver, tun_config)?;

    return new_client_with_device(tun_dev.compat(), token, sserver_port).await;
}

pub async fn new_client_with_device<D>(
    device: D,
    token: CancellationToken,
    sserver_port: watch::Receiver<u16>,
) -> Result<(), Box<dyn std::error::Error>>
where
    D: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let ipstack_config = IpStackConfig {
        mtu: MTU,
        packet_information: cfg!(any(target_os = "macos", target_os = "ios")),
        tcp_timeout: Duration::from_secs(60),
        udp_timeout: Duration::from_secs(30),
    };

    // Start pushing packets back and forth.
    debug!("ipstack");
    let ip_stack = ipstack::IpStack::new(ipstack_config, device);

    // The sidecar curently hardcodes the addreses, we should
    // just use `tun_dev.handle()` to query the address instead.
    let if_net4: Ipv4Net = "192.18.0.1/24".parse().unwrap();
    let if_addr4 = if_net4.addr();
    let if_net6: Ipv6Net = "2001::1/64".parse().unwrap();
    let if_addr6 = if_net6.addr();

    debug!("task");
    tokio::spawn(async move {
        let _ = ipstack_worker(ip_stack, if_addr4, if_addr6, token, sserver_port).await;
    });

    Ok(())
}

async fn ipstack_worker(
    mut ip_stack: ipstack::IpStack,
    if_addr4: Ipv4Addr,
    if_addr6: Ipv6Addr,
    token: CancellationToken,
    sserver_port: watch::Receiver<u16>,
) -> Result<(), Box<dyn std::error::Error>> {
    loop {
        let stream = select! {
            _ = token.cancelled() => {
                debug!("endpoint: shutdown requested");
                break;
            }
            result = ip_stack.accept() => {
                match result {
                    Ok(stream) => {
                        stream
                    }
                    Err(e) => {
                        warn!("endpoint: ipstack accept failed: {}", e);
                        break;
                    }
                }
            }
        };

        let ss_ctx = Context::new_shared(ServerType::Local);
        let sserver_port = sserver_port.borrow().clone();
        let ss_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), sserver_port);
        #[allow(unused_mut)]
        let mut ss_cfg = ServerConfig::new(ss_addr, "", CipherKind::NONE)?;
        #[cfg(feature = "snowstorm_client_udp")]
        ss_cfg.set_mode(Mode::TcpAndUdp);

        match stream {
            IpStackStream::Tcp(tcp) => {
                let dst_addr = tcp.peer_addr();

                debug!("TCP: New connection to {:?}", dst_addr);

                let s = match ProxyClientStream::connect(ss_ctx.clone(), &ss_cfg, dst_addr).await {
                    Ok(s) => s,
                    Err(e) => {
                        warn!("TCP: failed to connect to {:?}: {:?}", dst_addr, e);
                        continue;
                    }
                };

                tokio::spawn(async move {
                    match tcp_stream_worker(tcp, s, dst_addr).await {
                        Ok(_) => {}
                        Err(e) => {
                            warn!("TCP: stream failed {:?}: {:?}", dst_addr, e);
                        }
                    };
                });
            }
            IpStackStream::Udp(udp) => {
                let dst_addr = udp.peer_addr();

                #[cfg(not(feature = "snowstorm_client_udp"))]
                #[cfg(not(target_os = "android"))]
                debug!("UDP: Ignoring packet to {:?}", dst_addr);

                // Android is not great at falling back from regular UDP-powered DNS, same goes for Windows
                // So for the time being we relay DoH to google's DNS servers.
                #[cfg(not(feature = "snowstorm_client_udp"))]
                #[cfg(any(target_os = "android", target_os = "windows"))]
                if dst_addr.port() == 53 {
                    debug!("Android DNS relay via DoH");
                    tokio::spawn(async move {
                        // Ignoring the error here, since most of the "streams" will inevitably timeout.
                        // We usually get one read per "stream".
                        let _ = udp_doh_relay_worker(udp).await;
                    });
                }

                #[cfg(feature = "snowstorm_client_udp")]
                {
                    debug!("UDP: New packet stream to {:?}", dst_addr);

                    let s = match ProxySocket::connect(ss_ctx.clone(), &ss_cfg).await {
                        Ok(s) => s,
                        Err(e) => {
                            warn!("UDP: failed to connect to {:?}: {:?}", dst_addr, e);
                            continue;
                        }
                    };

                    tokio::spawn(async move {
                        match udp_stream_worker(udp, s, dst_addr).await {
                            Ok(_) => {}
                            Err(e) => {
                                warn!("UDP: stream failed {:?}: {:?}", dst_addr, e);
                            }
                        };
                    });
                }
            }
            IpStackStream::UnknownTransport(u) => {
                handle_unknown_transport(u, if_addr4, if_addr6)?;
            }
            IpStackStream::UnknownNetwork(pkt) => {
                debug!("Unknown: Not even IP? - {} bytes", pkt.len());
            }
        };
    }

    // Stop the ip_stack worker.
    debug!("ipstack abort");
    ip_stack.handle.abort();
    debug!("ipstack aborted");

    // XXX: We may need to close the fd out from underneath
    // the tunio crate...

    Ok(())
}

async fn tcp_stream_worker(
    local_tcp: IpStackTcpStream,
    remote_tcp: ProxyClientStream<TcpStream>,
    dst_addr: SocketAddr,
) -> Result<(), Box<dyn std::error::Error>> {
    debug!("TCP: started worker: {:?}", dst_addr);

    let (mut l_rx, mut l_tx) = tokio::io::split(local_tcp);
    let (mut r_rx, mut r_tx) = tokio::io::split(remote_tcp);

    // Per crates/shadowsocks_service/src/local/utils.rs, there
    // are server-speaks first protocols like FTP.  The helper
    // to deal with this is naturally `pub(crate)`.
    let _ = r_tx.write(&[]).await;

    let _ = tokio::join!(
        async move {
            let r = tokio::io::copy(&mut l_rx, &mut r_tx).await;
            let _ = r_tx.shutdown().await;
            r
        },
        async move {
            let r = tokio::io::copy(&mut r_rx, &mut l_tx).await;
            let _ = l_tx.shutdown().await;
            r
        }
    );

    debug!("TCP: Finished connection to {:?}", dst_addr);

    Ok(())
}

#[cfg(not(feature = "snowstorm_client_udp"))]
#[cfg(any(target_os = "android", target_os = "windows"))]
async fn udp_doh_relay_worker(mut udp: IpStackUdpStream) -> Result<(), Box<dyn std::error::Error>> {
    use reqwest::Client;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let google_dns_ips = ["8.8.8.8:443", "8.8.4.4:443"].map(|ip| ip.parse().unwrap());

    let client = Client::builder()
        .resolve_to_addrs("dns.google", &google_dns_ips)
        .build()?;

    let mut buf = vec![0; 2048];
    loop {
        let len = udp.read(&mut buf).await?;
        if len == 0 {
            return Ok(());
        }

        let dns_query = &buf[..len];

        let response = client
            .post("https://dns.google/dns-query")
            .header("Content-Type", "application/dns-message")
            .body(dns_query.to_vec())
            .send()
            .await?;

        if response.status().is_success() {
            let dns_response = response.bytes().await?;
            udp.write_all(&dns_response).await?;
        } else {
            warn!("DoH: failed to relay DNS query: {:?}", response.status());
        }
    }
}

// TODO: In theory there is only a need for a single ProxySocket that is
// shared across all UDP packet streams, since the ShadowSocks protocol
// can handle multiplexing.  Doing so would be more efficient, but this
// is a hell of a lot simpler.
//
// This is heavily inspired by https://github.com/tun2proxy/tun2proxy/blob/master/src/lib.rs
#[cfg(feature = "snowstorm_client_udp")]
async fn udp_stream_worker(
    mut local_udp: IpStackUdpStream,
    remote_udp: ProxySocket,
    dst_addr: SocketAddr,
) -> Result<(), Box<dyn std::error::Error>> {
    const UDP_MSS: usize = 8192; // 64 KiB in theory.

    let mut l2s_buf = [0_u8; UDP_MSS as usize];
    let mut s2l_buf = [0_u8; UDP_MSS as usize];

    let ss_dst_addr = Address::from(dst_addr);

    debug!("UDP: started worker: {:?}", dst_addr);

    loop {
        tokio::select! {
            len = local_udp.read(&mut l2s_buf) => {
                let len = len?;
                if len == 0 {
                    break;
                }
                let pkt = &l2s_buf[..len];

                debug!("UDP: local->remote: {:?} {} bytes", dst_addr, len);

                remote_udp.send(&ss_dst_addr, pkt).await?;
            }
            res = remote_udp.recv(&mut s2l_buf) => {
                let len = res?.0;
                if len == 0 {
                    break;
                }
                let pkt = &s2l_buf[..len];

                debug!("UDP: remote->local: {:?} {} bytes", dst_addr, len);

                local_udp.write_all(pkt).await?;
            }
        }
    }

    debug!("UDP: Finished connection to {:?}", dst_addr);

    Ok(())
}

fn handle_unknown_transport(
    u: IpStackUnknownTransport,
    if_addr4: Ipv4Addr,
    _if_addr6: Ipv6Addr,
) -> Result<(), Box<dyn std::error::Error>> {
    let dst_addr = u.dst_addr();

    if u.ip_protocol() != IpNumber::ICMP {
        debug!("IP: {:?} {:?}", dst_addr, u.ip_protocol());
        return Ok(());
    }

    // TODO: This could do IPv6, but the response header checksum
    // covers the src and destination addresses, and it's not obvious
    // to me how to convert from the rust Ipv6Addr type to a `[u8; 16]`.
    if !u.src_addr().is_ipv4() {
        debug!("ICMP: {:?} {:?}", dst_addr, u.ip_protocol());
        return Ok(());
    }

    let (icmp_header, req_payload) = Icmpv4Header::from_slice(u.payload())?;
    if let etherparse::Icmpv4Type::EchoRequest(req) = icmp_header.icmp_type {
        debug!("ICMPv4: echo {:?}", dst_addr);

        if dst_addr != if_addr4 {
            return Ok(());
        }

        let echo = IcmpEchoHeader {
            id: req.id,
            seq: req.seq,
        };
        let mut resp = Icmpv4Header::new(etherparse::Icmpv4Type::EchoReply(echo));
        resp.update_checksum(req_payload);
        let mut payload = resp.to_bytes().to_vec();
        payload.extend_from_slice(req_payload);
        u.send(payload)?;
    } else {
        debug!("ICMPv4: {:?} {:?}", icmp_header.icmp_type, dst_addr);
    }

    Ok(())
}