466 lines
16 KiB
Rust
466 lines
16 KiB
Rust
|
//! iroh-based P2P network driver for `DistributedNode`.
|
||
|
|
//!
|
||
|
|
//! Provides the same driver pattern as `NodeDriver` (TCP), but uses iroh's
|
||
|
|
//! QUIC-based peer-to-peer transport with built-in TLS, NAT hole-punching,
|
||
|
|
//! and relay server fallback.
|
||
|
|
//!
|
||
|
|
//! The driver owns a tokio runtime internally, exposing a synchronous API
|
||
|
|
//! (`tick()`, `recv()`, `join()`) to match the existing main loop pattern.
|
||
|
|
|
||
|
|
use std::collections::HashMap;
|
||
|
|
use std::time::Duration;
|
||
|
|
|
||
|
|
use iroh::endpoint::Connection;
|
||
|
|
use iroh::{Endpoint, PublicKey, RelayMode, SecretKey};
|
||
|
|
use tokio::runtime::Runtime as TokioRuntime;
|
||
|
|
|
||
|
|
|
||
|
|
use crate::crypto::Keypair;
|
||
|
|
use crate::messages::*;
|
||
|
|
use crate::node::{DistributedNode, DistributedNodeConfig};
|
||
|
|
use crate::snapshot::DistributionNodeSnapshot;
|
||
|
|
use crate::swim::node::NodeAction;
|
||
|
|
use crate::types::NodeId;
|
||
|
|
|
||
|
|
/// ALPN protocol identifier for SWIM messages over iroh.
|
||
|
|
const ALPN: &[u8] = b"swactor/swim/1";
|
||
|
|
|
||
|
|
// ─── Config ─────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
/// Configuration for the iroh-based driver.
|
||
|
|
pub struct IrohDriverConfig {
|
||
|
|
/// Secret key for the iroh endpoint.
|
||
|
|
/// If `None`, a fresh key is generated (node gets a random identity).
|
||
|
|
pub secret_key: Option<SecretKey>,
|
||
|
|
/// Relay server configuration.
|
||
|
|
/// Defaults to `RelayMode::Default` (n0 production relays).
|
||
|
|
pub relay_mode: RelayMode,
|
||
|
|
/// Protocol-layer configuration.
|
||
|
|
pub node: DistributedNodeConfig,
|
||
|
|
}
|
||
|
|
|
||
|
|
// ─── Driver ─────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
/// iroh P2P network driver.
|
||
|
|
///
|
||
|
|
/// Bridges the synchronous `DistributedNode` state machine with iroh's
|
||
|
|
/// async QUIC transport. Owns a tokio runtime internally.
|
||
|
|
pub struct IrohDriver {
|
||
|
|
node: DistributedNode,
|
||
|
|
endpoint: Endpoint,
|
||
|
|
rt: TokioRuntime,
|
||
|
|
connections: HashMap<NodeId, Connection>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl IrohDriver {
|
||
|
|
/// Create a new iroh driver.
|
||
|
|
///
|
||
|
|
/// Builds a tokio runtime, creates an iroh `Endpoint`, and initializes
|
||
|
|
/// the protocol-layer `DistributedNode`.
|
||
|
|
pub fn new(config: IrohDriverConfig) -> Result<Self, Box<dyn std::error::Error>> {
|
||
|
|
let rt = tokio::runtime::Builder::new_multi_thread()
|
||
|
|
.enable_all()
|
||
|
|
.build()?;
|
||
|
|
|
||
|
|
let endpoint = rt.block_on(async {
|
||
|
|
let mut builder = Endpoint::builder()
|
||
|
|
.alpns(vec![ALPN.to_vec()])
|
||
|
|
.relay_mode(config.relay_mode);
|
||
|
|
|
||
|
|
if let Some(key) = config.secret_key {
|
||
|
|
builder = builder.secret_key(key);
|
||
|
|
}
|
||
|
|
|
||
|
|
builder.bind().await
|
||
|
|
})?;
|
||
|
|
|
||
|
|
// Create a DistributedNode whose identity matches the iroh endpoint.
|
||
|
|
// Both use ed25519-dalek, so we can reconstruct our Keypair from iroh's secret key.
|
||
|
|
let iroh_secret = endpoint.secret_key().to_bytes();
|
||
|
|
let keypair = Keypair::from_bytes(&iroh_secret);
|
||
|
|
let node = DistributedNode::with_keypair(keypair, config.node);
|
||
|
|
|
||
|
|
Ok(Self {
|
||
|
|
node,
|
||
|
|
endpoint,
|
||
|
|
rt,
|
||
|
|
connections: HashMap::new(),
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
/// The node's identity.
|
||
|
|
pub fn node_id(&self) -> NodeId {
|
||
|
|
self.node.node_id()
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Access the underlying node (read-only).
|
||
|
|
pub fn node(&self) -> &DistributedNode {
|
||
|
|
&self.node
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Access the underlying node (mutable).
|
||
|
|
pub fn node_mut(&mut self) -> &mut DistributedNode {
|
||
|
|
&mut self.node
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Capture a snapshot enriched with iroh endpoint info.
|
||
|
|
pub fn snapshot(&self) -> DistributionNodeSnapshot {
|
||
|
|
let mut snap = self.node.snapshot();
|
||
|
|
// Use iroh endpoint address as the "listen address"
|
||
|
|
let addr_info = self.rt.block_on(async {
|
||
|
|
format!("{}", self.endpoint.id())
|
||
|
|
});
|
||
|
|
snap.listen_addr = Some(addr_info);
|
||
|
|
snap
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Join a cluster by connecting to seed nodes via iroh.
|
||
|
|
///
|
||
|
|
/// Each seed is identified by its iroh `PublicKey` (= our `NodeId`).
|
||
|
|
pub fn join(&mut self, seeds: &[PublicKey]) {
|
||
|
|
for seed_key in seeds {
|
||
|
|
let seed_node_id = NodeId(*seed_key.as_bytes());
|
||
|
|
if let Err(e) = self.send_join_request(*seed_key, seed_node_id) {
|
||
|
|
eprintln!("iroh driver: join error to {seed_key}: {e}");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn send_join_request(
|
||
|
|
&mut self,
|
||
|
|
seed_key: PublicKey,
|
||
|
|
seed_node_id: NodeId,
|
||
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
let msg = JoinRequest {
|
||
|
|
from: self.node.node_id(),
|
||
|
|
};
|
||
|
|
let payload = serde_json::to_vec(&msg)?;
|
||
|
|
let tag = <JoinRequest as swactor::transport::NetworkMessage>::type_tag();
|
||
|
|
|
||
|
|
let endpoint = self.endpoint.clone();
|
||
|
|
let conn = self.rt.block_on(async {
|
||
|
|
let conn = endpoint.connect(seed_key, ALPN).await?;
|
||
|
|
let mut send = conn.open_uni().await?;
|
||
|
|
write_message(&mut send, tag.as_bytes(), &payload).await?;
|
||
|
|
send.finish()?;
|
||
|
|
Ok::<_, Box<dyn std::error::Error>>(conn)
|
||
|
|
})?;
|
||
|
|
|
||
|
|
self.connections.insert(seed_node_id, conn);
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Advance the node by one tick.
|
||
|
|
pub fn tick(&mut self) {
|
||
|
|
let actions = self.node.tick();
|
||
|
|
self.send_actions(&actions);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Process incoming iroh connections and messages (non-blocking).
|
||
|
|
pub fn recv(&mut self) {
|
||
|
|
let (incoming, new_conns) = self.rt.block_on(async {
|
||
|
|
self.receive_pending().await
|
||
|
|
});
|
||
|
|
// Cache connections accepted from remote peers
|
||
|
|
for (node_id, conn) in new_conns {
|
||
|
|
self.connections.entry(node_id).or_insert(conn);
|
||
|
|
}
|
||
|
|
for (tag, payload, from_key) in incoming {
|
||
|
|
let from = NodeId(*from_key.as_bytes());
|
||
|
|
let response_actions = self.dispatch_incoming(&tag, &payload, from);
|
||
|
|
self.send_actions(&response_actions);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ─── Outgoing: NodeAction → iroh ─────────────────────────────────
|
||
|
|
|
||
|
|
fn send_actions(&mut self, actions: &[NodeAction]) {
|
||
|
|
for action in actions {
|
||
|
|
if let Err(e) = self.send_action(action) {
|
||
|
|
eprintln!("iroh driver: send error: {e}");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn send_action(&mut self, action: &NodeAction) -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
match action {
|
||
|
|
NodeAction::SendPing {
|
||
|
|
to,
|
||
|
|
sequence,
|
||
|
|
piggyback,
|
||
|
|
} => {
|
||
|
|
let msg = Ping {
|
||
|
|
from: self.node.node_id(),
|
||
|
|
sequence: *sequence,
|
||
|
|
piggyback: piggyback.clone(),
|
||
|
|
};
|
||
|
|
self.send_message(to, &msg)
|
||
|
|
}
|
||
|
|
|
||
|
|
NodeAction::SendAck {
|
||
|
|
to,
|
||
|
|
sequence,
|
||
|
|
piggyback,
|
||
|
|
} => {
|
||
|
|
let msg = Ack {
|
||
|
|
from: self.node.node_id(),
|
||
|
|
sequence: *sequence,
|
||
|
|
piggyback: piggyback.clone(),
|
||
|
|
};
|
||
|
|
self.send_message(to, &msg)
|
||
|
|
}
|
||
|
|
|
||
|
|
NodeAction::SendPingReq {
|
||
|
|
relay,
|
||
|
|
target,
|
||
|
|
sequence,
|
||
|
|
piggyback,
|
||
|
|
} => {
|
||
|
|
let msg = PingReq {
|
||
|
|
from: self.node.node_id(),
|
||
|
|
target: *target,
|
||
|
|
sequence: *sequence,
|
||
|
|
piggyback: piggyback.clone(),
|
||
|
|
};
|
||
|
|
self.send_message(relay, &msg)
|
||
|
|
}
|
||
|
|
|
||
|
|
NodeAction::SendJoinResponse { to, members } => {
|
||
|
|
let msg = JoinResponse {
|
||
|
|
members: members.clone(),
|
||
|
|
};
|
||
|
|
self.send_message(to, &msg)
|
||
|
|
}
|
||
|
|
|
||
|
|
NodeAction::MembershipChanged { .. } => Ok(()),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn send_message<M: swactor::transport::NetworkMessage + serde::Serialize>(
|
||
|
|
&mut self,
|
||
|
|
to: &NodeId,
|
||
|
|
msg: &M,
|
||
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
let tag = M::type_tag();
|
||
|
|
let payload = serde_json::to_vec(msg)?;
|
||
|
|
let target_key = PublicKey::from_bytes(&to.0)?;
|
||
|
|
|
||
|
|
let conn = self.get_or_connect(*to, target_key)?;
|
||
|
|
|
||
|
|
let result = self.rt.block_on(async {
|
||
|
|
let mut send = conn.open_uni().await?;
|
||
|
|
write_message(&mut send, tag.as_bytes(), &payload).await?;
|
||
|
|
send.finish()?;
|
||
|
|
Ok::<_, Box<dyn std::error::Error>>(())
|
||
|
|
});
|
||
|
|
|
||
|
|
if let Err(_) = result {
|
||
|
|
// Connection may be stale, remove and retry once
|
||
|
|
self.connections.remove(to);
|
||
|
|
let conn = self.get_or_connect(*to, target_key)?;
|
||
|
|
self.rt.block_on(async {
|
||
|
|
let mut send = conn.open_uni().await?;
|
||
|
|
write_message(&mut send, tag.as_bytes(), &payload).await?;
|
||
|
|
send.finish()?;
|
||
|
|
Ok::<_, Box<dyn std::error::Error>>(())
|
||
|
|
})?;
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn get_or_connect(
|
||
|
|
&mut self,
|
||
|
|
node_id: NodeId,
|
||
|
|
key: PublicKey,
|
||
|
|
) -> Result<Connection, Box<dyn std::error::Error>> {
|
||
|
|
// Check for cached connection that's still open
|
||
|
|
if let Some(conn) = self.connections.get(&node_id) {
|
||
|
|
if conn.close_reason().is_none() {
|
||
|
|
return Ok(conn.clone());
|
||
|
|
}
|
||
|
|
// Connection closed, remove it
|
||
|
|
self.connections.remove(&node_id);
|
||
|
|
}
|
||
|
|
|
||
|
|
let endpoint = self.endpoint.clone();
|
||
|
|
let conn = self.rt.block_on(async {
|
||
|
|
endpoint.connect(key, ALPN).await
|
||
|
|
})?;
|
||
|
|
|
||
|
|
self.connections.insert(node_id, conn.clone());
|
||
|
|
Ok(conn)
|
||
|
|
}
|
||
|
|
|
||
|
|
// ─── Incoming: iroh → handler ────────────────────────────────────
|
||
|
|
|
||
|
|
async fn receive_pending(&self) -> (Vec<(String, Vec<u8>, PublicKey)>, Vec<(NodeId, Connection)>) {
|
||
|
|
let mut messages = Vec::new();
|
||
|
|
let mut new_connections = Vec::new();
|
||
|
|
|
||
|
|
// Poll for incoming connections with a short timeout
|
||
|
|
loop {
|
||
|
|
let accept_fut = self.endpoint.accept();
|
||
|
|
let result = tokio::time::timeout(Duration::from_millis(1), accept_fut).await;
|
||
|
|
|
||
|
|
match result {
|
||
|
|
Ok(Some(incoming)) => {
|
||
|
|
if let Ok(conn) = incoming.await {
|
||
|
|
let remote_id = conn.remote_id();
|
||
|
|
self.read_streams(&conn, remote_id, &mut messages).await;
|
||
|
|
let node_id = NodeId(*remote_id.as_bytes());
|
||
|
|
new_connections.push((node_id, conn));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
_ => break,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Also read from existing cached connections
|
||
|
|
let conn_snapshot: Vec<(NodeId, Connection)> = self
|
||
|
|
.connections
|
||
|
|
.iter()
|
||
|
|
.map(|(id, c)| (*id, c.clone()))
|
||
|
|
.collect();
|
||
|
|
|
||
|
|
for (node_id, conn) in conn_snapshot {
|
||
|
|
let remote_id = PublicKey::from_bytes(&node_id.0).unwrap();
|
||
|
|
self.read_streams(&conn, remote_id, &mut messages).await;
|
||
|
|
}
|
||
|
|
|
||
|
|
(messages, new_connections)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn read_streams(
|
||
|
|
&self,
|
||
|
|
conn: &Connection,
|
||
|
|
remote_id: PublicKey,
|
||
|
|
messages: &mut Vec<(String, Vec<u8>, PublicKey)>,
|
||
|
|
) {
|
||
|
|
loop {
|
||
|
|
match tokio::time::timeout(Duration::from_millis(1), conn.accept_uni()).await {
|
||
|
|
Ok(Ok(mut recv)) => {
|
||
|
|
match read_message(&mut recv).await {
|
||
|
|
Ok((tag, payload)) => {
|
||
|
|
messages.push((tag, payload, remote_id));
|
||
|
|
}
|
||
|
|
Err(e) => {
|
||
|
|
eprintln!("iroh driver: read error: {e}");
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
_ => break,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn dispatch_incoming(
|
||
|
|
&mut self,
|
||
|
|
tag: &str,
|
||
|
|
payload: &[u8],
|
||
|
|
_from: NodeId,
|
||
|
|
) -> Vec<NodeAction> {
|
||
|
|
match tag {
|
||
|
|
"swactor_dist::Ping" => match serde_json::from_slice::<Ping>(payload) {
|
||
|
|
Ok(msg) => self.node.handle_ping(msg.from, msg.sequence, &msg.piggyback),
|
||
|
|
Err(e) => {
|
||
|
|
eprintln!("iroh driver: decode Ping: {e}");
|
||
|
|
Vec::new()
|
||
|
|
}
|
||
|
|
},
|
||
|
|
|
||
|
|
"swactor_dist::Ack" => match serde_json::from_slice::<Ack>(payload) {
|
||
|
|
Ok(msg) => self.node.handle_ack(msg.from, msg.sequence, &msg.piggyback),
|
||
|
|
Err(e) => {
|
||
|
|
eprintln!("iroh driver: decode Ack: {e}");
|
||
|
|
Vec::new()
|
||
|
|
}
|
||
|
|
},
|
||
|
|
|
||
|
|
"swactor_dist::PingReq" => match serde_json::from_slice::<PingReq>(payload) {
|
||
|
|
Ok(msg) => {
|
||
|
|
self.node
|
||
|
|
.handle_ping_req(msg.from, msg.target, msg.sequence, &msg.piggyback)
|
||
|
|
}
|
||
|
|
Err(e) => {
|
||
|
|
eprintln!("iroh driver: decode PingReq: {e}");
|
||
|
|
Vec::new()
|
||
|
|
}
|
||
|
|
},
|
||
|
|
|
||
|
|
"swactor_dist::JoinRequest" => {
|
||
|
|
match serde_json::from_slice::<JoinRequest>(payload) {
|
||
|
|
Ok(msg) => self.node.handle_join_request(msg.from),
|
||
|
|
Err(e) => {
|
||
|
|
eprintln!("iroh driver: decode JoinRequest: {e}");
|
||
|
|
Vec::new()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
"swactor_dist::JoinResponse" => {
|
||
|
|
match serde_json::from_slice::<JoinResponse>(payload) {
|
||
|
|
Ok(msg) => self.node.handle_join_response(msg.members),
|
||
|
|
Err(e) => {
|
||
|
|
eprintln!("iroh driver: decode JoinResponse: {e}");
|
||
|
|
Vec::new()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
other => {
|
||
|
|
eprintln!("iroh driver: unknown message type: {other}");
|
||
|
|
Vec::new()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Shut down the iroh endpoint.
|
||
|
|
pub fn shutdown(&self) {
|
||
|
|
self.rt.block_on(async {
|
||
|
|
self.endpoint.close().await;
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ─── Wire Framing Over QUIC Streams ─────────────────────────────────────────
|
||
|
|
|
||
|
|
/// Write a tagged message to a QUIC send stream.
|
||
|
|
///
|
||
|
|
/// Frame format: `[4B tag_len][tag_bytes][payload_bytes]`
|
||
|
|
async fn write_message(
|
||
|
|
send: &mut iroh::endpoint::SendStream,
|
||
|
|
tag: &[u8],
|
||
|
|
payload: &[u8],
|
||
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
let tag_len = (tag.len() as u32).to_be_bytes();
|
||
|
|
send.write_all(&tag_len).await?;
|
||
|
|
send.write_all(tag).await?;
|
||
|
|
send.write_all(payload).await?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Read a tagged message from a QUIC recv stream.
|
||
|
|
///
|
||
|
|
/// Returns `(type_tag, payload)`.
|
||
|
|
async fn read_message(
|
||
|
|
recv: &mut iroh::endpoint::RecvStream,
|
||
|
|
) -> Result<(String, Vec<u8>), Box<dyn std::error::Error>> {
|
||
|
|
let mut tag_len_buf = [0u8; 4];
|
||
|
|
recv.read_exact(&mut tag_len_buf).await?;
|
||
|
|
let tag_len = u32::from_be_bytes(tag_len_buf) as usize;
|
||
|
|
|
||
|
|
if tag_len > 1024 {
|
||
|
|
return Err("tag too large".into());
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut tag_buf = vec![0u8; tag_len];
|
||
|
|
recv.read_exact(&mut tag_buf).await?;
|
||
|
|
let tag = String::from_utf8(tag_buf)?;
|
||
|
|
|
||
|
|
let payload = recv.read_to_end(64 * 1024).await?;
|
||
|
|
|
||
|
|
Ok((tag, payload))
|
||
|
|
}
|