swarm.rs 89.8%

114 of 127 lines run, 25 of 65 functions · run #517 on main · run missed

//! The swarm's tracker (oxygenhub#144, oxygenflow docs/jetstream.md §8.1):
//! runners that share images tell the hub where peers reach them with each
//! claim; a runner fetching an image announces it and is given the peers
//! in its site that hold or are fetching that image, each with the
//! fingerprint of its TLS certificate, and a ticket, signed with the hub
//! key, that those peers accept.
//!
//! What the tracker knows lives in memory: runners repeat their address
//! with every claim and fetchers announce every `INTERVAL_S`, so a
//! restarted hub knows the swarm again within seconds.

use crate::api::ApiError;
use crate::db;
use crate::Shared;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::Json;
use flow_core::api::{Announce, Announced, SwarmAddr, SwarmPeer};
use jet::swarm::{Ticket, TICKET_LIFE_MS};
use std::collections::HashMap;
use std::sync::Mutex;

/// How often a fetcher announces again.
pub const INTERVAL_S: u64 = 15;
/// A runner whose address is older than this is not offered as a peer.
const ADDR_FRESH_MS: i64 = 2 * 60 * 1000;
/// A fetcher that has not announced for this long is no longer fetching.
const FETCHING_FRESH_MS: i64 = 3 * INTERVAL_S as i64 * 1000;
/// A ticket is made again once it has less than this to run.
const TICKET_RENEW_MS: u64 = 5 * 60 * 1000;

#[derive(Default)]
pub struct Tracker {
    /// Runner id to its swarm address and when it last said so.
    addrs: Mutex<HashMap<i64, (SwarmAddr, i64)>>,
    /// Image to the runners fetching it, with when each last announced.
    fetching: Mutex<HashMap<String, HashMap<i64, i64>>>,
    /// Runner id to its current ticket and when that expires.
    tickets: Mutex<HashMap<i64, (String, u64)>>,
}

impl Tracker {
    /// What a runner's claim said: where peers reach it, or that it is not
    /// in a swarm.
    pub fn claimed(&self, runner: i64, swarm: Option<&SwarmAddr>) {
        let mut addrs = self.addrs.lock().unwrap_or_else(|p| p.into_inner());
        // A peer is trusted by its certificate, so one without is not in
        // the swarm (an older flow).
        match swarm.filter(|s| plausible(&s.addr) && jet::images::is_hash(&s.cert)) {
            Some(s) => {
                addrs.insert(runner, (s.clone(), db::now_ms()));
            }
            None => {
                addrs.remove(&runner);
            }
        }
    }

    /// Where peers reach `runner`, if it is in a swarm now.
    pub fn addr(&self, runner: i64) -> Option<SwarmAddr> {
        let addrs = self.addrs.lock().unwrap_or_else(|p| p.into_inner());
        addrs.get(&runner).filter(|(_, seen)| db::now_ms() - seen < ADDR_FRESH_MS).map(|(a, _)| a.clone())
    }

    fn fetch(&self, image: &str, runner: i64, state: &str) {
        let mut fetching = self.fetching.lock().unwrap_or_else(|p| p.into_inner());
        let now = db::now_ms();
        for who in fetching.values_mut() {
            who.retain(|_, seen| now - *seen < FETCHING_FRESH_MS);
        }
        fetching.retain(|_, who| !who.is_empty());
        if state == "fetching" {
            fetching.entry(image.to_string()).or_default().insert(runner, now);
        } else if let Some(who) = fetching.get_mut(image) {
            who.remove(&runner);
        }
    }

    /// Who is fetching which image now, by image.
    pub fn fetching_now(&self) -> HashMap<String, Vec<i64>> {
        let fetching = self.fetching.lock().unwrap_or_else(|p| p.into_inner());
        let now = db::now_ms();
        fetching
            .iter()
            .map(|(image, who)| {
                (image.clone(), who.iter().filter(|(_, seen)| now - **seen < FETCHING_FRESH_MS).map(|(r, _)| *r).collect::<Vec<_>>())
            })
            .filter(|(_, who)| !who.is_empty())
            .collect()
    }

    fn fetchers(&self, image: &str) -> Vec<i64> {
        let fetching = self.fetching.lock().unwrap_or_else(|p| p.into_inner());
        fetching.get(image).map(|who| who.keys().copied().collect()).unwrap_or_default()
    }
}

/// An address a peer could dial: `host:port`, nothing else.
fn plausible(addr: &str) -> bool {
    addr.len() <= 255 && addr.rsplit_once(':').is_some_and(|(host, port)| !host.is_empty() && port.parse::<u16>().is_ok_and(|p| p > 0))
}

