anubis-wormhole 1.0.0

use crate::transit::hints::Hint;
use crate::transit::{relay_token_from_kdata, sender_handshake_string};
use futures::{stream::FuturesUnordered, StreamExt};
use std::pin::Pin;
use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::{TcpStream, TcpListener}, time::{timeout, Duration}};
use get_if_addrs::get_if_addrs;

pub async fn connect_direct(host: &str, port: u16) -> anyhow::Result<TcpStream> {
    let addr = format!("{}:{}", host, port);
    let s = TcpStream::connect(addr).await?;
    Ok(s)
}

pub async fn socks5_connect(socks_host: &str, socks_port: u16, target_host: &str, target_port: u16) -> anyhow::Result<TcpStream> {
    let mut s = TcpStream::connect(format!("{}:{}", socks_host, socks_port)).await?;
    // NO AUTH
    s.write_all(&[0x05, 0x01, 0x00]).await?;
    let mut m = [0u8;2]; s.read_exact(&mut m).await?; if m[1] != 0x00 { anyhow::bail!("SOCKS requires auth"); }
    let hostb = target_host.as_bytes(); if hostb.len() > 255 { anyhow::bail!("host too long"); }
    let mut req = Vec::with_capacity(7+hostb.len()); req.extend_from_slice(&[0x05, 0x01, 0x00, 0x03, hostb.len() as u8]); req.extend_from_slice(hostb); req.extend_from_slice(&target_port.to_be_bytes());
    s.write_all(&req).await?;
    let mut resp = [0u8; 4]; s.read_exact(&mut resp).await?; if resp[1] != 0x00 { anyhow::bail!("SOCKS connect failed"); }
    // read rest of reply addr
    let atyp = resp[3]; match atyp { 0x01 => { let mut b=[0u8;4]; s.read_exact(&mut b).await?; }, 0x03 => { let mut l=[0u8;1]; s.read_exact(&mut l).await?; let mut b=vec![0u8; l[0] as usize]; s.read_exact(&mut b).await?; }, 0x04 => { let mut b=[0u8;16]; s.read_exact(&mut b).await?; }, _=>{} }
    let mut prt=[0u8;2]; s.read_exact(&mut prt).await?;
    Ok(s)
}

pub async fn connect_relay(host: &str, port: u16, k_data: &[u8]) -> anyhow::Result<TcpStream> {
    let mut s = TcpStream::connect(format!("{}:{}", host, port)).await?;
    let token = relay_token_from_kdata(k_data);
    let hs = sender_handshake_string(&token);
    s.write_all(hs.as_bytes()).await?;
    let mut line = vec![0u8; 3];
    timeout(Duration::from_secs(5), s.read_exact(&mut line)).await??;
    if &line != b"ok\n" { anyhow::bail!("relay refused"); }
    Ok(s)
}

pub async fn race_connect(hints: Vec<Hint>, k_data: Option<&[u8]>, socks: Option<(&str,u16)>) -> anyhow::Result<TcpStream> {
    // one quick race, then a short backoff + a second try if nothing won
    let mut attempt = 0;
    loop {
        let mut futs: FuturesUnordered<Pin<Box<dyn std::future::Future<Output=anyhow::Result<TcpStream>> + Send>>> = FuturesUnordered::new();
        for h in hints.clone() {
            match h {
            Hint::DirectTcp { hostname, port, .. } => {
                let hn = hostname.clone();
                if let Some((sh,sp)) = socks {
                    futs.push(Box::pin(async move { socks5_connect(sh, sp, &hn, port).await }));
                } else {
                    futs.push(Box::pin(async move { connect_direct(&hn, port).await }));
                }
            }
            Hint::Relay { hostname, port, .. } => {
                if let Some(kd) = k_data {
                    if let Some((sh,sp)) = socks {
                        let hn = hostname.clone();
                        futs.push(Box::pin(async move {
                            let mut s = socks5_connect(sh, sp, &hn, port).await?;
                            let token = relay_token_from_kdata(kd);
                            let hs = sender_handshake_string(&token);
                            s.write_all(hs.as_bytes()).await?;
                            let mut line=[0u8;3]; timeout(Duration::from_secs(5), s.read_exact(&mut line)).await??;
                            if &line != b"ok\n" { anyhow::bail!("relay refused"); }
                            Ok(s)
                        }));
                    } else {
                        let hn = hostname.clone();
                        futs.push(Box::pin(async move { connect_relay(&hn, port, kd).await }));
                    }
                }
            }
            }
        }
        while let Some(res) = futs.next().await { if let Ok(s) = res { return Ok(s); } }
        attempt += 1;
        if attempt >= 2 { break; }
        tokio::time::sleep(Duration::from_millis(500)).await;
        // second attempt keeps same hints; a real implementation could reshuffle priorities
    }
    anyhow::bail!("no transit path available")
}

pub async fn start_direct_listener() -> anyhow::Result<(TcpListener, Vec<String>, u16)> {
    let listener = TcpListener::bind(("0.0.0.0", 0)).await?;
    let local_port = listener.local_addr()?.port();
    let mut addrs = vec![];
    for iface in get_if_addrs()? { if iface.ip().is_ipv4() { addrs.push(iface.ip().to_string()); } }
    Ok((listener, addrs, local_port))
}