use super::auto_connect_impl::BrowserDesiredPeer;
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct TransferPeerDemandRequest {
pub id: String,
pub device_id: String,
pub purpose: TransferPeerDemandPurpose,
}
impl TransferPeerDemandRequest {
pub fn validate(&self) -> anyhow::Result<()> {
for value in [&self.id, &self.device_id] {
Self::validate_identity(value)?;
}
Ok(())
}
pub fn validate_identity(value: &str) -> anyhow::Result<()> {
if value.is_empty()
|| value.len() > 256
|| value.trim() != value
|| value.chars().any(char::is_control)
{
anyhow::bail!("invalid transfer peer demand identity");
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum TransferPeerDemandPurpose {
FileTransfer,
DriveAccess,
ContentSync,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum TransferPeerDemandState {
WaitingForPeer,
Queued,
Admitted,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TransferPeerDemandStatus {
pub id: String,
pub device_id: String,
pub purpose: TransferPeerDemandPurpose,
pub generation: u64,
pub state: TransferPeerDemandState,
}
#[derive(Debug)]
pub(super) struct TransferPeerDemands {
budget: usize,
requests: BTreeMap<String, TransferPeerDemandRequest>,
order: VecDeque<String>,
active: HashSet<String>,
available: HashSet<String>,
}
impl Default for TransferPeerDemands {
fn default() -> Self {
Self {
budget: Self::platform_budget(),
requests: BTreeMap::new(),
order: VecDeque::new(),
active: HashSet::new(),
available: HashSet::new(),
}
}
}
impl TransferPeerDemands {
pub(super) fn platform_budget() -> usize {
if cfg!(any(
target_arch = "wasm32",
target_os = "ios",
target_os = "android"
)) {
8
} else {
32
}
}
pub(super) fn set_budget(&mut self, budget: usize) -> anyhow::Result<()> {
if budget == 0 || budget > Self::platform_budget() {
anyhow::bail!(
"invalid transfer peer budget: choose 1..={}",
Self::platform_budget()
);
}
if self.active.len() > budget {
anyhow::bail!("receiver_busy: release active peer demands before lowering the budget");
}
self.budget = budget;
Ok(())
}
pub(super) fn request(&mut self, request: TransferPeerDemandRequest) -> anyhow::Result<()> {
request.validate()?;
if let Some(current) = self.requests.get(&request.id) {
if current != &request {
anyhow::bail!("transfer peer demand identity conflicts with an accepted request");
}
return Ok(());
}
if self.requests.len() >= 4096 {
anyhow::bail!("receiver_busy: transfer peer demand input is full");
}
if !self.order.contains(&request.device_id) {
if self.order.len() >= 500 {
anyhow::bail!("receiver_busy: transfer peer demand directory is full");
}
self.order.push_back(request.device_id.clone());
}
self.requests.insert(request.id.clone(), request);
Ok(())
}
pub(super) fn release(&mut self, id: &str) -> bool {
let Some(request) = self.requests.remove(id) else {
return false;
};
if !self
.requests
.values()
.any(|other| other.device_id == request.device_id)
{
self.order.retain(|device| device != &request.device_id);
self.active.remove(&request.device_id);
}
true
}
pub(super) fn devices_with_requests(&self) -> Vec<String> {
self.order.iter().cloned().collect()
}
pub(super) fn has_requests_for(&self, device: &str) -> bool {
self.order.iter().any(|id| id == device)
}
pub(super) fn active(&self, device: &str) -> bool {
self.active.contains(device)
}
pub(super) fn yield_failed(&mut self, device: &str) {
if self.active.remove(device) {
self.order.retain(|id| id != device);
self.order.push_back(device.to_owned());
}
}
pub(super) fn reconcile(
&mut self,
peers: &BTreeMap<String, BrowserDesiredPeer>,
cooldown: &HashMap<String, i64>,
local_device: &str,
now: i64,
is_excluded: impl Fn(&str) -> bool,
) -> Vec<String> {
let previously_active = self.active.clone();
self.available = peers
.values()
.filter(|peer| {
peer.online
&& peer.device_id != local_device
&& !is_excluded(&peer.device_id)
&& !peer
.excluded_peers
.iter()
.any(|device| device.trim().eq_ignore_ascii_case(local_device.trim()))
&& peer.ticket.as_deref().is_some_and(|ticket| {
let (endpoint, suffix) = crate::session_token::split_ticket(ticket);
let Ok(address) = super::parse_endpoint_ticket(endpoint) else {
return false;
};
peer.node_id
.as_deref()
.is_none_or(|claimed| claimed == address.id.to_string())
&& suffix.is_none_or(|suffix| {
crate::session_token::decode_payload(endpoint, suffix).is_some()
})
})
})
.map(|peer| peer.device_id.clone())
.collect();
let ready = |device: &String| {
self.available.contains(device)
&& !cooldown.get(device).is_some_and(|deadline| *deadline > now)
};
self.active
.retain(|device| ready(device) && self.order.contains(device));
for device in &self.order {
if self.active.len() >= self.budget {
break;
}
if ready(device) {
self.active.insert(device.clone());
}
}
self.active
.difference(&previously_active)
.cloned()
.collect()
}
pub(super) fn status(&self, id: &str, generation: u64) -> Option<TransferPeerDemandStatus> {
self.requests
.get(id)
.map(|request| TransferPeerDemandStatus {
id: request.id.clone(),
device_id: request.device_id.clone(),
purpose: request.purpose,
generation,
state: if self.active(&request.device_id) {
TransferPeerDemandState::Admitted
} else if self.available.contains(&request.device_id) {
TransferPeerDemandState::Queued
} else {
TransferPeerDemandState::WaitingForPeer
},
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn invalid_or_excluded_directory_rows_do_not_hold_a_bulk_peer_slot() {
for kind in [
"node-mismatch",
"remote-exclusion",
"local-exclusion",
"malformed-token",
] {
let node = iroh::SecretKey::generate().public();
let ticket = iroh_tickets::endpoint::EndpointTicket::new(iroh::EndpointAddr::new(node))
.to_string();
let blocked: BrowserDesiredPeer = serde_json::from_value(serde_json::json!({
"deviceId":"blocked", "ticket":if kind == "malformed-token" { format!("{ticket}.invalid") } else { ticket.clone() }, "online":true, "connectionIntent":"observe",
"nodeId":if kind == "node-mismatch" { iroh::SecretKey::generate().public().to_string() } else { node.to_string() },
"excludedPeers":if kind == "remote-exclusion" { vec![" LOCAL "] } else { vec![] }
})).unwrap();
let ready: BrowserDesiredPeer = serde_json::from_value(serde_json::json!({
"deviceId":"ready", "ticket":ticket, "online":true, "connectionIntent":"observe"
}))
.unwrap();
let mut demands = TransferPeerDemands::default();
demands.set_budget(1).unwrap();
demands.request(request("blocked-job", "blocked")).unwrap();
demands.request(request("ready-job", "ready")).unwrap();
demands.reconcile(
&BTreeMap::from([("blocked".into(), blocked), ("ready".into(), ready)]),
&HashMap::new(),
"local",
100,
|id| kind == "local-exclusion" && id == "blocked",
);
assert_eq!(
demands.status("blocked-job", 1).unwrap().state,
TransferPeerDemandState::WaitingForPeer,
"{kind}"
);
assert_eq!(
demands.status("ready-job", 1).unwrap().state,
TransferPeerDemandState::Admitted,
"{kind}"
);
}
}
fn request(id: &str, device: &str) -> TransferPeerDemandRequest {
TransferPeerDemandRequest {
id: id.into(),
device_id: device.into(),
purpose: TransferPeerDemandPurpose::FileTransfer,
}
}
#[test]
fn transfer_peer_budget_is_fair_per_peer_and_counts_shared_demands_once() {
for budget in [1, 8, 32]
.into_iter()
.filter(|budget| *budget <= TransferPeerDemands::platform_budget())
{
let mut demands = TransferPeerDemands::default();
demands.set_budget(budget).unwrap();
let mut peers = BTreeMap::new();
for index in 0..=budget + 2 {
let device = format!("device-{index}");
let node = iroh::SecretKey::generate().public();
let ticket =
iroh_tickets::endpoint::EndpointTicket::new(iroh::EndpointAddr::new(node))
.to_string();
let peer: BrowserDesiredPeer = serde_json::from_value(serde_json::json!({
"deviceId":device, "ticket":ticket, "online":index != 0, "connectionIntent":"observe"
})).unwrap();
peers.insert(device.clone(), peer);
demands
.request(request(&format!("job-{index}"), &device))
.unwrap();
}
demands
.request(request("second-job-for-device-1", "device-1"))
.unwrap();
demands.reconcile(&peers, &HashMap::new(), "local", 100, |_| false);
assert_eq!(
demands.status("job-0", 7).unwrap().state,
TransferPeerDemandState::WaitingForPeer
);
for index in 1..=budget {
assert_eq!(
demands.status(&format!("job-{index}"), 7).unwrap().state,
TransferPeerDemandState::Admitted
);
}
assert_eq!(
demands
.status(&format!("job-{}", budget + 1), 7)
.unwrap()
.state,
TransferPeerDemandState::Queued
);
assert!(demands
.request(request("job-1", "different-device"))
.is_err());
assert!(demands.release("job-1"));
demands.reconcile(&peers, &HashMap::new(), "local", 100, |_| false);
assert_eq!(
demands.status("second-job-for-device-1", 7).unwrap().state,
TransferPeerDemandState::Admitted
);
assert_eq!(
demands
.status(&format!("job-{}", budget + 1), 7)
.unwrap()
.state,
TransferPeerDemandState::Queued
);
assert!(demands.release("second-job-for-device-1"));
demands.reconcile(&peers, &HashMap::new(), "local", 100, |_| false);
assert_eq!(
demands
.status(&format!("job-{}", budget + 1), 7)
.unwrap()
.state,
TransferPeerDemandState::Admitted
);
demands.yield_failed(&format!("device-{}", budget + 1));
let cooldown = HashMap::from([(format!("device-{}", budget + 1), 200)]);
demands.reconcile(&peers, &cooldown, "local", 100, |_| false);
assert_eq!(
demands
.status(&format!("job-{}", budget + 2), 7)
.unwrap()
.state,
TransferPeerDemandState::Admitted
);
}
}
}