//! atlibs-webrtc: `webrtc-peer --tunnel <tag>=<port> ... [--origin <page origin>] [--port 0] [--host 127.0.0.1]`
//! (src/main.rs, its own binary: `cargo build --release`; not in the atlibs CLI).
//! It needs only atlibs-host's Signal.

// The WebRTC guest side on this machine, for local checks (tests/loopback.test.mjs):
// what socialweb-guest-agent (pion) does on a won VM, with the bidder's records replaced by a local
// HTTP signal endpoint. One age X25519 identity and one DTLS certificate (so one fingerprint for every
// tag); each sealed offer gets its own peer (str0m, ice-lite, one host candidate on --host) whose
// datachannels bridge, by label, to 127.0.0.1:<port> of that tag's tunnel.
//
//   GET  /events   {events: [{kind, value}]}, the shape the requester `watch` entry returns:
//                  guest.publicKey {kind: age-x25519, tag, key} and vm.onNetwork {kind: webrtc, tag,
//                  address: "sha-256 AB:.."} per tunnel, then guest.signalWebRTC {kind: answer, tag,
//                  ciphertext} per answered offer
//   POST /signal   the requester `signal` input {receipt, tag, offer}: opens the armored offer
//                  ({sdp, replyTo}), answers, seals {sdp} to replyTo
//
// Prints {signal, key, fingerprint} as one JSON line on stdout once listening. Loopback only: the
// endpoint is unauthenticated, and the tunnels are as open as the guest agent's. So no other site's page
// gets in: a Host other than the listener's is refused (DNS rebinding), and an Origin other than --origin
// (the loopback page's, the only one CORS allows); at most MAX_ANSWERS offers are answered per run.

use std::cell::RefCell;
use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};
use std::rc::Rc;
use std::sync::Arc;
use std::time::Instant;

use atlibs_host::Signal;
use base64::Engine;
use base64::engine::general_purpose::{STANDARD, STANDARD_NO_PAD};
use http_body_util::{BodyExt, Full};
use hyper::body::{Bytes, Incoming};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper_util::rt::TokioIo;
use serde_json::{Value, json};
use str0m::change::SdpOffer;
use str0m::channel::ChannelId;
use str0m::crypto::{CryptoProvider, from_feature_flags};
use str0m::net::{Protocol, Receive};
use str0m::{Candidate, Event, IceConnectionState, Input, Output, Rtc, RtcConfig};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc;


const MAX_OFFER: usize = 16384; // the requester's limit on a sealed offer
const MAX_ANSWERS: usize = 64; // each binds a UDP socket and adds an event: a test needs a handful

struct Peer {
    provider: Arc<CryptoProvider>,
    cert: str0m::config::DtlsCert,
    host: IpAddr,
    identity: [u8; 32],
    tunnels: HashMap<String, u16>,
    events: RefCell<Vec<Value>>,
}

/// Why a request is refused (403): a Host that isn't the listener's (`bound`, or localhost on its port), or
/// an Origin other than `origin` (a page of another site: a text/plain POST needs no preflight). Missing or
/// non-ASCII headers fail closed; no Origin (curl, Node's fetch) is not a cross-site page.
fn refused(headers: &http::HeaderMap, bound: SocketAddr, origin: Option<&str>) -> Option<&'static str> {
    let host = headers.get("host").and_then(|v| v.to_str().ok());
    if !host.is_some_and(|h| h.eq_ignore_ascii_case(&bound.to_string()) || h.eq_ignore_ascii_case(&format!("localhost:{}", bound.port()))) {
        return Some("Host is not this listener's");
    }
    match headers.get("origin").map(|v| v.to_str()) {
        None => None,
        Some(Ok(o)) if Some(o) == origin => None,
        Some(_) => Some("cross-origin requests are refused (webrtc-peer --origin <the page's origin>)"),
    }
}

