141 lines
4.0 KiB
Rust
141 lines
4.0 KiB
Rust
use crate::orchestration::PeerHandle;
|
|
use crate::{discovery::DiscoveryMessage, node_service::HealthCheckRequest};
|
|
use std::collections::HashMap;
|
|
use std::net::SocketAddr;
|
|
use system_configuration::sys::libc::disconnectx;
|
|
use tokio::net::UdpSocket;
|
|
use tokio::select;
|
|
use tokio::sync::mpsc::UnboundedSender;
|
|
use tonic::transport::Error;
|
|
use tracing::{debug, error, info};
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct NodeInfo {
|
|
id: String,
|
|
listen_port: u16,
|
|
allowed_peer_ids: Option<Vec<String>>,
|
|
allowed_interfaces: Option<Vec<String>>,
|
|
}
|
|
|
|
pub async fn listen_for_discovery(
|
|
node_info: NodeInfo,
|
|
tx: UnboundedSender<(SocketAddr, DiscoveryMessage)>,
|
|
) {
|
|
let socket = UdpSocket::bind(format!("0.0.0.0:{}", node_info.listen_port))
|
|
.await
|
|
.unwrap();
|
|
let mut buf = vec![0u8; 65535];
|
|
|
|
loop {
|
|
let (len, addr) = socket.recv_from(&mut buf).await.unwrap();
|
|
if len == 0 {
|
|
continue;
|
|
}
|
|
|
|
let Ok(message) = String::from_utf8(buf[..len].to_vec()) else {
|
|
error!("Invalid UTF-8 message from {}", addr);
|
|
continue;
|
|
};
|
|
|
|
let Ok(message) = serde_json::from_str::<DiscoveryMessage>(&message) else {
|
|
error!("Invalid discovery message from {}", addr);
|
|
continue;
|
|
};
|
|
|
|
// Validate message
|
|
if message.message_type != "discovery" || message.node_id == node_info.id {
|
|
continue;
|
|
}
|
|
|
|
if node_info
|
|
.allowed_peer_ids
|
|
.as_ref()
|
|
.map(|ids| !ids.contains(&message.node_id))
|
|
.unwrap_or(false)
|
|
{
|
|
debug!(
|
|
"Ignoring peer {peer_id} as it's not in the allowed node IDs list",
|
|
peer_id = message.node_id
|
|
);
|
|
continue;
|
|
}
|
|
|
|
if node_info
|
|
.allowed_interfaces
|
|
.as_ref()
|
|
.map(|interfaces| !interfaces.contains(&message.interface_name))
|
|
.unwrap_or(false)
|
|
{
|
|
debug!("Ignoring peer {peer_id} as it's interface {interface} is not in the allowed interfaces list", peer_id = message.node_id, interface = message.interface_name);
|
|
continue;
|
|
}
|
|
|
|
tx.send((addr, message)).unwrap();
|
|
}
|
|
}
|
|
|
|
struct PeerInfo {
|
|
address: SocketAddr,
|
|
priority: u8,
|
|
}
|
|
|
|
pub async fn manage_discovery(node_info: NodeInfo) {
|
|
let mut peers: HashMap<String, PeerHandle> = HashMap::new();
|
|
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<(SocketAddr, DiscoveryMessage)>();
|
|
tokio::spawn(listen_for_discovery(node_info.clone(), tx));
|
|
|
|
while let Some((addr, message)) = rx.recv().await {
|
|
info!("Received discovery message from {}", message.node_id);
|
|
let existing = peers.get(&message.node_id);
|
|
let insert_new = match existing {
|
|
None => true,
|
|
Some(existing) => {
|
|
existing.address != addr && existing.address_priority < message.priority
|
|
}
|
|
};
|
|
|
|
if !insert_new {
|
|
continue;
|
|
}
|
|
|
|
let description = format!("{} ({})", message.interface_type, message.interface_name);
|
|
|
|
let a = PeerHandle::new(
|
|
message.node_id.clone(),
|
|
addr.clone(),
|
|
message.priority,
|
|
Some(description),
|
|
message.device_capabilities.clone(),
|
|
)
|
|
.await;
|
|
|
|
let a = match a {
|
|
Ok(a) => a,
|
|
Err(error) => {
|
|
error!(
|
|
"Failed to connect to new peer {} at {}: {}",
|
|
message.node_id, addr, error
|
|
);
|
|
continue;
|
|
}
|
|
};
|
|
|
|
let is_healthy = a
|
|
.client
|
|
.lock()
|
|
.await
|
|
.health_check(HealthCheckRequest::default())
|
|
.await
|
|
.ok()
|
|
.map(|x| x.into_inner().is_healthy)
|
|
.unwrap_or(false);
|
|
|
|
if !is_healthy {
|
|
error!("Peer {} is not healthy", message.node_id);
|
|
continue;
|
|
}
|
|
|
|
peers.insert(message.node_id, a);
|
|
}
|
|
}
|