use std::net::TcpStream;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use super::worker::{HardwareInfo, WorkerPricing, WorkerRegistration, WorkerResources};
use super::BrokerState;
use colored::Colorize;
use serde::{Deserialize, Serialize};
fn sanitize_price(raw: Option<f64>, default: f64, max: f64) -> f64 {
match raw {
Some(v) if v.is_finite() && v >= 0.0 && v <= max => v,
_ => default,
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum DiscoveryMode {
WireGuard { subnet: String },
Local,
}
#[derive(Debug, Clone)]
pub struct DiscoveryConfig {
pub subnet: String,
pub worker_port: u16,
pub extra_ports: Vec<u16>,
pub scan_port_range: Option<(u16, u16)>,
pub interval_secs: u64,
pub enable_scan: bool,
pub enable_dns: bool,
pub peers: Vec<String>,
pub local_workers: Vec<(String, u16)>,
}
impl Default for DiscoveryConfig {
fn default() -> Self {
let peers = std::env::var("ZAKURO_PEERS")
.unwrap_or_default()
.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
let worker_port = std::env::var("ZAKURO_WORKER_PORT")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(3960);
let scan_port_range = std::env::var("ZAKURO_SCAN_RANGE").ok().and_then(|v| {
let parts: Vec<&str> = v.splitn(2, '-').collect();
if parts.len() == 2 {
let start = parts[0].parse::<u16>().ok()?;
let end = parts[1].parse::<u16>().ok()?;
Some((start, end))
} else {
None
}
});
let local_workers = parse_local_workers(
&std::env::var("ZAKURO_WORKERS").unwrap_or_default(),
worker_port,
);
Self {
subnet: "10.13.13".to_string(),
worker_port,
extra_ports: vec![3961, 3962], scan_port_range,
interval_secs: std::env::var("ZAKURO_SCAN_INTERVAL")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(15), enable_scan: true,
enable_dns: true,
peers,
local_workers,
}
}
}
pub struct Discovery {
config: DiscoveryConfig,
state: Arc<BrokerState>,
mode: DiscoveryMode,
}
impl Discovery {
pub fn new(config: DiscoveryConfig, state: Arc<BrokerState>) -> Self {
let mode = detect_discovery_mode(&config.subnet);
Self {
config,
state,
mode,
}
}
pub fn mode(&self) -> &DiscoveryMode {
&self.mode
}
pub fn run(&self, verbose: bool) {
if verbose {
match &self.mode {
DiscoveryMode::WireGuard { subnet } => {
println!(
" {} WireGuard mode (subnet: {}.0/24)",
"[DISCOVERY]".cyan(),
subnet
);
}
DiscoveryMode::Local => {
println!(
" {} Local mode (scanning localhost)",
"[DISCOVERY]".yellow()
);
}
}
if !self.config.peers.is_empty() {
println!(
" {} Peers: {}",
"[DISCOVERY]".cyan(),
self.config.peers.join(", ")
);
}
}
if self.config.enable_scan {
self.discover_workers(verbose);
}
loop {
thread::sleep(Duration::from_secs(self.config.interval_secs));
if self.config.enable_scan {
self.discover_workers(verbose);
}
}
}
fn discover_workers(&self, verbose: bool) {
self.scan_localhost(verbose);
for (host, port) in self.config.local_workers.clone() {
self.try_register_worker(&host, port, verbose);
}
self.scan_peers(verbose);
}
fn scan_peers(&self, verbose: bool) {
for peer in &self.config.peers {
let (host, port) = if let Some((h, p)) = peer.rsplit_once(':') {
(h.to_string(), p.parse().unwrap_or(self.config.worker_port))
} else {
(peer.clone(), self.config.worker_port)
};
if host.is_empty() {
continue;
}
let broker_port = self.state.config.port;
if port == broker_port {
self.try_register_peer_broker(&host, broker_port, verbose);
self.fetch_workers_from_peer_broker(&host, broker_port, verbose);
} else {
self.try_register_worker(&host, port, verbose);
self.try_register_peer_broker(&host, broker_port, verbose);
self.fetch_workers_from_peer_broker(&host, broker_port, verbose);
}
}
}
fn fetch_workers_from_peer_broker(&self, host: &str, broker_port: u16, verbose: bool) {
let peer_key = self.state.peer_manager.peer_key();
let url = format!("http://{}:{}/peer/workers", host, broker_port);
let mut cfg = ureq::Agent::config_builder().timeout_global(Some(Duration::from_secs(3)));
if crate::vpn::is_mesh_ip(host) {
if let Some(addr) = crate::vpn::mesh_proxy_addr() {
if let Ok(proxy) = ureq::Proxy::new(&format!("http://{}", addr)) {
cfg = cfg.proxy(Some(proxy));
}
}
}
let agent = ureq::Agent::new_with_config(cfg.build());
let result = agent.get(&url).header("X-Peer-Key", peer_key).call();
let body = match result {
Ok(resp) => match resp.into_body().read_to_string() {
Ok(s) => s,
Err(_) => return,
},
Err(_) => return,
};
let json: serde_json::Value = match serde_json::from_str(&body) {
Ok(v) => v,
Err(_) => return,
};
let workers = match json["workers"].as_array() {
Some(arr) => arr,
None => return,
};
for w in workers {
let name = w["name"].as_str().unwrap_or("").to_string();
let uri = w["uri"].as_str().unwrap_or("").to_string();
let worker_type = w["worker_type"].as_str().unwrap_or("zakuro").to_string();
if name.is_empty() || uri.is_empty() {
continue;
}
let node_field = w["node"].as_str().unwrap_or("");
let source_node = {
let bare = strip_zc_node(node_field);
if bare.is_empty() {
format!("node-unknown-{}", host)
} else {
bare.to_string()
}
};
let rewritten_uri = if let Some(rest) = uri.strip_prefix("http://127.0.0.1:") {
format!(
"http://{}:{}",
host,
rest.split('/').next().unwrap_or("3960")
)
} else if let Some(rest) = uri.strip_prefix("http://localhost:") {
format!(
"http://{}:{}",
host,
rest.split('/').next().unwrap_or("3960")
)
} else {
uri.clone() };
let existing = self.state.workers.list();
if let Some(known) = existing
.iter()
.find(|e| is_same_peer_worker(e, &name, &rewritten_uri, &source_node))
{
self.state.workers.refresh_heartbeat(&known.id);
continue;
}
let (provider_type, served_models, price_per_mtok) = provider_fields_from_json(w);
let registration = WorkerRegistration {
name: name.clone(),
uri: rewritten_uri.clone(),
worker_type,
resources: WorkerResources {
cpus_available: w["cpus_available"].as_f64().unwrap_or(1.0),
cpus_total: w["cpus_total"].as_f64().unwrap_or(1.0),
memory_available: (w["memory_available_gib"].as_f64().unwrap_or(1.0)
* 1024.0
* 1024.0
* 1024.0) as u64,
memory_total: (w["memory_total_gib"].as_f64().unwrap_or(1.0)
* 1024.0
* 1024.0
* 1024.0) as u64,
gpus_available: w["gpus_available"].as_u64().unwrap_or(0) as u32,
gpus_total: w["gpus_total"].as_u64().unwrap_or(0) as u32,
},
pricing: WorkerPricing {
price_per_hour: sanitize_price(w["price_per_hour"].as_f64(), 3.6, 1000.0),
min_charge: sanitize_price(w["min_charge"].as_f64(), 0.001, 100.0),
},
tags: vec![],
max_timeout_secs: 0.0,
hardware: HardwareInfo {
cpu_model: w["cpu_model"].as_str().map(|s| s.to_string()),
gpu_model: w["gpu_model"].as_str().map(|s| s.to_string()),
gpu_vram_gb: w["gpu_vram_gb"].as_u64().map(|v| v as u32),
storage_gb: w["storage_gb"].as_u64().map(|v| v as u32),
},
wireguard_ip: Some(host.to_string()),
is_docker: w["is_docker"].as_bool(),
source_node: Some(source_node.clone()),
explicit_local: false,
provider_type,
served_models,
price_per_mtok,
};
let worker = self.state.workers.register(registration);
if let Some(fp) = peer_node_fp(&worker) {
self.state.workers.set_node_fp(&worker.id, &fp);
}
if verbose {
println!(
" {} Discovered peer worker {} at {} (via broker {}:{})",
"[DISCOVERY]".cyan(),
worker.name,
rewritten_uri,
host,
broker_port,
);
}
}
}
fn try_register_peer_broker(&self, host: &str, port: u16, verbose: bool) {
let addr = format!("{}:{}", host, port);
use std::net::ToSocketAddrs;
let sock_addr: std::net::SocketAddr = match addr.to_socket_addrs() {
Ok(mut addrs) => match addrs.next() {
Some(a) => a,
None => return,
},
Err(_) => return,
};
if TcpStream::connect_timeout(&sock_addr, Duration::from_millis(200)).is_err() {
return;
}
let base_url = format!("http://{}:{}", host, port);
self.state.peer_manager.register_peer(base_url.clone());
if let Some(client) = self.state.peer_manager.get_client(&base_url) {
if client.check_health() && verbose {
println!(
" {} Peer broker alive at {}:{}",
"[DISCOVERY]".cyan(),
host,
port
);
}
}
}
fn scan_localhost(&self, verbose: bool) {
let ports: Vec<u16> = if let Some((start, end)) = self.config.scan_port_range {
(start..=end).collect()
} else {
let mut p = vec![self.config.worker_port];
p.extend(&self.config.extra_ports);
p
};
let open_ports: Vec<u16> = {
let (tx, rx) = std::sync::mpsc::channel();
let mut handles = Vec::with_capacity(ports.len());
for port in &ports {
let port = *port;
let tx = tx.clone();
handles.push(thread::spawn(move || {
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap();
if TcpStream::connect_timeout(&addr, Duration::from_millis(50)).is_ok() {
let _ = tx.send(port);
}
}));
}
drop(tx);
let mut open: Vec<u16> = rx.iter().collect();
for h in handles {
let _ = h.join();
}
open.sort_unstable();
open
};
let state = Arc::clone(&self.state);
let config = self.config.clone();
let mut probe_handles = Vec::with_capacity(open_ports.len());
for port in open_ports {
let state = Arc::clone(&state);
let config = config.clone();
probe_handles.push(thread::spawn(move || {
let tmp = Discovery {
config,
state: Arc::clone(&state),
mode: DiscoveryMode::Local,
};
tmp.try_register_worker("127.0.0.1", port, verbose);
}));
}
for h in probe_handles {
let _ = h.join();
}
}
fn try_register_worker(&self, host: &str, port: u16, verbose: bool) {
if !is_loopback_host(host)
&& !is_explicit_local_worker(&self.config.local_workers, host, port)
{
return;
}
let addr = format!("{}:{}", host, port);
use std::net::ToSocketAddrs;
let sock_addr = match addr.to_socket_addrs() {
Ok(mut addrs) => match addrs.next() {
Some(a) => a,
None => return,
},
Err(_) => return,
};
if TcpStream::connect_timeout(&sock_addr, Duration::from_millis(1500)).is_err() {
return;
}
let uri = format!("http://{}:{}", host, port);
let existing = self.state.workers.list();
let existing_worker = existing.iter().find(|w| {
w.uri == uri
|| (host == "127.0.0.1" && w.uri.contains("localhost"))
|| (host == "localhost" && w.uri.contains("127.0.0.1"))
});
if let Some(worker) = existing_worker {
if let Some(info) = self.probe_worker(host, port) {
if let Some(resources) = info.resources {
self.state.workers.update_resources(
&worker.id,
resources,
info.hardware.unwrap_or_default(),
);
} else {
self.state.workers.refresh_heartbeat(&worker.id);
}
} else {
self.state.workers.refresh_heartbeat(&worker.id);
}
return;
}
if let Some(worker_info) = self.probe_worker(host, port) {
{
let registration = WorkerRegistration {
name: worker_info
.name
.unwrap_or_else(|| format!("worker-{}", host)),
uri,
worker_type: worker_info
.worker_type
.unwrap_or_else(|| "zakuro".to_string()),
resources: worker_info.resources.unwrap_or_default(),
pricing: worker_info.pricing.unwrap_or_default(),
tags: worker_info.tags.unwrap_or_default(),
max_timeout_secs: 0.0,
hardware: worker_info.hardware.unwrap_or_default(),
wireguard_ip: None,
is_docker: None,
source_node: None,
explicit_local: is_explicit_local_worker(
&self.config.local_workers,
host,
port,
),
provider_type: Default::default(),
served_models: vec![],
price_per_mtok: 0.0,
};
let worker = self.state.workers.register(registration);
let worker = self
.state
.workers
.set_node_fp(&worker.id, &self.state.node_key.fingerprint())
.unwrap_or(worker);
if verbose {
println!(
" {} Discovered worker {} at {}",
"[DISCOVERY]".cyan(),
worker.name,
worker.uri
);
}
if let Some(ref owner_id) = self.state.config.owner_user_id {
let node_name = self.state.config.node_name.as_deref();
let node_pubkey = self.state.node_key.public_b64();
if let (Some(ref api_url), Some(ref api_key)) =
(&self.state.config.api_url, &self.state.config.api_key)
{
match crate::broker::ledger::Ledger::sync_workers_via_api(
owner_id,
std::slice::from_ref(&worker),
api_url,
api_key,
node_name,
self.state.own_wireguard_ip.as_deref(),
Some(node_pubkey.as_str()),
) {
Ok(()) => {
if verbose {
println!(
" [WORKER_SYNC] Worker {} synced to dashboard",
worker.name
);
}
}
Err(e) => {
eprintln!(" [WORKER_SYNC] Failed to sync {}: {}", worker.name, e);
}
}
}
}
}
}
}
fn probe_worker(&self, ip: &str, port: u16) -> Option<WorkerProbeResult> {
let health_url = format!("http://{}:{}/health", ip, port);
match ureq::get(&health_url)
.config()
.timeout_global(Some(Duration::from_secs(2)))
.build()
.call()
{
Ok(response) if response.status().as_u16() == 200 => {
let info_url = format!("http://{}:{}/info", ip, port);
match ureq::get(&info_url)
.config()
.timeout_global(Some(Duration::from_secs(2)))
.build()
.call()
{
Ok(info_response) if info_response.status().as_u16() == 200 => {
if let Ok(body) = info_response.into_body().read_to_string() {
if let Ok(info) = serde_json::from_str::<WorkerProbeResult>(&body) {
if info.worker_type.is_some() {
return Some(info);
}
}
}
None
}
_ => None, }
}
_ => None,
}
}
}
#[derive(Debug, Clone, serde::Deserialize)]
struct GossipEntry {
fp: String,
url: String,
epoch: u64,
}
#[derive(Debug, Clone, serde::Deserialize)]
struct GossipResponse {
peers: Vec<GossipEntry>,
}
pub fn fetch_gossip(
base_url: &str,
roster: &super::roster_cache::RosterCache,
agent: &ureq::Agent,
node_key: &super::node_identity::NodeKey,
) -> Vec<(String, String, u64)> {
let url = format!("{}/peer/peers", base_url);
let mut rb = agent.get(&url);
for (k, v) in node_key.sign_headers("GET", "/peer/peers", b"") {
rb = rb.header(k.as_str(), v.as_str());
}
let resp = match rb.call() {
Ok(r) => r,
Err(_) => return Vec::new(),
};
if resp.status().as_u16() != 200 {
return Vec::new();
}
let node_id = resp
.headers()
.get("X-Node-Id")
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
let sig = resp
.headers()
.get("X-Node-Sig")
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
let (node_id, sig) = match (node_id, sig) {
(Some(n), Some(s)) => (n, s),
_ => return Vec::new(),
};
let body = match resp.into_body().read_to_string() {
Ok(s) => s.into_bytes(),
Err(_) => return Vec::new(),
};
verify_gossip_body(&body, &node_id, &sig, roster)
}
fn verify_gossip_body(
body: &[u8],
node_id: &str,
sig: &str,
roster: &super::roster_cache::RosterCache,
) -> Vec<(String, String, u64)> {
if !super::node_identity::verify_sig(node_id, body, sig) {
return Vec::new();
}
let fp = match super::node_identity::fingerprint_of_pubkey_b64(node_id) {
Some(fp) => fp,
None => return Vec::new(),
};
if !roster.fingerprint_authorized(&fp) {
return Vec::new();
}
let parsed: GossipResponse = match serde_json::from_slice(body) {
Ok(p) => p,
Err(_) => return Vec::new(),
};
parsed
.peers
.into_iter()
.map(|e| (e.fp, e.url, e.epoch))
.collect()
}
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Advert {
pub fp: String,
pub price_per_hour: f64,
pub resources: super::worker::BrokerResources,
pub epoch: u64,
}
pub fn fetch_advert(
base_url: &str,
roster: &super::roster_cache::RosterCache,
agent: &ureq::Agent,
node_key: &super::node_identity::NodeKey,
) -> Option<Advert> {
let url = format!("{}/peer/advert", base_url);
let mut rb = agent.get(&url);
for (k, v) in node_key.sign_headers("GET", "/peer/advert", b"") {
rb = rb.header(k.as_str(), v.as_str());
}
let resp = rb.call().ok()?;
if resp.status().as_u16() != 200 {
return None;
}
let node_id = resp
.headers()
.get("X-Node-Id")
.and_then(|v| v.to_str().ok())?
.to_string();
let sig = resp
.headers()
.get("X-Node-Sig")
.and_then(|v| v.to_str().ok())?
.to_string();
let body = resp.into_body().read_to_string().ok()?.into_bytes();
verify_advert_body(&body, &node_id, &sig, roster)
}
fn verify_advert_body(
body: &[u8],
node_id: &str,
sig: &str,
roster: &super::roster_cache::RosterCache,
) -> Option<Advert> {
if !super::node_identity::verify_sig(node_id, body, sig) {
return None;
}
let signer_fp = super::node_identity::fingerprint_of_pubkey_b64(node_id)?;
if !roster.fingerprint_authorized(&signer_fp) {
return None;
}
let advert: Advert = serde_json::from_slice(body).ok()?;
if advert.fp != signer_fp {
return None;
}
Some(advert)
}
pub fn accept_peer(
fp: &str,
url: &str,
roster: &super::roster_cache::RosterCache,
peer_key: &str,
) -> bool {
accept_peer_with(fp, url, roster, |u| fetch_signed_health(u, peer_key))
}
pub fn accept_peer_with<F>(
fp: &str,
url: &str,
roster: &super::roster_cache::RosterCache,
health_fetch: F,
) -> bool
where
F: Fn(&str) -> Option<(String, String, Vec<u8>)>,
{
if !roster.fingerprint_authorized(fp) {
return false;
}
let (node_id, sig, body) = match health_fetch(url) {
Some(v) => v,
None => return false,
};
if !super::node_identity::verify_sig(&node_id, &body, &sig) {
return false;
}
match super::node_identity::fingerprint_of_pubkey_b64(&node_id) {
Some(responder_fp) => responder_fp == fp,
None => false,
}
}
fn fetch_signed_health(url: &str, peer_key: &str) -> Option<(String, String, Vec<u8>)> {
let health_url = format!("{}/peer/health", url);
let host = url
.split("://")
.nth(1)
.and_then(|r| r.split([':', '/']).next())
.unwrap_or("");
let mut cfg = ureq::Agent::config_builder().timeout_global(Some(Duration::from_secs(3)));
if crate::vpn::is_mesh_ip(host) {
if let Some(addr) = crate::vpn::mesh_proxy_addr() {
if let Ok(proxy) = ureq::Proxy::new(&format!("http://{}", addr)) {
cfg = cfg.proxy(Some(proxy));
}
}
}
let agent = ureq::Agent::new_with_config(cfg.build());
let mut rb = agent.get(&health_url);
if !peer_key.is_empty() {
rb = rb.header("X-Peer-Key", peer_key);
}
let resp = rb.call().ok()?;
if resp.status().as_u16() != 200 {
return None;
}
let node_id = resp
.headers()
.get("X-Node-Id")
.and_then(|v| v.to_str().ok())?
.to_string();
let sig = resp
.headers()
.get("X-Node-Sig")
.and_then(|v| v.to_str().ok())?
.to_string();
let body = resp.into_body().read_to_string().ok()?.into_bytes();
Some((node_id, sig, body))
}
fn admit_verified_peer_with<F>(
state: &Arc<BrokerState>,
fp: &str,
url: &str,
roster: &super::roster_cache::RosterCache,
health_fetch: F,
) -> bool
where
F: Fn(&str) -> Option<(String, String, Vec<u8>)>,
{
if !accept_peer_with(fp, url, roster, health_fetch) {
return false;
}
state.peer_manager.register_peer(url.to_string());
state.peer_manager.set_peer_fingerprint(url, fp);
true
}
pub fn run_discovery_round(state: &Arc<BrokerState>, roster: &super::roster_cache::RosterCache) {
let agent = crate::vpn::mesh_agent(Duration::from_secs(3));
let peer_key = state.peer_manager.peer_key();
for peer_url in state.peer_manager.peer_urls() {
for (fp, url, _epoch) in fetch_gossip(&peer_url, roster, &agent, &state.node_key) {
admit_verified_peer_with(state, &fp, &url, roster, |u| {
fetch_signed_health(u, peer_key)
});
}
}
if state.peer_manager.peer_urls().is_empty() {
if let Some(self_ip) = get_mesh_ip() {
for url in discover_broker_peers_on_mesh_subnet(&self_ip, peer_key) {
if let Some((node_id, sig, body)) = fetch_signed_health(&url, peer_key) {
if super::node_identity::verify_sig(&node_id, &body, &sig) {
if let Some(fp) = super::node_identity::fingerprint_of_pubkey_b64(&node_id)
{
let fetched = (node_id.clone(), sig.clone(), body.clone());
admit_verified_peer_with(state, &fp, &url, roster, move |_| {
Some(fetched.clone())
});
}
}
}
}
}
}
state.peer_manager.persist();
refresh_adverts(state, roster);
}
pub(crate) fn refresh_adverts(state: &Arc<BrokerState>, roster: &super::roster_cache::RosterCache) {
let agent = crate::vpn::mesh_agent(Duration::from_secs(3));
for peer_url in state.peer_manager.peer_urls() {
if let Some(advert) = fetch_advert(&peer_url, roster, &agent, &state.node_key) {
state.peer_manager.set_peer_advert(&peer_url, advert);
}
}
}
#[derive(Debug, Clone, serde::Deserialize)]
struct WorkerProbeResult {
name: Option<String>,
worker_type: Option<String>,
resources: Option<WorkerResources>,
pricing: Option<WorkerPricing>,
tags: Option<Vec<String>>,
#[serde(default)]
hardware: Option<HardwareInfo>,
}
trait ColorExt {
fn cyan(&self) -> String;
fn yellow(&self) -> String;
fn green(&self) -> String;
}
impl ColorExt for &str {
fn cyan(&self) -> String {
format!("\x1b[36m{}\x1b[0m", self)
}
fn yellow(&self) -> String {
format!("\x1b[33m{}\x1b[0m", self)
}
fn green(&self) -> String {
format!("\x1b[32m{}\x1b[0m", self)
}
}
pub(crate) fn parse_local_workers(raw: &str, default_port: u16) -> Vec<(String, u16)> {
raw.split(',')
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.filter_map(|entry| match entry.rsplit_once(':') {
Some((host, port)) => {
let host = host.trim();
match port.trim().parse::<u16>() {
Ok(p) if !host.is_empty() => Some((host.to_string(), p)),
_ => None,
}
}
None => Some((entry.to_string(), default_port)),
})
.collect()
}
pub(crate) fn is_explicit_local_worker(
local_workers: &[(String, u16)],
host: &str,
port: u16,
) -> bool {
local_workers.iter().any(|(h, p)| h == host && *p == port)
}
fn is_loopback_host(host: &str) -> bool {
host == "127.0.0.1" || host == "localhost" || host == "::1"
}
pub(crate) fn strip_zc_node(s: &str) -> &str {
s.strip_prefix("zc://").unwrap_or(s)
}
fn provider_fields_from_json(
w: &serde_json::Value,
) -> (crate::broker::worker::ProviderType, Vec<String>, f64) {
let provider_type = w
.get("provider_type")
.cloned()
.and_then(|v| serde_json::from_value(v).ok())
.unwrap_or_default();
let served_models = w
.get("served_models")
.and_then(|v| v.as_array())
.map(|a| {
a.iter()
.filter_map(|m| m.as_str())
.map(|m| m.to_string())
.collect()
})
.unwrap_or_default();
let price_per_mtok = w
.get("price_per_mtok")
.and_then(|v| v.as_f64())
.unwrap_or(0.0);
(provider_type, served_models, price_per_mtok)
}
fn is_same_peer_worker(
existing: &super::worker::Worker,
name: &str,
uri: &str,
source_node: &str,
) -> bool {
match existing.source_node.as_deref() {
Some(existing_sn) => {
let existing_fp = existing_sn.strip_prefix("node-").unwrap_or(existing_sn);
let incoming_fp = source_node.strip_prefix("node-").unwrap_or(source_node);
existing_fp == incoming_fp && (existing.name == name || existing.uri == uri)
}
None => false,
}
}
fn peer_node_fp(worker: &super::worker::Worker) -> Option<String> {
let sn = worker.source_node.as_deref()?;
let fp = sn.strip_prefix("node-").unwrap_or(sn);
if fp.is_empty() {
None
} else {
Some(fp.to_string())
}
}
pub fn detect_discovery_mode(preferred_subnet: &str) -> DiscoveryMode {
if let Some(wireguard_ip) = get_mesh_ip() {
let parts: Vec<&str> = wireguard_ip.split('.').collect();
if parts.len() == 4 {
let subnet = format!("{}.{}.{}", parts[0], parts[1], parts[2]);
return DiscoveryMode::WireGuard { subnet };
}
}
let test_ip = format!("{}.1", preferred_subnet);
if let Ok(addr) = format!("{}:1", test_ip).parse() {
if TcpStream::connect_timeout(&addr, Duration::from_millis(50)).is_ok() {
return DiscoveryMode::WireGuard {
subnet: preferred_subnet.to_string(),
};
}
}
if std::env::var("ZAKURO_PEERS")
.map(|v| !v.is_empty())
.unwrap_or(false)
{
return DiscoveryMode::WireGuard {
subnet: "peers".to_string(),
};
}
DiscoveryMode::Local
}
pub fn get_effective_node_ip() -> Option<String> {
if let Some(ip) = get_mesh_ip() {
return Some(ip);
}
if let Ok(socket) = std::net::UdpSocket::bind("0.0.0.0:0") {
if socket.connect("8.8.8.8:80").is_ok() {
if let Ok(addr) = socket.local_addr() {
let ip = addr.ip().to_string();
if ip != "127.0.0.1" && ip != "::1" {
return Some(ip);
}
}
}
}
None
}
pub fn discover_broker_peers_on_localhost(
self_port: u16,
port_start: u16,
port_end: u16,
) -> Vec<String> {
let mut out = Vec::new();
let agent = ureq::Agent::new_with_config(
ureq::Agent::config_builder()
.timeout_global(Some(Duration::from_millis(500)))
.build(),
);
for port in port_start..=port_end {
if port == self_port {
continue;
}
let url = format!("http://127.0.0.1:{}/peer/health", port);
if agent
.get(&url)
.call()
.map(|r| r.status().as_u16() == 200)
.unwrap_or(false)
{
out.push(format!("http://127.0.0.1:{}", port));
}
}
out
}
pub fn discover_broker_peers_on_mesh_subnet(self_ip: &str, peer_key: &str) -> Vec<String> {
const PORT: u16 = 9000;
const BATCH_SIZE: usize = 32;
const TIMEOUT: Duration = Duration::from_millis(400);
let hosts: Vec<String> = (1u8..=254)
.map(|h| format!("10.13.13.{}", h))
.filter(|ip| ip != self_ip)
.collect();
let mut out = Vec::new();
for batch in hosts.chunks(BATCH_SIZE) {
let (tx, rx) = std::sync::mpsc::channel();
let mut handles = Vec::with_capacity(batch.len());
for host in batch {
let host = host.clone();
let tx = tx.clone();
let peer_key = peer_key.to_string();
handles.push(thread::spawn(move || {
let agent = ureq::Agent::new_with_config(
ureq::Agent::config_builder()
.timeout_global(Some(TIMEOUT))
.build(),
);
let url = format!("http://{}:{}/peer/health", host, PORT);
let mut rb = agent.get(&url);
if !peer_key.is_empty() {
rb = rb.header("X-Peer-Key", &peer_key);
}
let ok = rb
.call()
.map(|r| r.status().as_u16() == 200)
.unwrap_or(false);
if ok {
let _ = tx.send(format!("http://{}:{}", host, PORT));
}
}));
}
drop(tx);
out.extend(rx.iter());
for h in handles {
let _ = h.join();
}
}
out
}
pub fn get_mesh_ip() -> Option<String> {
for var in ["ZAKURO_MESH_IP", "ZAKURO_WIREGUARD_IP"] {
if let Ok(ip) = std::env::var(var) {
if !ip.is_empty() {
return Some(ip);
}
}
}
#[cfg(unix)]
{
for iface in ifaces::Interface::get_all().ok()?.into_iter() {
if iface.name == "zakuro0" || iface.name.starts_with("wg") {
if let Some(addr) = iface.addr {
let addr_str = addr.to_string();
let ip = addr_str.trim_end_matches(":0");
if ip.starts_with("10.13.13.") {
return Some(ip.to_string());
}
}
}
}
}
None
}
#[cfg(test)]
mod kdi_registry_dedup_tests {
use super::is_same_peer_worker;
use crate::broker::worker::{Worker, WorkerRegistration, WorkerRegistry};
fn peer_worker(id: &str, name: &str, uri: &str, source_node: &str) -> Worker {
let mut w = Worker::new(id.into(), name.into(), uri.into(), "zakuro".into());
w.source_node = Some(source_node.into());
w
}
#[test]
fn predicate_local_worker_never_matches_a_peer() {
let local = Worker::new(
"local1".into(),
"worker-lxd".into(),
"http://127.0.0.1:3960".into(),
"zakuro".into(),
);
assert!(local.source_node.is_none());
assert!(!is_same_peer_worker(
&local,
"worker-lxd",
"http://10.13.13.9:3960",
"node-cc1fingerprint"
));
}
#[test]
fn predicate_same_name_different_source_node_not_equal() {
let peer_a = peer_worker(
"a1",
"worker-lxd",
"http://10.13.13.9:3960",
"node-i9fingerprint",
);
assert!(!is_same_peer_worker(
&peer_a,
"worker-lxd",
"http://10.13.13.10:3960",
"node-cc1fingerprint",
));
}
#[test]
fn predicate_same_name_same_source_node_is_equal() {
let peer_a = peer_worker(
"a1",
"worker-lxd",
"http://10.13.13.9:3960",
"node-i9fingerprint",
);
assert!(is_same_peer_worker(
&peer_a,
"worker-lxd",
"http://10.13.13.9:3960",
"node-i9fingerprint",
));
}
#[test]
fn same_hostname_different_node_fp_yields_distinct_registry_entries() {
let registry = WorkerRegistry::new();
let local = registry.register(WorkerRegistration {
name: "worker-lxd".into(),
uri: "http://127.0.0.1:3960".into(),
worker_type: "zakuro".into(),
resources: Default::default(),
pricing: Default::default(),
tags: vec![],
max_timeout_secs: 0.0,
hardware: Default::default(),
wireguard_ip: None,
is_docker: None,
source_node: None,
explicit_local: false,
provider_type: Default::default(),
served_models: vec![],
price_per_mtok: 0.0,
});
let local = registry.set_node_fp(&local.id, "cc1fingerprint").unwrap();
let peer_i9 = registry.register(WorkerRegistration {
name: "worker-lxd".into(),
uri: "http://10.13.13.9:3960".into(),
worker_type: "zakuro".into(),
resources: Default::default(),
pricing: Default::default(),
tags: vec![],
max_timeout_secs: 0.0,
hardware: Default::default(),
wireguard_ip: Some("10.13.13.9".into()),
is_docker: None,
source_node: Some("node-i9fingerprint".into()),
explicit_local: false,
provider_type: Default::default(),
served_models: vec![],
price_per_mtok: 0.0,
});
let peer_i9 = registry.set_node_fp(&peer_i9.id, "i9fingerprint").unwrap();
let peer_msi = registry.register(WorkerRegistration {
name: "worker-lxd".into(),
uri: "http://10.13.13.11:3960".into(),
worker_type: "zakuro".into(),
resources: Default::default(),
pricing: Default::default(),
tags: vec![],
max_timeout_secs: 0.0,
hardware: Default::default(),
wireguard_ip: Some("10.13.13.11".into()),
is_docker: None,
source_node: Some("node-msifingerprint".into()),
explicit_local: false,
provider_type: Default::default(),
served_models: vec![],
price_per_mtok: 0.0,
});
let peer_msi = registry
.set_node_fp(&peer_msi.id, "msifingerprint")
.unwrap();
assert!(!is_same_peer_worker(
&local,
"worker-lxd",
"http://10.13.13.9:3960",
"node-i9fingerprint"
));
assert!(!is_same_peer_worker(
&local,
"worker-lxd",
"http://10.13.13.11:3960",
"node-msifingerprint"
));
assert!(!is_same_peer_worker(
&peer_i9,
"worker-lxd",
"http://10.13.13.11:3960",
"node-msifingerprint"
));
assert!(is_same_peer_worker(
&peer_i9,
"worker-lxd",
"http://10.13.13.9:3960",
"node-i9fingerprint"
));
let all = registry.list();
assert_eq!(
all.len(),
3,
"local + two same-named peers must all persist"
);
assert!(local.source_node.is_none());
assert_eq!(peer_i9.source_node.as_deref(), Some("node-i9fingerprint"));
assert_eq!(peer_msi.source_node.as_deref(), Some("node-msifingerprint"));
let uri_local = local.zc_uri();
let uri_i9 = peer_i9.zc_uri();
let uri_msi = peer_msi.zc_uri();
assert_ne!(uri_local, uri_i9);
assert_ne!(uri_local, uri_msi);
assert_ne!(uri_i9, uri_msi);
assert!(uri_i9.starts_with("zc://worker-i9fingerprint-"));
assert!(uri_msi.starts_with("zc://worker-msifingerprint-"));
assert!(uri_local.starts_with("zc://worker-cc1fingerprint-"));
}
}
#[cfg(test)]
mod price_sanitize_tests {
use super::sanitize_price;
#[test]
fn strip_zc_node_handles_both_forms() {
assert_eq!(super::strip_zc_node("zc://node-i9"), "node-i9");
assert_eq!(super::strip_zc_node("node-i9"), "node-i9");
}
#[test]
fn peer_synced_worker_node_fp_survives_round_trip() {
use crate::broker::worker::Worker;
let mut w = Worker::new(
"id1".into(),
"worker-x".into(),
"http://10.13.13.9:3960".into(),
"zakuro".into(),
);
w.source_node = Some("node-deadbeefcafebabe".into());
w.slot = "3960".into();
let fp = super::peer_node_fp(&w).expect("fingerprint recovered from source_node");
assert!(!fp.is_empty());
w.node_fp = fp;
let uri = w.zc_uri();
assert!(uri.starts_with("zc://worker-deadbeefcafebabe-"));
assert!(!uri.contains("--"), "malformed uri (empty node_fp): {uri}");
}
#[test]
fn peer_node_fp_none_without_source_node() {
use crate::broker::worker::Worker;
let w = Worker::new(
"id2".into(),
"worker-y".into(),
"http://127.0.0.1:3960".into(),
"zakuro".into(),
);
assert!(super::peer_node_fp(&w).is_none());
}
#[test]
fn sanitize_price_accepts_valid() {
assert!((sanitize_price(Some(2.5), 3.6, 1000.0) - 2.5).abs() < 1e-9);
assert_eq!(sanitize_price(Some(0.0), 3.6, 1000.0), 0.0);
}
#[test]
fn sanitize_price_rejects_bad_values() {
assert_eq!(sanitize_price(Some(-1.0), 3.6, 1000.0), 3.6); assert_eq!(sanitize_price(Some(f64::NAN), 3.6, 1000.0), 3.6); assert_eq!(sanitize_price(Some(f64::INFINITY), 3.6, 1000.0), 3.6); assert_eq!(sanitize_price(Some(5000.0), 3.6, 1000.0), 3.6); assert_eq!(sanitize_price(None, 3.6, 1000.0), 3.6); }
}
#[cfg(test)]
mod localhost_only_discovery_tests {
use super::*;
use crate::broker::BrokerState;
use std::io::Read;
use std::net::TcpListener;
use std::sync::atomic::{AtomicU16, Ordering};
#[test]
fn explicit_local_workers_are_parsed_from_the_env_list() {
let parsed = parse_local_workers("172.17.0.1:3960, 10.88.0.1 ,", 3960);
assert_eq!(
parsed,
vec![
("172.17.0.1".to_string(), 3960u16),
("10.88.0.1".to_string(), 3960u16)
]
);
assert!(parse_local_workers("", 3960).is_empty());
}
#[test]
fn explicit_local_workers_bypass_the_loopback_gate() {
assert!(!is_loopback_host("172.17.0.1"));
let explicit = parse_local_workers("172.17.0.1:3960", 3960);
assert!(is_explicit_local_worker(&explicit, "172.17.0.1", 3960));
assert!(!is_explicit_local_worker(&explicit, "10.13.13.9", 3960));
assert!(!is_explicit_local_worker(&explicit, "172.17.0.1", 3961));
}
#[test]
fn is_loopback_host_gate() {
assert!(is_loopback_host("127.0.0.1"));
assert!(is_loopback_host("localhost"));
assert!(is_loopback_host("::1"));
assert!(!is_loopback_host("127.0.0.2"));
assert!(!is_loopback_host("10.13.13.9"));
assert!(!is_loopback_host("100.64.0.5"));
}
fn reserve_port(bind_ip: &str) -> TcpListener {
static NEXT: AtomicU16 = AtomicU16::new(41000);
loop {
let port = NEXT.fetch_add(1, Ordering::Relaxed);
assert!(port < 50000, "reserve_port exhausted the test port range");
if let Ok(listener) = TcpListener::bind((bind_ip, port)) {
return listener;
}
}
}
fn spawn_fake_worker(listener: TcpListener, name: &str) {
let name = name.to_string();
std::thread::spawn(move || {
for stream in listener.incoming() {
let mut stream = match stream {
Ok(s) => s,
Err(_) => continue,
};
let mut buf = [0u8; 1024];
let n = stream.read(&mut buf).unwrap_or(0);
let req = String::from_utf8_lossy(&buf[..n]);
let path = req.split_whitespace().nth(1).unwrap_or("/");
let body = if path.starts_with("/info") {
format!(
r#"{{"name":"{}","worker_type":"zakuro","resources":{{"cpus_available":1.0,"cpus_total":1.0,"memory_available":1,"memory_total":1,"gpus_available":0,"gpus_total":0}}}}"#,
name
)
} else {
"{}".to_string()
};
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body.len(),
body
);
use std::io::Write;
let _ = stream.write_all(resp.as_bytes());
}
});
}
#[test]
fn only_loopback_worker_is_registered_locally() {
let local_listener = reserve_port("127.0.0.1");
let remote_listener = reserve_port("127.0.0.2");
let local_port = local_listener.local_addr().unwrap().port();
let remote_port = remote_listener.local_addr().unwrap().port();
spawn_fake_worker(local_listener, "worker-local");
spawn_fake_worker(remote_listener, "worker-remote");
std::thread::sleep(Duration::from_millis(50));
let state = Arc::new(BrokerState::new());
let config = DiscoveryConfig {
subnet: "10.13.13".to_string(),
worker_port: local_port,
extra_ports: vec![],
scan_port_range: None,
interval_secs: 9999,
enable_scan: false,
enable_dns: false,
peers: vec![],
local_workers: vec![],
};
let discovery = Discovery {
config,
state: Arc::clone(&state),
mode: DiscoveryMode::Local,
};
discovery.try_register_worker("127.0.0.1", local_port, false);
discovery.try_register_worker("127.0.0.2", remote_port, false);
let workers = state.workers.list();
assert_eq!(
workers.len(),
1,
"only the loopback worker should be registered: {:?}",
workers.iter().map(|w| &w.uri).collect::<Vec<_>>()
);
assert_eq!(workers[0].name, "worker-local");
assert!(
workers[0].source_node.is_none(),
"the sole registered worker must be local (source_node: None)"
);
assert!(workers[0].uri.contains("127.0.0.1"));
}
}
#[cfg(test)]
mod gossip_tests {
use super::verify_gossip_body;
use crate::broker::node_identity::NodeKey;
use crate::broker::roster_cache::RosterCache;
#[test]
fn fetch_gossip_rejects_unrostered_signer() {
let key = NodeKey::generate();
let body = br#"{"peers":[{"fp":"abc123","url":"http://10.13.13.9:9000","epoch":42}]}"#;
let sig = key.sign(body);
let roster = RosterCache::from_entries(vec![]); let peers = verify_gossip_body(body, &key.public_b64(), &sig, &roster);
assert!(peers.is_empty());
}
#[test]
fn fetch_gossip_accepts_rostered_signer() {
let key = NodeKey::generate();
let body = br#"{"peers":[{"fp":"abc123","url":"http://10.13.13.9:9000","epoch":42}]}"#;
let sig = key.sign(body);
let roster = RosterCache::from_entries(vec![(key.public_b64(), false)]);
let peers = verify_gossip_body(body, &key.public_b64(), &sig, &roster);
assert_eq!(
peers,
vec![(
"abc123".to_string(),
"http://10.13.13.9:9000".to_string(),
42u64
)]
);
}
#[test]
fn fetch_gossip_rejects_tampered_body() {
let key = NodeKey::generate();
let body = br#"{"peers":[{"fp":"abc123","url":"http://10.13.13.9:9000","epoch":42}]}"#;
let sig = key.sign(body);
let roster = RosterCache::from_entries(vec![(key.public_b64(), false)]);
let tampered = br#"{"peers":[{"fp":"evil","url":"http://10.13.13.99:9000","epoch":1}]}"#;
let peers = verify_gossip_body(tampered, &key.public_b64(), &sig, &roster);
assert!(peers.is_empty());
}
#[test]
fn fetch_gossip_rejects_revoked_signer() {
let key = NodeKey::generate();
let body = br#"{"peers":[{"fp":"abc123","url":"http://10.13.13.9:9000","epoch":42}]}"#;
let sig = key.sign(body);
let roster = RosterCache::from_entries(vec![(key.public_b64(), true)]); let peers = verify_gossip_body(body, &key.public_b64(), &sig, &roster);
assert!(peers.is_empty());
}
#[test]
fn fetch_gossip_returns_empty_on_unreachable_peer() {
let roster = RosterCache::from_entries(vec![]);
let agent = ureq::Agent::new_with_config(
ureq::Agent::config_builder()
.timeout_global(Some(std::time::Duration::from_millis(200)))
.build(),
);
let key = crate::broker::node_identity::NodeKey::generate();
let peers = super::fetch_gossip("http://127.0.0.1:1", &roster, &agent, &key);
assert!(peers.is_empty());
}
#[test]
fn fetch_gossip_signs_request_so_check_node_sig_gate_accepts_it() {
use crate::broker::node_identity::{verify_request, ReplayGuard};
let responder_key = NodeKey::generate();
let responder_pub = responder_key.public_b64();
let caller_key = NodeKey::generate();
let caller_pub = caller_key.public_b64();
let server = tiny_http::Server::http("127.0.0.1:0").unwrap();
let port = server.server_addr().to_ip().unwrap().port();
std::thread::spawn(move || {
let guard = ReplayGuard::new();
for req in server.incoming_requests() {
let headers: Vec<(String, String)> = req
.headers()
.iter()
.map(|h| (h.field.to_string(), h.value.to_string()))
.collect();
let get = |name: &str| -> Option<String> {
headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case(name))
.map(|(_, v)| v.clone())
};
let is_rostered = |id: &str| id == caller_pub;
let now = crate::broker::node_identity::now_secs();
let verified =
verify_request(&is_rostered, &guard, "GET", "/peer/peers", b"", &get, now);
let resp = match verified {
Ok(_) => {
let body = br#"{"peers":[{"fp":"abc123","url":"http://10.13.13.9:9000","epoch":1}]}"#.to_vec();
let sig = responder_key.sign(&body);
tiny_http::Response::from_data(body)
.with_status_code(200)
.with_header(
tiny_http::Header::from_bytes(
"X-Node-Id",
responder_key.public_b64(),
)
.unwrap(),
)
.with_header(tiny_http::Header::from_bytes("X-Node-Sig", sig).unwrap())
}
Err(_) => {
tiny_http::Response::from_string("unauthorized").with_status_code(401)
}
};
let _ = req.respond(resp);
}
});
let base_url = format!("http://127.0.0.1:{}", port);
let roster = RosterCache::from_entries(vec![(responder_pub, false)]);
let agent = ureq::Agent::new_with_config(
ureq::Agent::config_builder()
.timeout_global(Some(std::time::Duration::from_secs(3)))
.build(),
);
let peers = super::fetch_gossip(&base_url, &roster, &agent, &caller_key);
assert_eq!(
peers,
vec![(
"abc123".to_string(),
"http://10.13.13.9:9000".to_string(),
1u64
)],
"signed request must pass the server's check_node_sig gate and return gossip"
);
let unrostered_key = NodeKey::generate();
let peers2 = super::fetch_gossip(&base_url, &roster, &agent, &unrostered_key);
assert!(
peers2.is_empty(),
"an unrostered signer must be rejected by the server's gate"
);
}
}
#[cfg(test)]
mod accept_peer_tests {
use super::accept_peer_with;
use crate::broker::node_identity::NodeKey;
use crate::broker::roster_cache::RosterCache;
fn signed_health_ok(key: &NodeKey) -> impl Fn(&str) -> Option<(String, String, Vec<u8>)> + '_ {
move |_url| {
let body = b"{\"status\":\"healthy\"}".to_vec();
let sig = key.sign(&body);
Some((key.public_b64(), sig, body))
}
}
fn signed_health_by_other() -> impl Fn(&str) -> Option<(String, String, Vec<u8>)> {
let other = NodeKey::generate();
move |_url| {
let body = b"{\"status\":\"healthy\"}".to_vec();
let sig = other.sign(&body);
Some((other.public_b64(), sig, body))
}
}
#[test]
fn accept_peer_requires_roster_and_matching_signed_health() {
let key = NodeKey::generate();
let fp = key.fingerprint();
let roster = RosterCache::from_entries(vec![(key.public_b64(), false)]);
assert!(accept_peer_with(
&fp,
"http://10.13.13.17:9000",
&roster,
signed_health_ok(&key)
));
assert!(!accept_peer_with(
&fp,
"http://10.13.13.17:9000",
&roster,
signed_health_by_other()
));
let empty = RosterCache::from_entries(vec![]);
assert!(!accept_peer_with(
&fp,
"http://10.13.13.17:9000",
&empty,
signed_health_ok(&key)
));
}
#[test]
fn accept_peer_rejects_unreachable_or_unsigned_health() {
let key = NodeKey::generate();
let fp = key.fingerprint();
let roster = RosterCache::from_entries(vec![(key.public_b64(), false)]);
let never = |_url: &str| -> Option<(String, String, Vec<u8>)> { None };
assert!(!accept_peer_with(
&fp,
"http://10.13.13.17:9000",
&roster,
never
));
}
}
#[cfg(test)]
mod advert_tests {
use super::super::node_identity::NodeKey;
use super::super::roster_cache::RosterCache;
use super::super::worker::BrokerResources;
use super::{verify_advert_body, Advert};
fn advert_json(fp: &str) -> Vec<u8> {
let advert = Advert {
fp: fp.to_string(),
price_per_hour: 3.6,
resources: BrokerResources::default(),
epoch: 1,
};
serde_json::to_vec(&advert).unwrap()
}
fn sign_ok(key: &NodeKey, body: &[u8]) -> String {
key.sign(body)
}
#[test]
fn fetch_advert_rejects_wrong_signer_and_unrostered() {
let key = NodeKey::generate();
let fp = key.fingerprint();
let roster = RosterCache::from_entries(vec![(key.public_b64(), false)]);
let body_ok = advert_json(&fp);
assert!(verify_advert_body(
&body_ok,
&key.public_b64(),
&sign_ok(&key, &body_ok),
&roster
)
.is_some());
let other = NodeKey::generate();
let other_fp = other.fingerprint();
let body_bad = advert_json(&other_fp);
assert!(verify_advert_body(
&body_bad,
&key.public_b64(),
&sign_ok(&key, &body_bad),
&roster
)
.is_none());
let empty = RosterCache::from_entries(vec![]);
assert!(verify_advert_body(
&body_ok,
&key.public_b64(),
&sign_ok(&key, &body_ok),
&empty
)
.is_none());
}
#[test]
fn discovery_round_caches_peer_adverts() {
use crate::broker::node_identity::{verify_request, ReplayGuard};
use crate::broker::worker::BrokerResources;
use std::sync::Arc;
let responder_key = NodeKey::generate();
let responder_pub = responder_key.public_b64();
let responder_fp = responder_key.fingerprint();
let server = tiny_http::Server::http("127.0.0.1:0").unwrap();
let port = server.server_addr().to_ip().unwrap().port();
std::thread::spawn(move || {
let guard = ReplayGuard::new();
for req in server.incoming_requests() {
let headers: Vec<(String, String)> = req
.headers()
.iter()
.map(|h| (h.field.to_string(), h.value.to_string()))
.collect();
let get = |name: &str| -> Option<String> {
headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case(name))
.map(|(_, v)| v.clone())
};
let is_rostered = |_id: &str| true;
let now = crate::broker::node_identity::now_secs();
let verified =
verify_request(&is_rostered, &guard, "GET", "/peer/advert", b"", &get, now);
let resp = match verified {
Ok(_) => {
let advert = super::Advert {
fp: responder_fp.clone(),
price_per_hour: 7.5,
resources: BrokerResources::default(),
epoch: 1,
};
let body = serde_json::to_vec(&advert).unwrap();
let sig = responder_key.sign(&body);
tiny_http::Response::from_data(body)
.with_status_code(200)
.with_header(
tiny_http::Header::from_bytes(
"X-Node-Id",
responder_key.public_b64(),
)
.unwrap(),
)
.with_header(tiny_http::Header::from_bytes("X-Node-Sig", sig).unwrap())
}
Err(_) => {
tiny_http::Response::from_string("unauthorized").with_status_code(401)
}
};
let _ = req.respond(resp);
}
});
let peer_url = format!("http://127.0.0.1:{}", port);
let roster = RosterCache::from_entries(vec![(responder_pub, false)]);
let state = Arc::new(crate::broker::BrokerState::new());
state.peer_manager.register_peer(peer_url.clone());
super::refresh_adverts(&state, &roster);
let a = state
.peer_manager
.peer_advert(&peer_url)
.expect("advert should be cached after refresh");
assert_eq!(a.price_per_hour, 7.5);
}
}
#[cfg(test)]
mod eager_fingerprint_seeding_tests {
use super::super::node_identity::NodeKey;
use super::super::roster_cache::RosterCache;
use super::admit_verified_peer_with;
use crate::broker::BrokerState;
use std::sync::Arc;
#[test]
fn discovery_admission_seeds_fingerprint_without_any_health_probe() {
let key = NodeKey::generate();
let fp = key.fingerprint();
let roster = RosterCache::from_entries(vec![(key.public_b64(), false)]);
let state = Arc::new(BrokerState::new());
let url = "http://10.13.13.17:9000";
let body = b"{\"status\":\"healthy\"}".to_vec();
let sig = key.sign(&body);
let pubkey = key.public_b64();
let accepted = admit_verified_peer_with(&state, &fp, url, &roster, move |_| {
Some((pubkey.clone(), sig.clone(), body.clone()))
});
assert!(accepted);
assert_eq!(
state.peer_manager.peer_fingerprint(url).as_deref(),
Some(format!("zc://node-{fp}").as_str()),
"seeded identity must be canonical `zc://node-<fp>`, not bare hex"
);
assert_eq!(
state.peer_manager.get_url_for_fingerprint(&fp).as_deref(),
Some(url),
"earn path must resolve this peer with no health probe having run"
);
assert_eq!(
state
.peer_manager
.get_url_for_fingerprint(&format!("zc://node-{fp}"))
.as_deref(),
Some(url),
"query side is normalized too: the canonical zc://node-<fp> form must \
resolve, so a consumer of executor_node_uri cannot silently fall \
into the free-work arm"
);
}
#[test]
fn rejected_peer_is_neither_registered_nor_seeded() {
let key = NodeKey::generate();
let fp = key.fingerprint();
let empty_roster = RosterCache::from_entries(vec![]);
let state = Arc::new(BrokerState::new());
let url = "http://10.13.13.18:9000";
let body = b"{\"status\":\"healthy\"}".to_vec();
let sig = key.sign(&body);
let pubkey = key.public_b64();
let accepted = admit_verified_peer_with(&state, &fp, url, &empty_roster, move |_| {
Some((pubkey.clone(), sig.clone(), body.clone()))
});
assert!(!accepted, "unrostered fp must be rejected");
assert_eq!(state.peer_manager.peer_fingerprint(url), None);
assert_eq!(state.peer_manager.get_url_for_fingerprint(&fp), None);
}
}
#[cfg(test)]
mod provider_field_sync_tests {
use super::provider_fields_from_json;
use crate::broker::worker::ProviderType;
#[test]
fn parses_provider_fields_off_peer_entry() {
let w = serde_json::json!({
"name": "zc-serve-abc",
"provider_type": "general",
"served_models": ["*"],
"price_per_mtok": 2.0,
});
let (pt, models, price) = provider_fields_from_json(&w);
assert_eq!(pt, ProviderType::General);
assert_eq!(models, vec!["*".to_string()]);
assert_eq!(price, 2.0);
}
#[test]
fn absent_fields_fall_back_to_execute_worker_defaults() {
let w = serde_json::json!({ "name": "worker-lxd" });
let (pt, models, price) = provider_fields_from_json(&w);
assert_eq!(pt, ProviderType::Specialized);
assert!(models.is_empty());
assert_eq!(price, 0.0);
}
#[test]
fn malformed_fields_do_not_poison_the_sync() {
let w = serde_json::json!({
"provider_type": 7,
"served_models": "not-a-list",
"price_per_mtok": "free",
});
let (pt, models, price) = provider_fields_from_json(&w);
assert_eq!(pt, ProviderType::Specialized);
assert!(models.is_empty());
assert_eq!(price, 0.0);
}
}