/// The peer on host:port (0: any) with its tunnels (tag -> 127.0.0.1 port), until signal aborts. origin:
/// the page allowed to call it (CORS); None: no page, only clients that send no Origin.
pub async fn main(tunnels: HashMap<String, u16>, host: IpAddr, port: u16, origin: Option<String>, signal: Signal) -> Result<(), String> {
    let provider = Arc::new(from_feature_flags());
    let cert = provider.dtls_provider.generate_certificate().ok_or("no DTLS certificate")?;
    let fingerprint = RtcConfig::new().set_crypto_provider(provider.clone()).set_dtls_cert(cert.clone()).build(Instant::now()).direct_api().local_dtls_fingerprint().to_string();
    let identity: [u8; 32] = random();
    let key = age::recipient(&age::public(&identity));
    let mut events = vec![];
    for tag in tunnels.keys() {
        events.push(json!({ "kind": "guest.publicKey", "value": { "kind": "age-x25519", "tag": tag, "key": key } }));
        events.push(json!({ "kind": "vm.onNetwork", "value": { "kind": "webrtc", "tag": tag, "address": fingerprint } }));
    }
    let peer = Rc::new(Peer { provider, cert, host, identity, tunnels, events: RefCell::new(events) });
    let addr = SocketAddr::new(host, port);
    let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| format!("{addr}: {e}"))?;
    let bound = listener.local_addr().map_err(|e| e.to_string())?;
    println!("{}", json!({ "signal": format!("http://{bound}/"), "key": key, "fingerprint": fingerprint }));
    let local = tokio::task::LocalSet::new();
    local
        .run_until(async move {
            loop {
                let (stream, _) = tokio::select! {
                    accepted = listener.accept() => accepted.map_err(|e| e.to_string())?,
                    _ = signal.aborted() => return Ok(()),
                };
                let (peer, origin) = (peer.clone(), origin.clone());
                tokio::task::spawn_local(async move {
                    let service = service_fn(move |req: http::Request<Incoming>| {
                        let (peer, origin) = (peer.clone(), origin.clone());
                        async move {
                            let (parts, body) = req.into_parts();
                            let bytes = body.collect().await.map(|b| b.to_bytes()).unwrap_or_default();
                            let (status, body) = match refused(&parts.headers, bound, origin.as_deref()) {
                                Some(why) => (403, json!({ "error": "Forbidden", "message": why })),
                                None => match respond(&peer, &parts.method, parts.uri.path(), &bytes).await {
                                    Ok(v) => (200, v),
                                    Err(e) => (400, json!({ "error": "InvalidRequest", "message": e })),
                                },
                            };
                            let mut res = http::Response::builder().status(status).header("content-type", "application/json");
                            if let Some(origin) = &origin {
                                res = res.header("access-control-allow-origin", origin.as_str()); // the page is on another loopback port
                            }
                            let res = res.body(Full::new(Bytes::from(body.to_string()))).expect("valid response");
                            Ok::<_, std::convert::Infallible>(res)
                        }
                    });
                    let _ = http1::Builder::new().serve_connection(TokioIo::new(stream), service).await;
                });
            }
        })
        .await
}

async fn respond(peer: &Rc<Peer>, method: &http::Method, path: &str, body: &[u8]) -> Result<Value, String> {
    match (method.as_str(), path) {
        ("GET", "/events") => Ok(json!({ "events": *peer.events.borrow() })),
        ("POST", "/signal") => {
            let input: Value = serde_json::from_slice(body).map_err(|e| format!("signal input: {e}"))?;
            let tag = input["tag"].as_str().ok_or("tag is required")?.to_string();
            if peer.events.borrow().len() >= peer.tunnels.len() * 2 + MAX_ANSWERS {
                return Err(format!("{MAX_ANSWERS} offers answered: restart webrtc-peer"));
            }
            let offer = input["offer"].as_str().filter(|o| o.len() <= MAX_OFFER).ok_or("offer must be an age armored string of at most 16384 bytes")?;
            let opened: Value = serde_json::from_slice(&age::open(&peer.identity, offer)?).map_err(|e| format!("offer: {e}"))?;
            let (sdp, reply_to) = (opened["sdp"].as_str().ok_or("offer has no sdp")?, opened["replyTo"].as_str().ok_or("offer has no replyTo")?);
            let reply_to = age::parse_recipient(reply_to)?;
            let answer = answer(peer, sdp).await?;
            let ciphertext = age::seal(&reply_to, json!({ "sdp": answer }).to_string().as_bytes());
            peer.events.borrow_mut().push(json!({ "kind": "guest.signalWebRTC", "value": { "kind": "answer", "tag": tag, "ciphertext": ciphertext } }));
            Ok(json!({}))
        }
        _ => Err(format!("{method} {path}: try GET /events or POST /signal")),
    }
}

