swarm.rs 89.7%

663 of 739 lines run, 70 of 131 functions · run #195 on main · run missed
//! The swarm (oxygenflow#47, docs/jetstream.md §8.1): a runner with a
//! `[runner.swarm]` section serves the image chunks it holds to its peers,
//! and fetches an image it lacks from its peers and the hub together.
//!
//! The server is a small HTTPS endpoint on the swarm address: a peer
//! shows the ticket its hub signed and gets a session, then asks which
//! chunks of an image this runner has and for chunks by hash. A chunk is
//! served from the chunk cache, or read out of a whole image by its
//! manifest. Every chunk is checked by the side that asked for it, so a
//! peer can only fail to help.
//!
//! Each runner makes its own TLS certificate on first use and reports its
//! BLAKE3 with its swarm address; the hub hands that on with the peer, and
//! a fetcher talks only to a peer that shows that certificate and holds
//! its key. Nobody on the network reads the chunks, tickets or sessions.

use crate::home::{Home, RunnerConfig};
use crate::images::{cache_path, chunk_cache, manifest_path, save_manifest, write_image, Client, Moved};
use anyhow::{bail, Context, Result};
use flow_core::api::{SwarmAddr, SwarmPeer};
use jet::images::{is_hash, Manifest};
use jet::swarm::{Bits, Pick, Picker, Source, Ticket};
use oxy_core::identity::Identity;
use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::crypto::CryptoProvider;
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName, UnixTime};
use rustls::{DigitallySignedStruct, SignatureScheme};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};
use std::net::{TcpListener, TcpStream};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

/// Connections served at once; more are told 503 and come back shortly.
const MAX_CONNS: usize = 16;
/// Connections past this are closed without a word.
const MAX_OPEN: usize = 4 * MAX_CONNS;
/// Requests in flight to one peer, and to the hub.
const PEER_WORKERS: usize = 4;
const HUB_WORKERS: usize = 8;
/// How often a fetcher asks its peers what they have.
const REFRESH: Duration = Duration::from_millis(500);
/// Hub failures a fetch survives before it gives up.
const HUB_FAILURES: usize = 5;
/// The largest request body the server reads: a ticket.
const MAX_BODY: usize = 64 * 1024;
/// Names the part files of chunks two sources fetch at once apart.
static PART: AtomicU64 = AtomicU64::new(0);

fn now_ms() -> u64 {
    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis() as u64).unwrap_or(0)
}

/// An image this runner offers: which of its needed chunks are here
/// (`None`: all of them, the image is whole) and which it is fetching now.
struct Offer {
    bits: Option<Bits>,
    fetching: Bits,
}

struct State {
    root: PathBuf,
    site: String,
    client: Client,
    /// The hub's identity, which signs tickets; fetched when first needed.
    hub: Mutex<Option<Identity>>,
    offers: Mutex<HashMap<String, Offer>>,
    /// Where a chunk lies in a whole image: the image, offset and length.
    located: Mutex<HashMap<String, (PathBuf, u64, u32)>>,
    /// Sessions: token to (peer name, expiry).
    sessions: Mutex<HashMap<String, (String, u64)>>,
    conns: AtomicUsize,
    /// The upload allowance: bytes in hand (negative: owed) and when it
    /// was last topped up.
    bucket: Mutex<(f64, Instant)>,
    rate: f64,
    tls: Arc<rustls::ServerConfig>,
}

/// This runner's place in the swarm.
pub struct Swarm {
    state: Arc<State>,
    pub addr: SwarmAddr,
}

impl Swarm {
    /// Start serving peers when the runner is configured for it.
    pub fn start(home: &Home, cfg: &RunnerConfig) -> Result<Option<Swarm>> {
        let Some(sc) = &cfg.swarm else { return Ok(None) };
        let listener = TcpListener::bind(&sc.listen).with_context(|| format!("swarm: listen on {}", sc.listen))?;
        let client = Client::new(&cfg.hub, &cfg.token, cfg.ca.as_deref())?;
        let hub = client.hub_identity().map_err(|e| eprintln!("swarm: the hub's key: {e:#}; trying again when a peer comes")).ok();
        let rate = sc.upload_mib_s.unwrap_or(50).max(1) as f64 * (1 << 20) as f64;
        let (cert, key) = certificate(&home.root)?;
        let fingerprint = blake3::hash(&cert).to_hex().to_string();
        let tls = server_tls(cert, key)?;
        let state = Arc::new(State {
            root: home.root.clone(),
            site: sc.site.clone(),
            client,
            hub: Mutex::new(hub),
            offers: Mutex::new(HashMap::new()),
            located: Mutex::new(HashMap::new()),
            sessions: Mutex::new(HashMap::new()),
            conns: AtomicUsize::new(0),
            bucket: Mutex::new((rate, Instant::now())),
            rate,
            tls,
        });
        // Whole images: offered from their manifests; one without a
        // manifest (imported here, not fetched) is scanned for one, in the
        // background and one at a time.
        let mut unscanned = Vec::new();
        for h in home.images() {
            match std::fs::read(manifest_path(&home.root, &h)).ok().and_then(|b| serde_json::from_slice::<Manifest>(&b).ok()) {
                Some(m) if m.hash == h => offer_whole(&state, &m, &home.rootfs(&h)),
                _ => unscanned.push((h.clone(), home.rootfs(&h), home.image_arch(&h))),
            }
        }
        let scanning = state.clone();
        std::thread::spawn(move || {
            for (h, path, arch) in unscanned {
                match jet::images::scan_file(&path, arch, |_, _| Ok(())) {
                    Ok(m) if m.hash == h => {
                        save_manifest(&scanning.root, &m).ok();
                        offer_whole(&scanning, &m, &path);
                    }
                    Ok(_) => eprintln!("swarm: {} does not hold the image its name says; not offering it", path.display()),
                    Err(e) => eprintln!("swarm: scan {}: {e}", path.display()),
                }
            }
        });
        let serving = state.clone();
        std::thread::spawn(move || accept(listener, serving));
        let addr = SwarmAddr { addr: sc.advertise.clone().unwrap_or_else(|| sc.listen.clone()), site: sc.site.clone(), cert: fingerprint };
        eprintln!(
            "swarm: serving image chunks on {} over TLS (certificate {}, site {:?}, up to {} MiB/s)",
            sc.listen,
            &addr.cert[..12],
            sc.site,
            sc.upload_mib_s.unwrap_or(50)
        );
        Ok(Some(Swarm { state, addr }))
    }

