108 lines
3.1 KiB
Rust
108 lines
3.1 KiB
Rust
use crate::network::get_broadcast_creation_info;
|
|
use crate::orchestration::PeerHandle;
|
|
use crate::topology::DeviceCapabilities;
|
|
use serde::{Deserialize, Serialize};
|
|
use socket2::{Domain, Protocol, Socket, Type};
|
|
use std::collections::HashMap;
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use tokio::net::UdpSocket;
|
|
use tokio::sync::Mutex;
|
|
use tokio::task::JoinHandle;
|
|
use tracing::{debug, info};
|
|
use uuid::Uuid;
|
|
|
|
mod broadcast;
|
|
mod udp_listen;
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
pub struct DiscoveryMessage {
|
|
#[serde(rename = "type")]
|
|
pub message_type: String,
|
|
pub node_id: String,
|
|
pub grpc_port: u16,
|
|
pub device_capabilities: DeviceCapabilities,
|
|
pub priority: u8,
|
|
pub interface_name: String,
|
|
pub interface_type: String,
|
|
}
|
|
|
|
pub fn bind_to_address(address: SocketAddr) -> UdpSocket {
|
|
let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP)).unwrap();
|
|
socket.set_broadcast(true).unwrap();
|
|
socket.set_reuse_address(true).unwrap();
|
|
|
|
#[cfg(not(target_os = "windows"))]
|
|
socket.set_reuse_port(true).unwrap();
|
|
|
|
socket.bind(&address.into()).unwrap();
|
|
UdpSocket::from_std(socket.into()).unwrap()
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct NodeInfo {
|
|
pub node_id: String,
|
|
pub discovery_listen_port: u16,
|
|
pub broadcast_port: u16,
|
|
pub broadcast_interval: Duration,
|
|
pub grpc_port: u16,
|
|
pub allowed_peer_ids: Option<Vec<String>>,
|
|
pub allowed_interfaces: Option<Vec<String>>,
|
|
pub discovery_timeout: Duration,
|
|
pub device_capabilities: DeviceCapabilities,
|
|
}
|
|
|
|
impl Default for NodeInfo {
|
|
fn default() -> Self {
|
|
NodeInfo {
|
|
node_id: Uuid::new_v4().to_string(),
|
|
discovery_listen_port: 5678,
|
|
broadcast_port: 5678,
|
|
broadcast_interval: Duration::from_secs_f32(2.5),
|
|
grpc_port: 49152,
|
|
allowed_peer_ids: None,
|
|
allowed_interfaces: None,
|
|
discovery_timeout: Duration::from_secs(30),
|
|
device_capabilities: DeviceCapabilities::determine(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct UdpDiscovery {
|
|
discovery_handle: JoinHandle<()>,
|
|
presence_handle: JoinHandle<()>,
|
|
peer_manager_handle: JoinHandle<()>,
|
|
pub peers: Arc<Mutex<HashMap<String, PeerHandle>>>,
|
|
}
|
|
|
|
impl UdpDiscovery {
|
|
pub fn new(node_info: NodeInfo) -> Self {
|
|
let broadcast_creation_info = get_broadcast_creation_info();
|
|
info!("Found addresses: {:?}", broadcast_creation_info);
|
|
|
|
let peers = Arc::new(Mutex::new(HashMap::new()));
|
|
|
|
let discovery_handle = tokio::spawn(broadcast::listen_all(
|
|
node_info.clone(),
|
|
broadcast_creation_info,
|
|
));
|
|
let (presence_handle, peer_manager_handle) =
|
|
udp_listen::manage_discovery(node_info.clone(), peers.clone());
|
|
|
|
UdpDiscovery {
|
|
discovery_handle,
|
|
presence_handle,
|
|
peer_manager_handle,
|
|
peers,
|
|
}
|
|
}
|
|
|
|
pub fn stop(&self) {
|
|
self.discovery_handle.abort();
|
|
self.presence_handle.abort();
|
|
self.peer_manager_handle.abort();
|
|
}
|
|
}
|