// A peer for one offer: its UDP socket, one host candidate, ice-lite (the browser does the checks).
async fn answer(peer: &Rc<Peer>, sdp: &str) -> Result<String, String> {
    let offer = SdpOffer::from_sdp_string(sdp).map_err(|e| format!("offer sdp: {e}"))?;
    let socket = tokio::net::UdpSocket::bind(SocketAddr::new(peer.host, 0)).await.map_err(|e| e.to_string())?;
    let local = socket.local_addr().map_err(|e| e.to_string())?;
    let mut rtc = RtcConfig::new().set_crypto_provider(peer.provider.clone()).set_dtls_cert(peer.cert.clone()).set_ice_lite(true).build(Instant::now());
    rtc.add_local_candidate(Candidate::host(local, "udp").map_err(|e| e.to_string())?);
    let answer = rtc.sdp_api().accept_offer(offer).map_err(|e| format!("accept offer: {e}"))?;
    tokio::task::spawn_local(run(rtc, socket, local, peer.tunnels.clone()));
    Ok(answer.to_sdp_string())
}

enum Woke {
    Udp(usize, SocketAddr),
    Tcp(ChannelId, Option<Vec<u8>>),
    Timeout,
}

// The sans-IO loop (str0m's examples/http-post.rs on tokio): UDP in, transmits out, and each
// datachannel bridged to its tag's TCP port by a task (bytes up via `up`, down via its sender).
async fn run(mut rtc: Rtc, socket: tokio::net::UdpSocket, local: SocketAddr, tunnels: HashMap<String, u16>) {
    let mut buf = vec![0u8; 2000];
    let (up, mut from_tcp) = mpsc::unbounded_channel::<(ChannelId, Option<Vec<u8>>)>();
    let mut to_tcp: HashMap<ChannelId, mpsc::UnboundedSender<Vec<u8>>> = HashMap::new();
    loop {
        let deadline = match rtc.poll_output() {
            Ok(Output::Timeout(t)) => t,
            Ok(Output::Transmit(t)) => {
                let _ = socket.send_to(&t.contents, t.destination).await;
                continue;
            }
            Ok(Output::Event(event)) => {
                match event {
                    Event::IceConnectionStateChange(IceConnectionState::Disconnected) => return,
                    Event::ChannelOpen(id, label) => match tunnels.get(&label) {
                        Some(port) => {
                            let (down, rx) = mpsc::unbounded_channel();
                            to_tcp.insert(id, down);
                            tokio::task::spawn_local(bridge(id, *port, rx, up.clone()));
                        }
                        None => rtc.direct_api().close_data_channel(id),
                    },
                    Event::ChannelData(d) => {
                        if let Some(down) = to_tcp.get(&d.id) {
                            let _ = down.send(d.data);
                        }
                    }
                    Event::ChannelClose(id) => drop(to_tcp.remove(&id)),
                    _ => {}
                }
                continue;
            }
            Err(e) => return eprintln!("webrtc-peer: {e}"),
        };
        if !rtc.is_alive() {
            return;
        }
        let woke = tokio::select! {
            r = socket.recv_from(&mut buf) => match r {
                Ok((n, source)) => Woke::Udp(n, source),
                Err(_) => return,
            },
            Some((id, data)) = from_tcp.recv() => Woke::Tcp(id, data),
            _ = tokio::time::sleep_until(deadline.into()) => Woke::Timeout,
        };
        let input = match woke {
            Woke::Udp(n, source) => match buf[..n].try_into() {
                Ok(contents) => Input::Receive(Instant::now(), Receive { proto: Protocol::Udp, source, destination: local, contents }),
                Err(_) => continue, // not STUN, DTLS or RTP
            },
            Woke::Tcp(id, Some(data)) => {
                // ponytail: write's false (SCTP buffer over its max) is not waited on; fine for ssh-sized traffic.
                if let Some(mut channel) = rtc.channel(id) {
                    let _ = channel.write(true, &data);
                }
                continue;
            }
            Woke::Tcp(id, None) => {
                rtc.direct_api().close_data_channel(id);
                continue;
            }
            Woke::Timeout => Input::Timeout(Instant::now()),
        };
        if let Err(e) = rtc.handle_input(input) {
            return eprintln!("webrtc-peer: {e}");
        }
    }
}