    /// Fetch image `hash` into `home` from the peers and the hub (§8.1):
    /// the rarest chunks from peers first, the hub for what no peer has.
    /// `Ok(None)` when the hub does not hold it.
    pub fn fetch(&self, home: &Home, hash: &str) -> Result<Option<Moved>> {
        let start = Instant::now();
        let client = &self.state.client;
        let Some(m) = client.manifest(hash)? else { return Ok(None) };
        if m.hash != hash {
            bail!("the hub's manifest for {} names {}", &hash[..12], &m.hash[..12]);
        }
        save_manifest(&home.root, &m)?;
        let needed: Vec<String> = m.needed().into_iter().map(String::from).collect();
        let cache = chunk_cache(home);
        let mut have = Bits::new(needed.len());
        for (i, h) in needed.iter().enumerate() {
            if cache_path(&cache, h).is_file() {
                have.set(i);
            }
        }
        self.state.offers.lock().unwrap().insert(hash.into(), Offer { bits: Some(have.clone()), fetching: Bits::new(needed.len()) });
        let announced = client
            .announce(hash, "fetching")
            .map_err(|e| eprintln!("swarm: announce {}: {e:#}; fetching from the hub alone", &hash[..12]))
            .ok();
        let seed = now_ms() ^ (std::process::id() as u64) << 32;
        let picker = Mutex::new(Picker::new(needed.len(), have, seed));
        let peers: Mutex<Vec<Arc<Peer>>> = Mutex::new(Vec::new());
        if let Some(a) = &announced {
            self.join(&a.peers, &a.ticket, hash, needed.len(), &peers, &picker);
        }
        let (bytes, peer_bytes, got) = (AtomicU64::new(0), AtomicU64::new(0), AtomicUsize::new(0));
        let (failed, hub_failures) = (Mutex::new(None::<anyhow::Error>), AtomicUsize::new(0));
        {
            let (picker, peers, needed, cache) = (&picker, &peers, &needed, &cache);
            let (bytes, peer_bytes, got, failed, hub_failures) = (&bytes, &peer_bytes, &got, &failed, &hub_failures);
            let (state, m) = (&self.state, &m);
            let worker = move |src: Source| loop {
                if failed.lock().unwrap().is_some() {
                    break;
                }
                let peer = match src {
                    Source::Peer(p) => Some(peers.lock().unwrap()[p].clone()),
                    Source::Hub => None,
                };
                if peer.as_ref().is_some_and(|p| !p.alive.load(Ordering::Relaxed)) {
                    break;
                }
                // Taken apart from the match, so the lock is not held
                // while waiting.
                let pick = picker.lock().unwrap().pick(src);
                let i = match pick {
                    Pick::Done => break,
                    Pick::Wait => {
                        std::thread::sleep(Duration::from_millis(20));
                        continue;
                    }
                    Pick::Chunk(i) => i,
                };
                if let Some(o) = state.offers.lock().unwrap().get_mut(&m.hash) {
                    o.fetching.set(i);
                }
                let h = &needed[i];
                let fetched = match &peer {
                    Some(p) => self.peer_chunk(p, h),
                    None => client.chunk(h),
                };
                let kept = fetched.and_then(|b| {
                    if blake3::hash(&b).to_hex().as_str() != h {
                        bail!("chunk {} came back with other bytes", &h[..12]);
                    }
                    let p = cache_path(cache, h);
                    std::fs::create_dir_all(p.parent().expect("a chunk directory"))?;
                    let tmp = p.with_extension(format!("part{}", PART.fetch_add(1, Ordering::Relaxed)));
                    std::fs::write(&tmp, &b)?;
                    std::fs::rename(&tmp, &p)?;
                    Ok(b.len() as u64)
                });
                match kept {
                    Ok(n) => {
                        if let Some(Offer { bits: Some(b), fetching }) = state.offers.lock().unwrap().get_mut(&m.hash) {
                            b.set(i);
                            fetching.clear(i);
                        }
                        picker.lock().unwrap().got(i);
                        bytes.fetch_add(n, Ordering::Relaxed);
                        got.fetch_add(1, Ordering::Relaxed);
                        if peer.is_some() {
                            peer_bytes.fetch_add(n, Ordering::Relaxed);
                        }
                    }
                    Err(e) => {
                        picker.lock().unwrap().failed(i);
                        if let Some(o) = state.offers.lock().unwrap().get_mut(&m.hash) {
                            o.fetching.clear(i);
                        }
                        match (&peer, src) {
                            // A busy peer is asked again shortly.
                            (Some(_), _) if e.is::<Busy>() => std::thread::sleep(Duration::from_millis(200)),
                            (Some(p), Source::Peer(idx)) => {
                                eprintln!("swarm: peer {} failed ({e:#}); carrying on without it", p.name);
                                p.alive.store(false, Ordering::Relaxed);
                                picker.lock().unwrap().drop_peer(idx);
                                break;
                            }
                            _ => {
                                if hub_failures.fetch_add(1, Ordering::Relaxed) + 1 >= HUB_FAILURES {
                                    *failed.lock().unwrap() = Some(e);
                                    break;
                                }
                                std::thread::sleep(Duration::from_millis(500));
                            }
                        }
                    }
                }
            };
            std::thread::scope(|s| {
                for _ in 0..HUB_WORKERS {
                    s.spawn(move || worker(Source::Hub));
                }
                let mut spawned = 0;
                let mut spawn_new = || {
                    let n = peers.lock().unwrap().len();
                    for p in spawned..n {
                        for _ in 0..PEER_WORKERS {
                            s.spawn(move || worker(Source::Peer(p)));
                        }
                    }
                    spawned = n;
                };
                spawn_new();
                // Keep the peers' maps fresh, and find new peers while
                // fetching, until every chunk is here.
                let mut asked_hub = Instant::now();
                let interval = Duration::from_secs(announced.as_ref().map(|a| a.interval_s).unwrap_or(30).max(5));
                let mut last = Instant::now();
                while picker.lock().unwrap().missing() > 0 && failed.lock().unwrap().is_none() {
                    std::thread::sleep(Duration::from_millis(100));
                    if last.elapsed() < REFRESH {
                        continue;
                    }
                    last = Instant::now();
                    let now: Vec<(usize, Arc<Peer>)> =
                        peers.lock().unwrap().iter().cloned().enumerate().filter(|(_, p)| p.alive.load(Ordering::Relaxed)).collect();
                    for (idx, p) in now {
                        match self.peer_have(&p, hash, needed.len()) {
                            Ok((bits, fetching)) => picker.lock().unwrap().peer_fetching(idx, bits, fetching),
                            Err(e) if e.is::<Busy>() => {}
                            Err(_) => {
                                p.alive.store(false, Ordering::Relaxed);
                                picker.lock().unwrap().drop_peer(idx);
                            }
                        }
                    }
                    if announced.is_some() && asked_hub.elapsed() >= interval {
                        asked_hub = Instant::now();
                        if let Ok(a) = client.announce(hash, "fetching") {
                            self.join(&a.peers, &a.ticket, hash, needed.len(), peers, picker);
                            spawn_new();
                        }
                    }
                }
            });
        }
        if let Some(e) = failed.into_inner().unwrap() {
            client.announce(hash, "stopped").ok();
            return Err(e);
        }
        write_image(home, &m)?;
        offer_whole(&self.state, &m, &home.rootfs(hash));
        client.announce(hash, "done").ok();
        Ok(Some(Moved {
            chunks: got.into_inner(),
            bytes: bytes.into_inner(),
            peer_bytes: peer_bytes.into_inner(),
            elapsed: start.elapsed(),
        }))
    }