/// `POST /flow/swarm/announce`: a runner starting, continuing or ending a
/// fetch. The peers it may ask, in its site, and its ticket.
pub async fn announce(State(app): State<Shared>, headers: HeaderMap, Json(a): Json<Announce>) -> Result<Json<Announced>, ApiError> {
    let me = crate::flow::runner_auth(&app, &headers).await?;
    if !jet::images::is_hash(&a.image) || !matches!(a.state.as_str(), "fetching" | "done" | "stopped") {
        return Err(ApiError(StatusCode::BAD_REQUEST, "announce an image's hash with state fetching, done or stopped".into()));
    }
    let image = a.image.to_ascii_lowercase();
    let Some(mine) = app.swarm.addr(me.id) else {
        return Err(ApiError(
            StatusCode::CONFLICT,
            "this runner has not said where peers reach it: give it a [runner.swarm] section and let it claim".into(),
        ));
    };
    app.swarm.fetch(&image, me.id, &a.state);
    let fetchers = app.swarm.fetchers(&image);
    let mut peers = Vec::new();
    if a.state == "fetching" {
        for r in db::runners::list(&app.db).await? {
            if r.id == me.id || r.state == "disabled" || !r.online() {
                continue;
            }
            let Some(addr) = app.swarm.addr(r.id).filter(|s| s.site == mine.site) else { continue };
            let complete = r.image_list().is_some_and(|held| held.iter().any(|h| h.eq_ignore_ascii_case(&image)));
            if complete || fetchers.contains(&r.id) {
                peers.push(SwarmPeer { runner: r.id, name: r.name.clone(), addr: addr.addr, cert: addr.cert, complete });
            }
        }
    }
    let ticket = ticket(&app, &me, &mine.site)?;
    // An image only runners hold (oxygenflow#94): the hub has its manifest
    // and none of its chunks, so the fetcher asks peers alone.
    let seed = Some(
        sqlx::query_scalar::<_, String>("SELECT source FROM jet_images WHERE hash = $1")
            .bind(&image)
            .fetch_optional(&app.db)
            .await?
            .is_none_or(|s| s != "peer"),
    );
    Ok(Json(Announced { ticket, peers, interval_s: INTERVAL_S, seed }))
}

/// The runner's ticket, made again when it nears its end or its site
/// changed.
fn ticket(app: &Shared, me: &db::runners::Runner, site: &str) -> Result<String, ApiError> {
    let now = db::now_ms() as u64;
    let mut tickets = app.swarm.tickets.lock().unwrap_or_else(|p| p.into_inner());
    if let Some((t, until)) = tickets.get(&me.id) {
        if *until > now + TICKET_RENEW_MS && Ticket::open(t, &app.hub_key.identity, now).is_ok_and(|t| t.site == site) {
            return Ok(t.clone());
        }
    }
    let (key, secret) = crate::keys::load_secret(&app.config.keys_dir())
        .map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("the hub key: {e}")))?;
    let expires_ms = now + TICKET_LIFE_MS;
    let t = Ticket { hub: app.hub_key.key.to_hex(), runner: me.id, name: me.name.clone(), site: site.to_string(), expires_ms }
        .seal(&key, &secret);
    tickets.insert(me.id, (t.clone(), expires_ms));
    Ok(t)
}

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

    #[test]
    fn only_host_and_port_addresses_are_kept() {
        assert!(plausible("10.0.0.5:7401"));
        assert!(plausible("[::1]:7401"));
        assert!(plausible("runner.lan:7401"));
        assert!(!plausible("10.0.0.5"));
        assert!(!plausible(":7401"));
        assert!(!plausible("10.0.0.5:0"));
        assert!(!plausible("10.0.0.5:99999"));
    }
}