// One datachannel <-> 127.0.0.1:port. Bytes the channel sent before the connect are queued in `down`.
async fn bridge(id: ChannelId, port: u16, mut down: mpsc::UnboundedReceiver<Vec<u8>>, up: mpsc::UnboundedSender<(ChannelId, Option<Vec<u8>>)>) {
    let Ok(stream) = tokio::net::TcpStream::connect(("127.0.0.1", port)).await else {
        return drop(up.send((id, None)));
    };
    let (mut read, mut write) = stream.into_split();
    let mut buf = vec![0u8; 16384];
    loop {
        tokio::select! {
            n = read.read(&mut buf) => match n {
                Ok(n) if n > 0 => drop(up.send((id, Some(buf[..n].to_vec())))),
                _ => return drop(up.send((id, None))),
            },
            data = down.recv() => match data {
                Some(data) if write.write_all(&data).await.is_ok() => {}
                _ => return drop(write.shutdown().await),
            },
        }
    }
}

fn random<const N: usize>() -> [u8; N] {
    let mut bytes = [0u8; N];
    rand_core::RngCore::fill_bytes(&mut rand_core::OsRng, &mut bytes);
    bytes
}

/// age v1 (age-encryption.org/v1), X25519 recipients only, armored: what the page's typage seals the
/// offer with and opens the answer with. Rolled here over crates str0m already brings (x25519-dalek,
/// chacha20poly1305, hkdf, hmac, sha2 0.11) instead of the `age` crate and its ~50 more (i18n, plugins).
// ponytail: test-only (webrtc-peer, never in a release build); interop with the page's typage is
// covered by tests/loopback.test.mjs. Swap for the `age` crate if this ever leaves test tooling.
mod age {
    use super::*;
    use chacha20poly1305::aead::{Aead, KeyInit};
    use chacha20poly1305::{ChaCha20Poly1305, Nonce};
    use hmac::{Hmac, Mac};
    use sha2_11::Sha256;

    const INTRO: &str = "age-encryption.org/v1\n";
    const X25519_INFO: &[u8] = b"age-encryption.org/v1/X25519";
    const CHUNK: usize = 64 * 1024;
    const BEGIN: &str = "-----BEGIN AGE ENCRYPTED FILE-----";
    const END: &str = "-----END AGE ENCRYPTED FILE-----";

    pub fn public(identity: &[u8; 32]) -> [u8; 32] {
        x25519_dalek::x25519(*identity, x25519_dalek::X25519_BASEPOINT_BYTES)
    }

    pub fn recipient(public: &[u8; 32]) -> String {
        bech32::encode::<bech32::Bech32>(bech32::Hrp::parse("age").expect("valid hrp"), public).expect("short enough")
    }

    pub fn parse_recipient(s: &str) -> Result<[u8; 32], String> {
        let (hrp, data) = bech32::decode(s).map_err(|e| format!("recipient {s}: {e}"))?;
        if hrp.as_str() != "age" {
            return Err(format!("recipient {s}: not an age1 X25519 recipient"));
        }
        data.try_into().map_err(|_| format!("recipient {s}: not 32 bytes"))
    }

    fn hkdf(ikm: &[u8], salt: &[u8], info: &[u8]) -> [u8; 32] {
        let mut out = [0u8; 32];
        hkdf::Hkdf::<Sha256>::new(Some(salt), ikm).expand(info, &mut out).expect("32 bytes is a valid length");
        out
    }