    /// Admit this runner to the peers it does not know yet: a session from
    /// each, and what each has.
    fn join(&self, found: &[SwarmPeer], ticket: &str, hash: &str, len: usize, peers: &Mutex<Vec<Arc<Peer>>>, picker: &Mutex<Picker>) {
        for sp in found {
            if peers.lock().unwrap().iter().any(|p| p.runner == sp.runner) {
                continue;
            }
            if !is_hash(&sp.cert) {
                eprintln!("swarm: peer {} gave no certificate (an older flow); not asking it", sp.name);
                continue;
            }
            let joined = hello(&sp.addr, &sp.cert, ticket).and_then(|(agent, session)| {
                let p = Peer {
                    runner: sp.runner,
                    name: sp.name.clone(),
                    base: format!("https://{}/swarm/v1", sp.addr),
                    agent,
                    session,
                    alive: AtomicBool::new(true),
                };
                let (bits, fetching) = self.peer_have(&p, hash, len)?;
                Ok((p, bits, fetching))
            });
            match joined {
                Ok((p, bits, fetching)) => {
                    let mut list = peers.lock().unwrap();
                    picker.lock().unwrap().peer_fetching(list.len(), bits, fetching);
                    list.push(Arc::new(p));
                }
                Err(e) => eprintln!("swarm: peer {} ({}): {e:#}", sp.name, sp.addr),
            }
        }
    }

    /// What a peer has of image `hash`, and what it is fetching now:
    /// nothing yet when it has not begun (the tracker may name it a moment
    /// before it starts).
    fn peer_have(&self, p: &Peer, hash: &str, len: usize) -> Result<(Bits, Bits)> {
        let url = format!("{}/images/{hash}/have", p.base);
        let r = match p.agent.get(&url).set("Authorization", &format!("Bearer {}", p.session)).call() {
            Ok(r) => r,
            Err(ureq::Error::Status(404, _)) => return Ok((Bits::new(len), Bits::new(len))),
            Err(e) => return Err(peer_error(&url, e)),
        };
        let mut body = Vec::new();
        r.into_reader().take(len as u64 / 4 + 64).read_to_end(&mut body)?;
        if body == b"complete" {
            return Ok((Bits::full(len), Bits::new(len)));
        }
        let (have, fetching) = body.split_at(len.div_ceil(8).min(body.len()));
        Ok((Bits::from_bytes(len, have), Bits::from_bytes(len, fetching)))
    }