1		//! The swarm's tracker (oxygenhub#144, oxygenflow docs/jetstream.md §8.1):
2		//! runners that share images tell the hub where peers reach them with each
3		//! claim; a runner fetching an image announces it and is given the peers
4		//! in its site that hold or are fetching that image, each with the
5		//! fingerprint of its TLS certificate, and a ticket, signed with the hub
6		//! key, that those peers accept.
7		//!
8		//! What the tracker knows lives in memory: runners repeat their address
9		//! with every claim and fetchers announce every `INTERVAL_S`, so a
10		//! restarted hub knows the swarm again within seconds.
11
12		use crate::api::ApiError;
13		use crate::db;
14		use crate::Shared;
15		use axum::extract::State;
16		use axum::http::{HeaderMap, StatusCode};
17		use axum::Json;
18		use flow_core::api::{Announce, Announced, SwarmAddr, SwarmPeer};
19		use jet::swarm::{Ticket, TICKET_LIFE_MS};
20		use std::collections::HashMap;
21		use std::sync::Mutex;
22
23		/// How often a fetcher announces again.
24		pub const INTERVAL_S: u64 = 15;
25		/// A runner whose address is older than this is not offered as a peer.
26		const ADDR_FRESH_MS: i64 = 2 * 60 * 1000;
27		/// A fetcher that has not announced for this long is no longer fetching.
28		const FETCHING_FRESH_MS: i64 = 3 * INTERVAL_S as i64 * 1000;
29		/// A ticket is made again once it has less than this to run.
30		const TICKET_RENEW_MS: u64 = 5 * 60 * 1000;
31
32		#[derive(Default)]
33		pub struct Tracker {
34		/// Runner id to its swarm address and when it last said so.
35		addrs: Mutex<HashMap<i64, (SwarmAddr, i64)>>,
36		/// Image to the runners fetching it, with when each last announced.
37		fetching: Mutex<HashMap<String, HashMap<i64, i64>>>,
38		/// Runner id to its current ticket and when that expires.
39		tickets: Mutex<HashMap<i64, (String, u64)>>,
40		}
41
42		impl Tracker {
43		/// What a runner's claim said: where peers reach it, or that it is not
44		/// in a swarm.
45	224×	pub fn claimed(&self, runner: i64, swarm: Option<&SwarmAddr>) {
46	224×	let mut addrs = self.addrs.lock().unwrap_or_else(\|p\| p.into_inner());
47		// A peer is trusted by its certificate, so one without is not in
48		// the swarm (an older flow).
49	224×	match swarm.filter(\|s\| plausible(&s.addr) && jet::images::is_hash(&s.cert)) {
50	5×	Some(s) => {
51	5×	addrs.insert(runner, (s.clone(), db::now_ms()));
52	5×	}
53	219×	None => {
54	219×	addrs.remove(&runner);
55	219×	}
56		}
57	224×	}
58
59		/// Where peers reach `runner`, if it is in a swarm now.
60	335×	pub fn addr(&self, runner: i64) -> Option<SwarmAddr> {
61	335×	let addrs = self.addrs.lock().unwrap_or_else(\|p\| p.into_inner());
62	335×	addrs.get(&runner).filter(\|(_, seen)\| db::now_ms() - seen < ADDR_FRESH_MS).map(\|(a, _)\| a.clone())
63	335×	}
64
65	4×	fn fetch(&self, image: &str, runner: i64, state: &str) {
66	4×	let mut fetching = self.fetching.lock().unwrap_or_else(\|p\| p.into_inner());
67	4×	let now = db::now_ms();
68	4×	for who in fetching.values_mut() {
69	4×	who.retain(\|_, seen\| now - *seen < FETCHING_FRESH_MS);
70		}
71	4×	fetching.retain(\|_, who\| !who.is_empty());
72	4×	if state == "fetching" {
73	3×	fetching.entry(image.to_string()).or_default().insert(runner, now);
74	3×	} else if let Some(who) = fetching.get_mut(image) {
75	1×	who.remove(&runner);
76	1×	}
77	4×	}
78
79		/// Who is fetching which image now, by image.
80	2×	pub fn fetching_now(&self) -> HashMap<String, Vec<i64>> {
81	2×	let fetching = self.fetching.lock().unwrap_or_else(\|p\| p.into_inner());
82	2×	let now = db::now_ms();
83	2×	fetching
84	2×	.iter()
85	2×	.map(\|(image, who)\| {
86	0×	(image.clone(), who.iter().filter(\|(_, seen)\| now - *seen < FETCHING_FRESH_MS).map(\|(r, _)\| r).collect::<Vec<_>>())
87	0×	})
88	2×	.filter(\|(_, who)\| !who.is_empty())
89	2×	.collect()
90	2×	}
91
92	4×	fn fetchers(&self, image: &str) -> Vec<i64> {
93	4×	let fetching = self.fetching.lock().unwrap_or_else(\|p\| p.into_inner());
94	4×	fetching.get(image).map(\|who\| who.keys().copied().collect()).unwrap_or_default()
95	4×	}
96		}
97
98		/// An address a peer could dial: `host:port`, nothing else.
99	13×	fn plausible(addr: &str) -> bool {
100	13×	addr.len() <= 255 && addr.rsplit_once(':').is_some_and(\|(host, port)\| !host.is_empty() && port.parse::<u16>().is_ok_and(\|p\| p > 0))
101	13×	}
102
103		/// `POST /flow/swarm/announce`: a runner starting, continuing or ending a
104		/// fetch. The peers it may ask, in its site, and its ticket.
105	8×	pub async fn announce(State(app): State<Shared>, headers: HeaderMap, Json(a): Json<Announce>) -> Result<Json<Announced>, ApiError> {
106	8×	let me = crate::flow::runner_auth(&app, &headers).await?;
107	7×	if !jet::images::is_hash(&a.image) \|\| !matches!(a.state.as_str(), "fetching" \| "done" \| "stopped") {
108	1×	return Err(ApiError(StatusCode::BAD_REQUEST, "announce an image's hash with state fetching, done or stopped".into()));
109	6×	}
110	6×	let image = a.image.to_ascii_lowercase();
111	6×	let Some(mine) = app.swarm.addr(me.id) else {
112	2×	return Err(ApiError(
113	2×	StatusCode::CONFLICT,
114	2×	"this runner has not said where peers reach it: give it a [runner.swarm] section and let it claim".into(),
115	2×	));
116		};
117	4×	app.swarm.fetch(&image, me.id, &a.state);
118	4×	let fetchers = app.swarm.fetchers(&image);
119	4×	let mut peers = Vec::new();
120	4×	if a.state == "fetching" {
121	9×	for r in db::runners::list(&app.db).await? {
122	9×	if r.id == me.id \|\| r.state == "disabled" \|\| !r.online() {
123	3×	continue;
124	6×	}
125	6×	let Some(addr) = app.swarm.addr(r.id).filter(\|s\| s.site == mine.site) else { continue };
126	4×	let complete = r.image_list().is_some_and(\|held\| held.iter().any(\|h\| h.eq_ignore_ascii_case(&image)));
127	4×	if complete \|\| fetchers.contains(&r.id) {
128	3×	peers.push(SwarmPeer { runner: r.id, name: r.name.clone(), addr: addr.addr, cert: addr.cert, complete });
129	3×	}
130		}
131	1×	}
132	4×	let ticket = ticket(&app, &me, &mine.site)?;
133		// An image only runners hold (oxygenflow#94): the hub has its manifest
134		// and none of its chunks, so the fetcher asks peers alone.
135	4×	let seed = Some(
136	4×	sqlx::query_scalar::<_, String>("SELECT source FROM jet_images WHERE hash = $1")
137	4×	.bind(&image)
138	4×	.fetch_optional(&app.db)
139	4×	.await?
140	4×	.is_none_or(\|s\| s != "peer"),
141		);
142	4×	Ok(Json(Announced { ticket, peers, interval_s: INTERVAL_S, seed }))
143	8×	}
144
145		/// The runner's ticket, made again when it nears its end or its site
146		/// changed.
147	4×	fn ticket(app: &Shared, me: &db::runners::Runner, site: &str) -> Result<String, ApiError> {
148	4×	let now = db::now_ms() as u64;
149	4×	let mut tickets = app.swarm.tickets.lock().unwrap_or_else(\|p\| p.into_inner());
150	4×	if let Some((t, until)) = tickets.get(&me.id) {
151	2×	if *until > now + TICKET_RENEW_MS && Ticket::open(t, &app.hub_key.identity, now).is_ok_and(\|t\| t.site == site) {
152	1×	return Ok(t.clone());
153	1×	}
154	2×	}
155	3×	let (key, secret) = crate::keys::load_secret(&app.config.keys_dir())
156	3×	.map_err(\|e\| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("the hub key: {e}")))?;
157	3×	let expires_ms = now + TICKET_LIFE_MS;
158	3×	let t = Ticket { hub: app.hub_key.key.to_hex(), runner: me.id, name: me.name.clone(), site: site.to_string(), expires_ms }
159	3×	.seal(&key, &secret);
160	3×	tickets.insert(me.id, (t.clone(), expires_ms));
161	3×	Ok(t)
162	4×	}
163
164		#[cfg(test)]
165		mod tests {
166		use super::*;
167
168		#[test]
169	1×	fn only_host_and_port_addresses_are_kept() {
170	1×	assert!(plausible("10.0.0.5:7401"));
171	1×	assert!(plausible("[::1]:7401"));
172	1×	assert!(plausible("runner.lan:7401"));
173	1×	assert!(!plausible("10.0.0.5"));
174	1×	assert!(!plausible(":7401"));
175	1×	assert!(!plausible("10.0.0.5:0"));
176	1×	assert!(!plausible("10.0.0.5:99999"));
177	1×	}
178		}

`/`	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/hub/src/swarm.rs 89.8%

Keyboard