    fn aead(key: &[u8; 32]) -> ChaCha20Poly1305 {
        ChaCha20Poly1305::new_from_slice(key).expect("32-byte key")
    }

    fn mac(file_key: &[u8], header: &[u8]) -> Vec<u8> {
        let mut mac = <Hmac<Sha256> as KeyInit>::new_from_slice(&hkdf(file_key, b"", b"header")).expect("any key length");
        mac.update(header);
        mac.finalize().into_bytes().to_vec()
    }

    // STREAM: 64 KiB chunks, nonce = 11-byte big-endian counter + 1 on the last chunk.
    fn stream(key: &[u8; 32], data: &[u8], seal: bool) -> Result<Vec<u8>, String> {
        let size = if seal { CHUNK } else { CHUNK + 16 };
        let chunks: Vec<&[u8]> = if data.is_empty() { vec![data] } else { data.chunks(size).collect() };
        let mut out = vec![];
        for (i, chunk) in chunks.iter().enumerate() {
            let mut nonce = [0u8; 12];
            nonce[3..11].copy_from_slice(&(i as u64).to_be_bytes());
            nonce[11] = (i + 1 == chunks.len()) as u8;
            let nonce = Nonce::from(nonce);
            let part = if seal { aead(key).encrypt(&nonce, *chunk) } else { aead(key).decrypt(&nonce, *chunk) };
            out.extend(part.map_err(|_| "age: payload authentication failed")?);
        }
        Ok(out)
    }

    fn wrap(b64: &str) -> String {
        let lines: Vec<&str> = b64.as_bytes().chunks(64).map(|c| std::str::from_utf8(c).expect("base64 is ascii")).collect();
        // a stanza body ends on a line shorter than 64 columns (an empty one if need be)
        if b64.len() % 64 == 0 { format!("{}\n\n", lines.join("\n")) } else { format!("{}\n", lines.join("\n")) }
    }

    pub fn seal(to: &[u8; 32], plaintext: &[u8]) -> String {
        let file_key: [u8; 16] = random();
        let ephemeral: [u8; 32] = random();
        let share = public(&ephemeral);
        let shared = x25519_dalek::x25519(ephemeral, *to);
        let wrap_key = hkdf(&shared, &[share.as_slice(), to.as_slice()].concat(), X25519_INFO);
        let body = aead(&wrap_key).encrypt(&Nonce::from([0u8; 12]), file_key.as_slice()).expect("encrypts");
        let mut header = format!("{INTRO}-> X25519 {}\n{}---", STANDARD_NO_PAD.encode(share), wrap(&STANDARD_NO_PAD.encode(body)));
        header = format!("{header} {}\n", STANDARD_NO_PAD.encode(mac(&file_key, header.as_bytes())));
        let nonce: [u8; 16] = random();
        let payload = stream(&hkdf(&file_key, &nonce, b"payload"), plaintext, true).expect("seals");
        let file = [header.as_bytes(), &nonce, &payload].concat();
        let b64 = STANDARD.encode(file);
        let lines: Vec<&str> = b64.as_bytes().chunks(64).map(|c| std::str::from_utf8(c).expect("base64 is ascii")).collect();
        format!("{BEGIN}\n{}\n{END}\n", lines.join("\n"))
    }