    fn peer_chunk(&self, p: &Peer, hash: &str) -> Result<Vec<u8>> {
        let url = format!("{}/chunks/{hash}", p.base);
        let r = p.agent.get(&url).set("Authorization", &format!("Bearer {}", p.session)).call().map_err(|e| peer_error(&url, e))?;
        let mut out = Vec::new();
        r.into_reader().take(oxy_core::chunk::MAX_SIZE as u64 + 1).read_to_end(&mut out)?;
        Ok(out)
    }
}

/// A peer this fetch asks for chunks, over a connection that trusts only
/// its certificate.
struct Peer {
    runner: i64,
    name: String,
    base: String,
    agent: ureq::Agent,
    session: String,
    alive: AtomicBool,
}

/// A peer at its connection limit: ask it again shortly.
#[derive(Debug)]
struct Busy;

impl std::fmt::Display for Busy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("the peer is busy")
    }
}

impl std::error::Error for Busy {}

fn peer_error(url: &str, e: ureq::Error) -> anyhow::Error {
    match e {
        ureq::Error::Status(503, _) => anyhow::Error::new(Busy),
        ureq::Error::Status(code, r) => anyhow::anyhow!("{url}: HTTP {code}: {}", r.into_string().unwrap_or_default().trim()),
        e => anyhow::anyhow!("{url}: {e}"),
    }
}

/// Show `ticket` to the peer at `addr`, trusting only the certificate
/// `cert`: a connection to it and a session.
fn hello(addr: &str, cert: &str, ticket: &str) -> Result<(ureq::Agent, String)> {
    let agent = peer_agent(cert)?;
    let url = format!("https://{addr}/swarm/v1/hello");
    let r: serde_json::Value = agent.post(&url).send_string(ticket).map_err(|e| peer_error(&url, e))?.into_json()?;
    let session = r["session"].as_str().map(String::from).context("the peer gave no session")?;
    Ok((agent, session))
}

fn provider() -> Arc<CryptoProvider> {
    Arc::new(rustls::crypto::ring::default_provider())
}

/// A client that trusts one certificate, by its BLAKE3, whatever name
/// the peer goes by: the one the hub said this peer has.
fn peer_agent(cert: &str) -> Result<ureq::Agent> {
    let provider = provider();
    let cfg = rustls::ClientConfig::builder_with_provider(provider.clone())
        .with_safe_default_protocol_versions()?
        .dangerous()
        .with_custom_certificate_verifier(Arc::new(Pinned { cert: cert.to_string(), provider }))
        .with_no_client_auth();
    Ok(ureq::AgentBuilder::new().tls_config(Arc::new(cfg)).timeout_connect(Duration::from_secs(2)).timeout(Duration::from_secs(20)).build())
}

/// Trusts the certificate with this BLAKE3 and no other. The handshake's
/// signature is still checked against it, so only the holder of its key
/// passes.
#[derive(Debug)]
struct Pinned {
    cert: String,
    provider: Arc<CryptoProvider>,
}

impl ServerCertVerifier for Pinned {
    fn verify_server_cert(
        &self,
        end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _name: &ServerName<'_>,
        _ocsp: &[u8],
        _now: UnixTime,
    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
        if blake3::hash(end_entity.as_ref()).to_hex().as_str() == self.cert {
            Ok(ServerCertVerified::assertion())
        } else {
            Err(rustls::Error::General("not the peer the hub named: its certificate is another".into()))
        }
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls12_signature(message, cert, dss, &self.provider.signature_verification_algorithms)
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls13_signature(message, cert, dss, &self.provider.signature_verification_algorithms)
    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        self.provider.signature_verification_algorithms.supported_schemes()
    }
}

/// The runner's swarm certificate and key (`swarm/cert.der`, and
/// `swarm/key.der` readable by the runner alone), made on first use and
/// kept, so its fingerprint holds across restarts.
fn certificate(root: &Path) -> Result<(Vec<u8>, Vec<u8>)> {
    let dir = root.join("swarm");
    let (cert_path, key_path) = (dir.join("cert.der"), dir.join("key.der"));
    if let (Ok(c), Ok(k)) = (std::fs::read(&cert_path), std::fs::read(&key_path)) {
        return Ok((c, k));
    }
    std::fs::create_dir_all(&dir)?;
    let made = rcgen::generate_simple_self_signed(vec!["flow-swarm".to_string()]).context("make the swarm's certificate")?;
    let (cert, key) = (made.cert.der().to_vec(), made.signing_key.serialize_der());
    {
        let mut f = crate::portable::private(std::fs::OpenOptions::new().write(true).create(true).truncate(true)).open(&key_path)?;
        f.write_all(&key)?;
    }
    std::fs::write(&cert_path, &cert)?;
    Ok((cert, key))
}

fn server_tls(cert: Vec<u8>, key: Vec<u8>) -> Result<Arc<rustls::ServerConfig>> {
    let cfg = rustls::ServerConfig::builder_with_provider(provider())
        .with_safe_default_protocol_versions()?
        .with_no_client_auth()
        .with_single_cert(vec![CertificateDer::from(cert)], PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key)))
        .context("the swarm's certificate")?;
    Ok(Arc::new(cfg))
}

/// Offer a whole image: every chunk of it can be read out of `path`.
fn offer_whole(state: &State, m: &Manifest, path: &Path) {
    let mut located = state.located.lock().unwrap();
    let mut at = 0u64;
    for c in &m.chunks {
        if !c.is_zeros() {
            located.entry(c.hash().to_string()).or_insert_with(|| (path.to_path_buf(), at, c.1));
        }
        at += c.1 as u64;
    }
    state.offers.lock().unwrap().insert(m.hash.clone(), Offer { bits: None, fetching: Bits::new(0) });
}

fn accept(listener: TcpListener, state: Arc<State>) {
    for conn in listener.incoming() {
        let Ok(conn) = conn else { continue };
        if state.conns.load(Ordering::Relaxed) >= MAX_OPEN {
            continue;
        }
        let state = state.clone();
        std::thread::spawn(move || {
            // Past the limit a peer is still answered, with 503, so it
            // comes back rather than giving up on this runner.
            let busy = state.conns.fetch_add(1, Ordering::Relaxed) >= MAX_CONNS;
            serve_tls(conn, &state, busy).ok();
            state.conns.fetch_sub(1, Ordering::Relaxed);
        });
    }
}

fn serve_tls(conn: TcpStream, state: &State, busy: bool) -> std::io::Result<()> {
    conn.set_read_timeout(Some(Duration::from_secs(30)))?;
    conn.set_write_timeout(Some(Duration::from_secs(30)))?;
    conn.set_nodelay(true)?;
    let tls = rustls::ServerConnection::new(state.tls.clone()).map_err(std::io::Error::other)?;
    serve_conn(rustls::StreamOwned::new(tls, conn), state, busy)
}

fn respond(conn: &mut impl Write, code: u16, kind: &str, body: &[u8]) -> std::io::Result<()> {
    let reason = match code {
        200 => "OK",
        400 => "Bad Request",
        401 => "Unauthorized",
        403 => "Forbidden",
        404 => "Not Found",
        503 => "Service Unavailable",
        _ => "Error",
    };
    write!(conn, "HTTP/1.1 {code} {reason}\r\nContent-Type: {kind}\r\nContent-Length: {}\r\n\r\n", body.len())?;
    conn.write_all(body)?;
    conn.flush()
}

/// One connection: requests one after another while the peer keeps it.
fn serve_conn<S: Read + Write>(conn: S, state: &State, busy: bool) -> std::io::Result<()> {
    let mut r = BufReader::new(conn);
    loop {
        let mut line = String::new();
        if r.read_line(&mut line)? == 0 {
            return Ok(());
        }
        let mut parts = line.split_whitespace();
        let (method, path) = (parts.next().unwrap_or("").to_string(), parts.next().unwrap_or("").to_string());
        let (mut length, mut auth, mut close) = (0usize, String::new(), false);
        loop {
            let mut h = String::new();
            if r.read_line(&mut h)? == 0 {
                return Ok(());
            }
            let h = h.trim_end();
            if h.is_empty() {
                break;
            }
            if let Some((k, v)) = h.split_once(':') {
                match k.trim().to_ascii_lowercase().as_str() {
                    "content-length" => length = v.trim().parse().unwrap_or(usize::MAX),
                    "authorization" => auth = v.trim().to_string(),
                    "connection" => close = v.trim().eq_ignore_ascii_case("close"),
                    _ => {}
                }
            }
        }
        if length > MAX_BODY {
            return respond(r.get_mut(), 400, "text/plain", b"too large");
        }
        let mut body = vec![0; length];
        r.read_exact(&mut body)?;
        if busy {
            return respond(r.get_mut(), 503, "text/plain", b"busy: ask again shortly");
        }
        let (code, kind, reply) = route(state, &method, &path, &auth, &body);
        if code == 200 && kind == "application/octet-stream" {
            throttle(state, reply.len());
        }
        respond(r.get_mut(), code, kind, &reply)?;
        if close {
            return Ok(());
        }
    }
}

fn route(state: &State, method: &str, path: &str, auth: &str, body: &[u8]) -> (u16, &'static str, Vec<u8>) {
    let text = |code, s: &str| (code, "text/plain", s.as_bytes().to_vec());
    if method == "POST" && path == "/swarm/v1/hello" {
        return match admit(state, &String::from_utf8_lossy(body)) {
            Ok(session) => (200, "application/json", serde_json::json!({ "session": session }).to_string().into_bytes()),
            Err(why) => text(403, &why),
        };
    }
    let session = auth.strip_prefix("Bearer ").unwrap_or("");
    let admitted = {
        let mut sessions = state.sessions.lock().unwrap();
        let now = now_ms();
        sessions.retain(|_, (_, until)| *until > now);
        sessions.contains_key(session)
    };
    if !admitted {
        return text(401, "show a ticket at /swarm/v1/hello first");
    }
    let rest = path.strip_prefix("/swarm/v1/").unwrap_or("");
    match (method, rest.split('/').collect::<Vec<_>>().as_slice()) {
        ("GET", ["images", image, "have"]) => match state.offers.lock().unwrap().get(*image) {
            Some(Offer { bits: None, .. }) => (200, "text/plain", b"complete".to_vec()),
            Some(Offer { bits: Some(b), fetching }) => (200, "application/x-have", [b.to_bytes(), fetching.to_bytes()].concat()),
            None => text(404, "this runner does not hold that image"),
        },
        ("GET", ["chunks", chunk]) if is_hash(chunk) => match read_chunk(state, chunk) {
            Some(b) => (200, "application/octet-stream", b),
            None => text(404, "no such chunk here"),
        },
        _ => text(404, "not a swarm path"),
    }
}