    pub fn open(identity: &[u8; 32], armored: &str) -> Result<Vec<u8>, String> {
        let inner = armored.trim().strip_prefix(BEGIN).and_then(|s| s.strip_suffix(END)).ok_or("age: not armored")?;
        let file = STANDARD.decode(inner.split_whitespace().collect::<String>()).map_err(|e| format!("age armor: {e}"))?;
        let text = String::from_utf8_lossy(&file);
        let end = text.find("\n---").ok_or("age: no header MAC")? + 4;
        let header = text[..end].strip_prefix(INTRO).ok_or("age: not age-encryption.org/v1")?;
        let mac_line_end = text[end..].find('\n').ok_or("age: truncated header")? + end;
        let mut lines = header.lines().peekable();
        let ours = public(identity);
        let mut file_key = None;
        while let Some(line) = lines.next() {
            let Some(args) = line.strip_prefix("-> ") else { break };
            let mut body = String::new();
            while let Some(l) = lines.next() {
                body.push_str(l);
                if l.len() < 64 {
                    break;
                }
            }
            let args: Vec<&str> = args.split(' ').collect();
            if args.len() != 2 || args[0] != "X25519" || file_key.is_some() {
                continue; // another recipient type (or grease)
            }
            let (Ok(share), Ok(body)) = (STANDARD_NO_PAD.decode(args[1]), STANDARD_NO_PAD.decode(&body)) else { continue };
            let Ok(share) = <[u8; 32]>::try_from(share) else { continue };
            let shared = x25519_dalek::x25519(*identity, share);
            let wrap_key = hkdf(&shared, &[share.as_slice(), ours.as_slice()].concat(), X25519_INFO);
            file_key = aead(&wrap_key).decrypt(&Nonce::from([0u8; 12]), body.as_slice()).ok();
        }
        let file_key = file_key.ok_or("age: not sealed to this identity")?;
        let got = STANDARD_NO_PAD.decode(text[end..mac_line_end].trim()).map_err(|_| "age: bad header MAC")?;
        if got != mac(&file_key, &file[..end]) {
            return Err("age: header MAC mismatch".into());
        }
        let rest = &file[mac_line_end + 1..];
        if rest.len() < 16 {
            return Err("age: truncated payload".into());
        }
        stream(&hkdf(&file_key, &rest[..16], b"payload"), &rest[16..], false)
    }

    #[cfg(test)]
    mod tests {
        use super::*;

        #[test]
        fn only_the_listener_host_and_the_page_origin_get_in() {
            let bound: SocketAddr = "127.0.0.1:4000".parse().unwrap();
            let page = Some("http://127.0.0.1:5000");
            let headers = |pairs: &[(&str, &[u8])]| {
                let mut h = http::HeaderMap::new();
                for (k, v) in pairs {
                    h.insert(http::HeaderName::from_bytes(k.as_bytes()).unwrap(), http::HeaderValue::from_bytes(v).unwrap());
                }
                h
            };
            for ok in [headers(&[("host", b"127.0.0.1:4000")]), headers(&[("host", b"LOCALHOST:4000"), ("origin", b"http://127.0.0.1:5000")])] {
                assert_eq!(refused(&ok, bound, page), None, "{ok:?}");
            }
            for bad in [
                headers(&[]),
                headers(&[("host", b"evil.example:4000")]),
                headers(&[("host", b"127.0.0.1:4001")]),
                headers(&[("host", b"\xff")]),
                headers(&[("host", b"127.0.0.1:4000"), ("origin", b"https://evil.example")]),
                headers(&[("host", b"127.0.0.1:4000"), ("origin", b"\xff")]),
            ] {
                assert!(refused(&bad, bound, page).is_some(), "{bad:?}");
            }
            assert!(refused(&headers(&[("host", b"127.0.0.1:4000"), ("origin", b"http://127.0.0.1:5000")]), bound, None).is_some(), "no --origin: no page");
        }

        #[test]
        fn seal_open_round_trip_and_wrong_identity() {
            let (me, other): ([u8; 32], [u8; 32]) = (random(), random());
            let message = vec![7u8; 70000]; // two STREAM chunks
            let sealed = seal(&public(&me), &message);
            assert!(sealed.starts_with(BEGIN) && sealed.trim_end().ends_with(END), "{sealed}");
            assert_eq!(open(&me, &sealed).unwrap(), message);
            assert_eq!(open(&me, &seal(&public(&me), b"")).unwrap(), b"");
            assert_eq!(open(&other, &sealed).unwrap_err(), "age: not sealed to this identity");
            let r = recipient(&public(&me));
            assert!(r.starts_with("age1") && r.len() == 62, "{r}");
            assert_eq!(parse_recipient(&r).unwrap(), public(&me));
        }
    }
}