/// Admit a peer that shows a ticket from this runner's hub, for this
/// site: a session until the ticket expires.
fn admit(state: &State, ticket: &str) -> std::result::Result<String, String> {
    let hub = {
        let mut hub = state.hub.lock().unwrap();
        if hub.is_none() {
            *hub = state.client.hub_identity().ok();
        }
        hub.clone().ok_or("this runner cannot reach its hub to check tickets")?
    };
    let t = Ticket::open(ticket, &hub, now_ms())?;
    if t.site != state.site {
        return Err(format!("this runner is in site {:?}, the ticket's runner in {:?}", state.site, t.site));
    }
    let session = random_token().map_err(|e| format!("make a session: {e}"))?;
    state.sessions.lock().unwrap().insert(session.clone(), (t.name, t.expires_ms));
    Ok(session)
}

fn random_token() -> std::io::Result<String> {
    let mut b = [0u8; 32];
    std::fs::File::open("/dev/urandom")?.read_exact(&mut b)?;
    Ok(hex::encode(b))
}

fn read_chunk(state: &State, hash: &str) -> Option<Vec<u8>> {
    if let Ok(b) = std::fs::read(cache_path(&state.root.join("chunks"), hash)) {
        return Some(b);
    }
    let (path, at, len) = state.located.lock().unwrap().get(hash).cloned()?;
    let mut f = std::fs::File::open(path).ok()?;
    f.seek(SeekFrom::Start(at)).ok()?;
    let mut b = vec![0; len as usize];
    f.read_exact(&mut b).ok()?;
    Some(b)
}

/// Hold back until sending `n` more bytes keeps within the upload rate.
fn throttle(state: &State, n: usize) {
    let wait = {
        let mut b = state.bucket.lock().unwrap();
        let now = Instant::now();
        b.0 = (b.0 + now.duration_since(b.1).as_secs_f64() * state.rate).min(state.rate);
        b.1 = now;
        b.0 -= n as f64;
        if b.0 < 0.0 {
            -b.0 / state.rate
        } else {
            0.0
        }
    };
    if wait > 0.0 {
        std::thread::sleep(Duration::from_secs_f64(wait));
    }
}

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

    /// A hub that holds one image and runs the tracker: it serves the
    /// manifest and chunks (counting the chunks), and answers announces
    /// with every other runner as a peer and a ticket for the caller,
    /// known by its token.
    struct FakeHub {
        url: String,
        chunks_served: Arc<AtomicUsize>,
        runners: Arc<Mutex<Vec<(String, SwarmPeer)>>>,
    }

    fn fake_hub(m: Manifest, image: Vec<u8>, site: &str) -> FakeHub {
        slow_hub(m, image, site, Duration::ZERO)
    }

    /// A hub that takes `delay` over each chunk.
    fn slow_hub(m: Manifest, image: Vec<u8>, site: &str, delay: Duration) -> FakeHub {
        let sk = SecretKey::generate();
        let id = sk.identity("hub", "hub@example", 1);
        let key = id.hash();
        let mut chunks = HashMap::new();
        let mut at = 0usize;
        for c in &m.chunks {
            chunks.insert(c.hash().to_string(), image[at..at + c.len()].to_vec());
            at += c.len();
        }
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        let served = Arc::new(AtomicUsize::new(0));
        let runners: Arc<Mutex<Vec<(String, SwarmPeer)>>> = Arc::new(Mutex::new(Vec::new()));
        let (count, known, site) = (served.clone(), runners.clone(), site.to_string());
        let shared = Arc::new((m, chunks, id, key, sk));
        std::thread::spawn(move || {
            for conn in listener.incoming() {
                let Ok(conn) = conn else { continue };
                let (shared, count, known, site) = (shared.clone(), count.clone(), known.clone(), site.clone());
                std::thread::spawn(move || {
                    let (m, chunks, id, key, sk) = &*shared;
                    let mut out = conn.try_clone().unwrap();
                    let mut r = BufReader::new(conn);
                    loop {
                        let mut line = String::new();
                        if r.read_line(&mut line).unwrap_or(0) == 0 {
                            return;
                        }
                        let path = line.split_whitespace().nth(1).unwrap_or("").to_string();
                        let (mut len, mut token) = (0, String::new());
                        loop {
                            let mut h = String::new();
                            r.read_line(&mut h).unwrap();
                            let h = h.trim_end().to_string();
                            if h.is_empty() {
                                break;
                            }
                            let lower = h.to_ascii_lowercase();
                            if let Some(v) = lower.strip_prefix("content-length:") {
                                len = v.trim().parse().unwrap();
                            }
                            if let Some(v) = h.strip_prefix("Authorization: Bearer ") {
                                token = v.to_string();
                            }
                        }
                        let mut body = vec![0; len];
                        r.read_exact(&mut body).unwrap();
                        let (code, reply): (u16, Vec<u8>) = if path == "/.well-known/oxyhub-key" {
                            (200, id.encode())
                        } else if path == format!("/api/v1/flow/images/{}", m.hash) {
                            (200, serde_json::to_vec(m).unwrap())
                        } else if let Some(c) = path.strip_prefix("/api/v1/flow/images/chunks/") {
                            count.fetch_add(1, Ordering::Relaxed);
                            std::thread::sleep(delay);
                            (200, chunks[c].clone())
                        } else if path == "/api/v1/flow/swarm/announce" {
                            let runners = known.lock().unwrap();
                            let me = runners.iter().find(|(t, _)| *t == token).map(|(_, p)| p.clone()).unwrap();
                            let ticket = Ticket {
                                hub: key.to_hex(),
                                runner: me.runner,
                                name: me.name.clone(),
                                site: site.clone(),
                                expires_ms: now_ms() + 60_000,
                            };
                            let peers: Vec<SwarmPeer> = runners.iter().filter(|(t, _)| *t != token).map(|(_, p)| p.clone()).collect();
                            let a = flow_core::api::Announced { ticket: ticket.seal(key, sk), peers, interval_s: 5 };
                            (200, serde_json::to_vec(&a).unwrap())
                        } else {
                            (404, b"no".to_vec())
                        };
                        let head = format!("HTTP/1.1 {code} X\r\nContent-Length: {}\r\n\r\n", reply.len());
                        if out.write_all(head.as_bytes()).and_then(|_| out.write_all(&reply)).is_err() {
                            return;
                        }
                    }
                });
            }
        });
        FakeHub { url, chunks_served: served, runners }
    }

    fn free_port() -> u16 {
        TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port()
    }

    /// A runner home in `dir` with a swarm on `port`, known to `hub` by
    /// `token` as runner `id`.
    fn runner(dir: &Path, hub: &FakeHub, id: i64, token: &str, port: u16, site: &str) -> (Home, Swarm) {
        let toml = format!(
            "hub = \"{}\"\ntoken = \"{token}\"\nname = \"r{id}\"\n[swarm]\nlisten = \"127.0.0.1:{port}\"\nsite = \"{site}\"\n",
            hub.url
        );
        std::fs::write(dir.join(crate::home::RUNNER_TOML), toml).unwrap();
        let home = Home::locate(Some(dir)).unwrap();
        let cfg = home.config.runner.clone().unwrap();
        let swarm = Swarm::start(&home, &cfg).unwrap().unwrap();
        hub.runners.lock().unwrap().push((
            token.into(),
            SwarmPeer {
                runner: id,
                name: format!("r{id}"),
                addr: format!("127.0.0.1:{port}"),
                cert: swarm.addr.cert.clone(),
                complete: true,
            },
        ));
        (home, swarm)
    }

    fn image() -> Vec<u8> {
        // 3 MiB of varied bytes, with a run of zeros the swarm never moves.
        let mut x = 0x9e3779b97f4a7c15u64;
        let mut b: Vec<u8> = (0..3 << 20)
            .map(|_| {
                x ^= x << 13;
                x ^= x >> 7;
                x ^= x << 17;
                x as u8
            })
            .collect();
        b[1 << 20..(1 << 20) + (300 << 10)].fill(0);
        b
    }

    #[test]
    fn a_runner_fetches_an_image_from_its_peer_and_not_the_hub() {
        let bytes = image();
        let m = jet::images::scan(&bytes[..], Some("x86_64".into()), |_, _| Ok(())).unwrap();
        let hub = fake_hub(m.clone(), bytes.clone(), "lab");
        let (a_dir, b_dir) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap());
        // Runner A holds the image, imported rather than fetched: its swarm
        // scans it for a manifest before offering it.
        std::fs::create_dir_all(a_dir.path().join("cas")).unwrap();
        std::fs::write(a_dir.path().join("cas").join(format!("{}.ext4", m.hash)), &bytes).unwrap();
        let (_a_home, a) = runner(a_dir.path(), &hub, 1, "ta", free_port(), "lab");
        let deadline = Instant::now() + Duration::from_secs(20);
        while !a.state.offers.lock().unwrap().contains_key(&m.hash) {
            assert!(Instant::now() < deadline, "runner A never offered its image");
            std::thread::sleep(Duration::from_millis(20));
        }
        let (b_home, b) = runner(b_dir.path(), &hub, 2, "tb", free_port(), "lab");
        let moved = b.fetch(&b_home, &m.hash).unwrap().unwrap();
        assert_eq!(std::fs::read(b_home.rootfs(&m.hash)).unwrap(), bytes, "the image came back whole");
        assert_eq!(moved.peer_bytes, moved.bytes, "every chunk came from the peer");
        assert_eq!(hub.chunks_served.load(Ordering::Relaxed), 0, "the hub sent no chunks");
        assert_eq!(moved.chunks, m.needed().len());
        // B now offers the whole image too, read out of its rootfs.
        assert!(matches!(b.state.offers.lock().unwrap().get(&m.hash), Some(Offer { bits: None, .. })));
    }

    #[test]
    fn runners_starting_together_share_what_each_has_fetched() {
        let bytes = image();
        let m = jet::images::scan(&bytes[..], None, |_, _| Ok(())).unwrap();
        let needed = m.needed().len();
        // A slow hub, as a remote one is: the runners have time to swap.
        let hub = slow_hub(m.clone(), bytes.clone(), "lab", Duration::from_millis(400));
        let (b_dir, c_dir) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap());
        let (b_home, b) = runner(b_dir.path(), &hub, 2, "tb", free_port(), "lab");
        let (c_home, c) = runner(c_dir.path(), &hub, 3, "tc", free_port(), "lab");
        let (mb, mc) = std::thread::scope(|s| {
            let fb = s.spawn(|| b.fetch(&b_home, &m.hash).unwrap().unwrap());
            let fc = s.spawn(|| c.fetch(&c_home, &m.hash).unwrap().unwrap());
            (fb.join().unwrap(), fc.join().unwrap())
        });
        assert_eq!(std::fs::read(b_home.rootfs(&m.hash)).unwrap(), bytes);
        assert_eq!(std::fs::read(c_home.rootfs(&m.hash)).unwrap(), bytes);
        let served = hub.chunks_served.load(Ordering::Relaxed);
        assert!(served < 2 * needed, "the hub sent {served} chunks for two runners needing {needed} each");
        assert!(mb.peer_bytes + mc.peer_bytes > 0, "they shared");
    }

    #[test]
    fn without_peers_the_hub_sends_everything_and_a_ticket_from_another_site_is_refused() {
        let bytes = image();
        let m = jet::images::scan(&bytes[..], None, |_, _| Ok(())).unwrap();
        let hub = fake_hub(m.clone(), bytes.clone(), "lab");
        let dir = tempfile::tempdir().unwrap();
        let (home, s) = runner(dir.path(), &hub, 1, "t1", free_port(), "lab");
        let moved = s.fetch(&home, &m.hash).unwrap().unwrap();
        assert_eq!(moved.peer_bytes, 0);
        assert_eq!(hub.chunks_served.load(Ordering::Relaxed), m.needed().len());
        assert_eq!(std::fs::read(home.rootfs(&m.hash)).unwrap(), bytes);

        // A runner in another site shows a ticket for "lab": refused.
        let other = tempfile::tempdir().unwrap();
        let port = free_port();
        let (_h, o) = runner(other.path(), &hub, 3, "t3", port, "elsewhere");
        let ticket = s.state.client.announce(&m.hash, "fetching").unwrap().ticket;
        let err = hello(&format!("127.0.0.1:{port}"), &o.addr.cert, &ticket).unwrap_err();
        assert!(format!("{err:#}").contains("HTTP 403"), "{err:#}");
        // And no session, no chunks.
        let url = format!("https://127.0.0.1:{port}/swarm/v1/chunks/{}", m.chunks[0].hash());
        assert!(matches!(peer_agent(&o.addr.cert).unwrap().get(&url).call(), Err(ureq::Error::Status(401, _))));
        // Nor any answer at all without TLS.
        assert!(ureq::get(&url.replacen("https", "http", 1)).call().is_err());
    }

    #[test]
    fn a_peer_with_another_certificate_is_not_trusted() {
        let bytes = image();
        let m = jet::images::scan(&bytes[..], None, |_, _| Ok(())).unwrap();
        let hub = fake_hub(m.clone(), bytes.clone(), "lab");
        let (a_dir, b_dir) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap());
        std::fs::create_dir_all(a_dir.path().join("cas")).unwrap();
        std::fs::write(a_dir.path().join("cas").join(format!("{}.ext4", m.hash)), &bytes).unwrap();
        let (_a_home, a) = runner(a_dir.path(), &hub, 1, "ta", free_port(), "lab");
        while !a.state.offers.lock().unwrap().contains_key(&m.hash) {
            std::thread::sleep(Duration::from_millis(20));
        }
        let (b_home, b) = runner(b_dir.path(), &hub, 2, "tb", free_port(), "lab");
        // The hub names A, but with B's certificate: whoever answers at A's
        // address is not who the hub vouched for.
        hub.runners.lock().unwrap()[0].1.cert = b.addr.cert.clone();
        let moved = b.fetch(&b_home, &m.hash).unwrap().unwrap();
        assert_eq!(moved.peer_bytes, 0, "nothing was taken from the stranger");
        assert_eq!(hub.chunks_served.load(Ordering::Relaxed), m.needed().len());
        assert_eq!(std::fs::read(b_home.rootfs(&m.hash)).unwrap(), bytes);
        // The certificate is kept: the same runner, the same fingerprint.
        let mut cfg = b_home.config.runner.clone().unwrap();
        cfg.swarm.as_mut().unwrap().listen = "127.0.0.1:0".into();
        let again = Swarm::start(&b_home, &cfg).unwrap().unwrap();
        assert_eq!(again.addr.cert, b.addr.cert);
    }

    #[test]
    fn the_upload_rate_holds_back_a_burst() {
        let dir = tempfile::tempdir().unwrap();
        let (cert, key) = certificate(dir.path()).unwrap();
        let state = State {
            tls: server_tls(cert, key).unwrap(),
            root: PathBuf::new(),
            site: String::new(),
            client: Client::new("http://127.0.0.1:1", "t", None).unwrap(),
            hub: Mutex::new(None),
            offers: Mutex::new(HashMap::new()),
            located: Mutex::new(HashMap::new()),
            sessions: Mutex::new(HashMap::new()),
            conns: AtomicUsize::new(0),
            bucket: Mutex::new((1_000_000.0, Instant::now())),
            rate: 1_000_000.0,
        };
        let start = Instant::now();
        throttle(&state, 1_000_000);
        assert!(start.elapsed() < Duration::from_millis(100), "the first second's worth goes at once");
        throttle(&state, 300_000);
        assert!(start.elapsed() >= Duration::from_millis(250), "then it waits for the rate");
    }
}
`/`	Search
`g` `c`	Code
`g` `h`	History
`g` `b`	Bookmarks
`g` `t`	Trust
`y`	Permalink to this state (files and trees)
`?`	This sheet
Coverage / crates/flow/src/swarm.rs 89.7%

Keyboard
