mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-09-02 17:15:43 +00:00
fanyang
13412895c5
hotpath adds noticeable overhead on hot paths, so remove the optional profiling integration while keeping direct quanta::Instant timing.
695 lines
26 KiB
Rust
695 lines
26 KiB
Rust
use std::sync::Arc;
|
|
|
|
use dashmap::DashMap;
|
|
use prost::Message;
|
|
use quanta::Instant;
|
|
use snow::params::NoiseParams;
|
|
use tokio::sync::{Mutex, OwnedMutexGuard, oneshot};
|
|
use tokio::time::{Duration, timeout};
|
|
|
|
use crate::peers::foreign_network_client::ForeignNetworkClient;
|
|
use crate::{
|
|
common::error::Error,
|
|
common::{PeerId, global_ctx::ArcGlobalCtx, shrink_dashmap},
|
|
peers::peer_map::PeerMap,
|
|
peers::peer_session::{PeerSession, PeerSessionAction, PeerSessionStore, SessionKey},
|
|
peers::route_trait::NextHopPolicy,
|
|
peers::traffic_metrics::AggregateTrafficMetrics,
|
|
proto::peer_rpc::{PeerConnSessionActionPb, RelayNoiseMsg1Pb, RelayNoiseMsg2Pb},
|
|
tunnel::packet_def::{PacketType, ZCPacket},
|
|
};
|
|
|
|
const RELAY_NOISE_VERSION: u32 = 1;
|
|
const RELAY_NOISE_PROLOGUE: &[u8] = b"easytier-relay-noise";
|
|
const HANDSHAKE_TIMEOUT_SECS: u64 = 5;
|
|
const HANDSHAKE_RETRY_BASE_MS: u64 = 200;
|
|
const HANDSHAKE_MAX_ATTEMPTS: u32 = 3;
|
|
const MAX_PENDING_PACKETS_PER_PEER: usize = 32;
|
|
|
|
#[derive(Clone)]
|
|
pub struct RelayPeerState {
|
|
pub last_active_at: Instant,
|
|
pub failure_count: u32,
|
|
pub next_retry_at: Option<Instant>,
|
|
}
|
|
|
|
impl Default for RelayPeerState {
|
|
fn default() -> Self {
|
|
Self {
|
|
last_active_at: Instant::now(),
|
|
failure_count: 0,
|
|
next_retry_at: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct RelayPeerMap {
|
|
peer_map: Arc<PeerMap>,
|
|
foreign_network_client: Option<Arc<ForeignNetworkClient>>,
|
|
global_ctx: ArcGlobalCtx,
|
|
my_peer_id: PeerId,
|
|
peer_session_store: Arc<PeerSessionStore>,
|
|
states: DashMap<PeerId, RelayPeerState>,
|
|
pending_handshakes: DashMap<PeerId, oneshot::Sender<ZCPacket>>,
|
|
handshake_locks: DashMap<PeerId, Arc<Mutex<()>>>,
|
|
pub(crate) pending_packets: DashMap<PeerId, Vec<(ZCPacket, NextHopPolicy)>>,
|
|
|
|
is_secure_mode_enabled: bool,
|
|
control_metrics: AggregateTrafficMetrics,
|
|
}
|
|
|
|
impl RelayPeerMap {
|
|
pub fn new(
|
|
peer_map: Arc<PeerMap>,
|
|
foreign_network_client: Option<Arc<ForeignNetworkClient>>,
|
|
global_ctx: ArcGlobalCtx,
|
|
my_peer_id: PeerId,
|
|
peer_session_store: Arc<PeerSessionStore>,
|
|
) -> Arc<Self> {
|
|
let is_secure_mode_enabled = global_ctx
|
|
.config
|
|
.get_secure_mode()
|
|
.map(|cfg| cfg.enabled)
|
|
.unwrap_or(false);
|
|
Arc::new(Self {
|
|
control_metrics: AggregateTrafficMetrics::control(
|
|
global_ctx.stats_manager().clone(),
|
|
global_ctx.get_network_name(),
|
|
),
|
|
peer_map,
|
|
foreign_network_client,
|
|
global_ctx,
|
|
my_peer_id,
|
|
peer_session_store,
|
|
states: DashMap::new(),
|
|
pending_handshakes: DashMap::new(),
|
|
handshake_locks: DashMap::new(),
|
|
pending_packets: DashMap::new(),
|
|
is_secure_mode_enabled,
|
|
})
|
|
}
|
|
|
|
pub fn is_secure_mode_enabled(&self) -> bool {
|
|
self.is_secure_mode_enabled
|
|
}
|
|
|
|
fn get_local_keypair(&self) -> Result<(Vec<u8>, Vec<u8>), Error> {
|
|
let cfg = self
|
|
.global_ctx
|
|
.config
|
|
.get_secure_mode()
|
|
.ok_or_else(|| Error::RouteError(Some("secure mode config not set".to_string())))?;
|
|
let private = cfg
|
|
.private_key()
|
|
.map_err(|e| Error::RouteError(Some(format!("invalid private key: {e:?}"))))?;
|
|
let public = cfg
|
|
.public_key()
|
|
.map_err(|e| Error::RouteError(Some(format!("invalid public key: {e:?}"))))?;
|
|
Ok((private.as_bytes().to_vec(), public.as_bytes().to_vec()))
|
|
}
|
|
|
|
async fn get_remote_static_pubkey(&self, peer_id: PeerId) -> Result<Vec<u8>, Error> {
|
|
let info = self
|
|
.peer_map
|
|
.get_route_peer_info(peer_id)
|
|
.await
|
|
.ok_or_else(|| Error::RouteError(Some("route peer info not found".to_string())))?;
|
|
if info.noise_static_pubkey.is_empty() {
|
|
return Err(Error::RouteError(Some(
|
|
"remote static pubkey not found".to_string(),
|
|
)));
|
|
}
|
|
Ok(info.noise_static_pubkey)
|
|
}
|
|
|
|
fn get_handshake_lock(&self, peer_id: PeerId) -> Arc<Mutex<()>> {
|
|
self.handshake_locks
|
|
.entry(peer_id)
|
|
.or_insert_with(|| Arc::new(Mutex::new(())))
|
|
.clone()
|
|
}
|
|
|
|
async fn send_handshake_packet(
|
|
&self,
|
|
payload: Vec<u8>,
|
|
packet_type: PacketType,
|
|
dst_peer_id: PeerId,
|
|
policy: NextHopPolicy,
|
|
) -> Result<(), Error> {
|
|
let mut pkt = ZCPacket::new_with_payload(&payload);
|
|
pkt.fill_peer_manager_hdr(self.my_peer_id, dst_peer_id, packet_type as u8);
|
|
let pkt_len = pkt.buf_len() as u64;
|
|
self.send_via_next_hop(pkt, dst_peer_id, policy).await?;
|
|
self.control_metrics.record_tx(pkt_len);
|
|
Ok(())
|
|
}
|
|
|
|
async fn send_via_next_hop(
|
|
&self,
|
|
msg: ZCPacket,
|
|
dst_peer_id: PeerId,
|
|
policy: NextHopPolicy,
|
|
) -> Result<(), Error> {
|
|
let Some(next_hop) = self.peer_map.get_gateway_peer_id(dst_peer_id, policy).await else {
|
|
return Err(Error::RouteError(Some(format!(
|
|
"next hop not found in route for peer {dst_peer_id:?}"
|
|
))));
|
|
};
|
|
if self.peer_map.has_peer(next_hop) {
|
|
self.peer_map.send_msg_directly(msg, next_hop).await
|
|
} else if let Some(foreign_network_client) = &self.foreign_network_client {
|
|
foreign_network_client.send_msg(msg, next_hop).await
|
|
} else {
|
|
Err(Error::RouteError(Some(format!(
|
|
"next hop not found in direct peer map: {next_hop:?}"
|
|
))))
|
|
}
|
|
}
|
|
|
|
pub async fn send_msg(
|
|
self: &Arc<Self>,
|
|
mut msg: ZCPacket,
|
|
dst_peer_id: PeerId,
|
|
policy: NextHopPolicy,
|
|
) -> Result<(), Error> {
|
|
let now = Instant::now();
|
|
|
|
self.states.entry(dst_peer_id).or_default().last_active_at = now;
|
|
|
|
if self.is_secure_mode_enabled() {
|
|
match self.ensure_session(dst_peer_id, policy.clone()).await {
|
|
Ok(session) => {
|
|
let my_peer_id = self.my_peer_id;
|
|
session
|
|
.encrypt_payload(my_peer_id, dst_peer_id, &mut msg)
|
|
.map_err(|e| Error::RouteError(Some(format!("{e:?}"))))?;
|
|
}
|
|
Err(_) => {
|
|
// Handshake in progress, buffer the packet instead of dropping it
|
|
self.buffer_pending_packet(dst_peer_id, msg, policy);
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
|
|
self.send_via_next_hop(msg, dst_peer_id, policy).await
|
|
}
|
|
|
|
fn buffer_pending_packet(&self, dst_peer_id: PeerId, pkt: ZCPacket, policy: NextHopPolicy) {
|
|
let mut entry = self.pending_packets.entry(dst_peer_id).or_default();
|
|
if entry.len() < MAX_PENDING_PACKETS_PER_PEER {
|
|
entry.push((pkt, policy));
|
|
}
|
|
// silently drop when buffer is full
|
|
}
|
|
|
|
async fn flush_pending_packets(&self, dst_peer_id: PeerId, session: Arc<PeerSession>) {
|
|
let packets = self.pending_packets.remove(&dst_peer_id).map(|(_, v)| v);
|
|
let Some(packets) = packets else { return };
|
|
if packets.is_empty() {
|
|
return;
|
|
}
|
|
|
|
tracing::debug!(
|
|
?dst_peer_id,
|
|
count = packets.len(),
|
|
"flushing pending packets after relay handshake"
|
|
);
|
|
|
|
for (mut pkt, policy) in packets {
|
|
if session
|
|
.encrypt_payload(self.my_peer_id, dst_peer_id, &mut pkt)
|
|
.is_err()
|
|
{
|
|
continue;
|
|
}
|
|
let _ = self.send_via_next_hop(pkt, dst_peer_id, policy).await;
|
|
}
|
|
}
|
|
|
|
pub fn has_session(&self, dst_peer_id: PeerId) -> bool {
|
|
self.peer_session_store
|
|
.get(&SessionKey::new(
|
|
self.global_ctx.get_network_identity().network_name.clone(),
|
|
dst_peer_id,
|
|
))
|
|
.is_some()
|
|
}
|
|
|
|
pub async fn ensure_session(
|
|
self: &Arc<Self>,
|
|
dst_peer_id: PeerId,
|
|
policy: NextHopPolicy,
|
|
) -> Result<Arc<PeerSession>, Error> {
|
|
let network = self.global_ctx.get_network_identity();
|
|
let key = SessionKey::new(network.network_name.clone(), dst_peer_id);
|
|
if let Some(session) = self.peer_session_store.get(&key) {
|
|
return Ok(session);
|
|
}
|
|
|
|
let lock = self.get_handshake_lock(dst_peer_id);
|
|
if let Ok(guard) = lock.try_lock_owned() {
|
|
let self_clone = self.clone();
|
|
tokio::spawn(async move {
|
|
self_clone
|
|
.handshake_session(dst_peer_id, policy, Some(guard))
|
|
.await
|
|
});
|
|
};
|
|
Err(Error::RouteError(Some(
|
|
"relay handshake in progress".to_string(),
|
|
)))
|
|
}
|
|
|
|
#[tracing::instrument(skip(self, _lock_guard), level = "debug", ret)]
|
|
pub async fn handshake_session(
|
|
&self,
|
|
dst_peer_id: PeerId,
|
|
policy: NextHopPolicy,
|
|
_lock_guard: Option<OwnedMutexGuard<()>>,
|
|
) -> Result<(), Error> {
|
|
let network = self.global_ctx.get_network_identity();
|
|
let key = SessionKey::new(network.network_name.clone(), dst_peer_id);
|
|
if let Some(session) = self.peer_session_store.get(&key) {
|
|
self.flush_pending_packets(dst_peer_id, session).await;
|
|
return Ok(());
|
|
}
|
|
|
|
if let Some(next_retry_at) = self.states.get(&dst_peer_id).and_then(|v| v.next_retry_at)
|
|
&& Instant::now() < next_retry_at
|
|
{
|
|
self.pending_packets.remove(&dst_peer_id);
|
|
return Err(Error::RouteError(Some(
|
|
"relay handshake backoff".to_string(),
|
|
)));
|
|
}
|
|
|
|
let mut last_err = None;
|
|
for attempt in 0..HANDSHAKE_MAX_ATTEMPTS {
|
|
let ret = self
|
|
.handshake_session_once(dst_peer_id, policy.clone())
|
|
.await;
|
|
match ret {
|
|
Ok(session) => {
|
|
self.register_handshake_success(dst_peer_id);
|
|
self.flush_pending_packets(dst_peer_id, session).await;
|
|
return Ok(());
|
|
}
|
|
Err(e) => {
|
|
last_err = Some(e);
|
|
self.register_handshake_failure(dst_peer_id, attempt);
|
|
if attempt + 1 < HANDSHAKE_MAX_ATTEMPTS {
|
|
let backoff = HANDSHAKE_RETRY_BASE_MS.saturating_mul(1 << attempt);
|
|
tokio::time::sleep(Duration::from_millis(backoff)).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// All attempts failed, drop buffered packets
|
|
self.pending_packets.remove(&dst_peer_id);
|
|
|
|
Err(last_err
|
|
.unwrap_or_else(|| Error::RouteError(Some("relay handshake failed".to_string()))))
|
|
}
|
|
|
|
#[tracing::instrument(skip(self), level = "debug", ret)]
|
|
async fn handshake_session_once(
|
|
&self,
|
|
dst_peer_id: PeerId,
|
|
policy: NextHopPolicy,
|
|
) -> Result<Arc<PeerSession>, Error> {
|
|
let network = self.global_ctx.get_network_identity();
|
|
let session_key = SessionKey::new(network.network_name.clone(), dst_peer_id);
|
|
let (local_private_key, _local_public_key) = self.get_local_keypair()?;
|
|
let remote_static = self.get_remote_static_pubkey(dst_peer_id).await?;
|
|
let params: NoiseParams = "Noise_IK_25519_ChaChaPoly_SHA256"
|
|
.parse()
|
|
.map_err(|e| Error::RouteError(Some(format!("parse noise params failed: {e:?}"))))?;
|
|
|
|
let builder = snow::Builder::new(params);
|
|
let mut hs = builder
|
|
.prologue(RELAY_NOISE_PROLOGUE)
|
|
.map_err(|e| Error::RouteError(Some(format!("set prologue failed: {e:?}"))))?
|
|
.local_private_key(&local_private_key)
|
|
.map_err(|e| Error::RouteError(Some(format!("set local key failed: {e:?}"))))?
|
|
.remote_public_key(&remote_static)
|
|
.map_err(|e| Error::RouteError(Some(format!("set remote key failed: {e:?}"))))?
|
|
.build_initiator()
|
|
.map_err(|e| Error::RouteError(Some(format!("build initiator failed: {e:?}"))))?;
|
|
|
|
let a_session_generation = self
|
|
.peer_session_store
|
|
.get(&session_key)
|
|
.map(|s| s.session_generation());
|
|
let a_conn_id = uuid::Uuid::new_v4();
|
|
let msg1_pb = RelayNoiseMsg1Pb {
|
|
version: RELAY_NOISE_VERSION,
|
|
a_session_generation,
|
|
a_conn_id: Some(a_conn_id.into()),
|
|
client_encryption_algorithm: self.global_ctx.get_flags().encryption_algorithm.clone(),
|
|
};
|
|
let payload = msg1_pb.encode_to_vec();
|
|
let mut out = vec![0u8; 4096];
|
|
let out_len = hs
|
|
.write_message(&payload, &mut out)
|
|
.map_err(|e| Error::RouteError(Some(format!("noise write msg1 failed: {e:?}"))))?;
|
|
let (tx, rx) = oneshot::channel();
|
|
self.pending_handshakes.insert(dst_peer_id, tx);
|
|
|
|
let send_res = self
|
|
.send_handshake_packet(
|
|
out[..out_len].to_vec(),
|
|
PacketType::RelayHandshake,
|
|
dst_peer_id,
|
|
policy,
|
|
)
|
|
.await;
|
|
|
|
if send_res.is_err() {
|
|
self.pending_handshakes.remove(&dst_peer_id);
|
|
}
|
|
send_res?;
|
|
let msg2_pkt = match timeout(Duration::from_secs(HANDSHAKE_TIMEOUT_SECS), rx).await {
|
|
Ok(Ok(pkt)) => pkt,
|
|
Ok(Err(_)) => {
|
|
self.pending_handshakes.remove(&dst_peer_id);
|
|
return Err(Error::RouteError(Some(
|
|
"relay handshake canceled".to_string(),
|
|
)));
|
|
}
|
|
Err(_) => {
|
|
self.pending_handshakes.remove(&dst_peer_id);
|
|
return Err(Error::RouteError(Some(
|
|
"relay handshake timeout".to_string(),
|
|
)));
|
|
}
|
|
};
|
|
|
|
let msg2_pb = self.decode_handshake_message::<RelayNoiseMsg2Pb>(
|
|
PacketType::RelayHandshakeAck,
|
|
&mut hs,
|
|
msg2_pkt,
|
|
)?;
|
|
if msg2_pb.a_conn_id_echo != Some(a_conn_id.into()) {
|
|
return Err(Error::RouteError(Some(
|
|
"relay msg2 conn_id_echo mismatch".to_string(),
|
|
)));
|
|
}
|
|
|
|
let action = PeerConnSessionActionPb::try_from(msg2_pb.action)
|
|
.map_err(|_| Error::RouteError(Some("invalid session action".to_string())))?;
|
|
let session_action = match action {
|
|
PeerConnSessionActionPb::Join => PeerSessionAction::Join,
|
|
PeerConnSessionActionPb::Sync => PeerSessionAction::Sync,
|
|
PeerConnSessionActionPb::Create => PeerSessionAction::Create,
|
|
};
|
|
let remote_static_key = if remote_static.len() == 32 {
|
|
let mut key = [0u8; 32];
|
|
key.copy_from_slice(&remote_static);
|
|
Some(key)
|
|
} else {
|
|
None
|
|
};
|
|
let root_key_bytes = msg2_pb
|
|
.root_key_32
|
|
.as_deref()
|
|
.filter(|v| v.len() == 32)
|
|
.map(|v| {
|
|
let mut key_bytes = [0u8; 32];
|
|
key_bytes.copy_from_slice(v);
|
|
key_bytes
|
|
});
|
|
let algo = self.global_ctx.get_flags().encryption_algorithm.clone();
|
|
let session = self
|
|
.peer_session_store
|
|
.apply_initiator_action(
|
|
&session_key,
|
|
session_action,
|
|
msg2_pb.b_session_generation,
|
|
root_key_bytes,
|
|
msg2_pb.initial_epoch,
|
|
algo,
|
|
msg2_pb.server_encryption_algorithm.clone(),
|
|
remote_static_key,
|
|
)
|
|
.map_err(|e| Error::RouteError(Some(format!("{e:?}"))))?;
|
|
|
|
Ok(session)
|
|
}
|
|
|
|
fn register_handshake_success(&self, dst_peer_id: PeerId) {
|
|
let mut entry = self.states.entry(dst_peer_id).or_default();
|
|
entry.failure_count = 0;
|
|
entry.next_retry_at = None;
|
|
}
|
|
|
|
fn register_handshake_failure(&self, dst_peer_id: PeerId, attempt: u32) {
|
|
let mut entry = self.states.entry(dst_peer_id).or_default();
|
|
entry.failure_count = entry.failure_count.saturating_add(1);
|
|
let backoff = HANDSHAKE_RETRY_BASE_MS.saturating_mul(1 << attempt);
|
|
entry.next_retry_at = Some(Instant::now() + Duration::from_millis(backoff));
|
|
}
|
|
|
|
fn decode_handshake_message<MsgT: Message + Default>(
|
|
&self,
|
|
expected_type: PacketType,
|
|
hs: &mut snow::HandshakeState,
|
|
pkt: ZCPacket,
|
|
) -> Result<MsgT, Error> {
|
|
let hdr = pkt.peer_manager_header().ok_or_else(|| {
|
|
Error::RouteError(Some("packet without peer manager header".to_string()))
|
|
})?;
|
|
if hdr.packet_type != expected_type as u8 {
|
|
return Err(Error::RouteError(Some("packet type mismatch".to_string())));
|
|
}
|
|
let mut out = vec![0u8; 4096];
|
|
let out_len = hs
|
|
.read_message(pkt.payload(), &mut out)
|
|
.map_err(|e| Error::RouteError(Some(format!("noise read msg failed: {e:?}"))))?;
|
|
let msg = MsgT::decode(&out[..out_len])
|
|
.map_err(|e| Error::RouteError(Some(format!("decode message failed: {e:?}"))))?;
|
|
Ok(msg)
|
|
}
|
|
|
|
pub async fn handle_handshake_packet(&self, packet: ZCPacket) -> Result<(), Error> {
|
|
let hdr = packet
|
|
.peer_manager_header()
|
|
.ok_or_else(|| Error::RouteError(Some("packet without header".to_string())))?;
|
|
let src_peer_id = hdr.from_peer_id.get();
|
|
self.control_metrics.record_rx(packet.buf_len() as u64);
|
|
match hdr.packet_type {
|
|
x if x == PacketType::RelayHandshake as u8 => {
|
|
tracing::debug!("handle_relay_msg1 from {:?}", src_peer_id);
|
|
self.handle_relay_msg1(packet, src_peer_id).await
|
|
}
|
|
x if x == PacketType::RelayHandshakeAck as u8 => {
|
|
if let Some((_, sender)) = self.pending_handshakes.remove(&src_peer_id) {
|
|
let _ = sender.send(packet);
|
|
}
|
|
Ok(())
|
|
}
|
|
_ => Ok(()),
|
|
}
|
|
}
|
|
|
|
async fn handle_relay_msg1(&self, msg1: ZCPacket, remote_peer_id: PeerId) -> Result<(), Error> {
|
|
// Check for bidirectional handshake race condition.
|
|
// If we are also waiting for a RelayHandshakeAck from this peer,
|
|
// use deterministic rule: the peer with smaller peer_id becomes initiator.
|
|
if self.pending_handshakes.contains_key(&remote_peer_id) {
|
|
// We have a pending handshake as initiator.
|
|
// If remote_peer_id < my_peer_id, remote should be initiator, we should be responder.
|
|
// Cancel our pending handshake and proceed as responder.
|
|
if remote_peer_id < self.my_peer_id {
|
|
tracing::debug!(
|
|
?remote_peer_id,
|
|
my_peer_id = ?self.my_peer_id,
|
|
"bidirectional handshake race: yielding initiator role to smaller peer_id"
|
|
);
|
|
// Remove our pending handshake
|
|
self.pending_handshakes.remove(&remote_peer_id);
|
|
} else {
|
|
// We have smaller peer_id, we should remain initiator.
|
|
// Ignore this RelayHandshake and let our initiator flow complete.
|
|
tracing::debug!(
|
|
?remote_peer_id,
|
|
my_peer_id = ?self.my_peer_id,
|
|
"bidirectional handshake race: keeping initiator role due to smaller peer_id"
|
|
);
|
|
return Err(Error::RouteError(Some(
|
|
"bidirectional handshake race: we are initiator".to_string(),
|
|
)));
|
|
}
|
|
}
|
|
|
|
let (local_private_key, _local_public_key) = self.get_local_keypair()?;
|
|
let params: NoiseParams = "Noise_IK_25519_ChaChaPoly_SHA256"
|
|
.parse()
|
|
.map_err(|e| Error::RouteError(Some(format!("parse noise params failed: {e:?}"))))?;
|
|
let builder = snow::Builder::new(params);
|
|
let mut hs = builder
|
|
.prologue(RELAY_NOISE_PROLOGUE)
|
|
.map_err(|e| Error::RouteError(Some(format!("set prologue failed: {e:?}"))))?
|
|
.local_private_key(&local_private_key)
|
|
.map_err(|e| Error::RouteError(Some(format!("set local key failed: {e:?}"))))?
|
|
.build_responder()
|
|
.map_err(|e| Error::RouteError(Some(format!("build responder failed: {e:?}"))))?;
|
|
|
|
let msg1_pb = self.decode_handshake_message::<RelayNoiseMsg1Pb>(
|
|
PacketType::RelayHandshake,
|
|
&mut hs,
|
|
msg1,
|
|
)?;
|
|
let remote_static = hs
|
|
.get_remote_static()
|
|
.map(|x: &[u8]| x.to_vec())
|
|
.unwrap_or_default();
|
|
let remote_static_key = if remote_static.len() == 32 {
|
|
let mut key = [0u8; 32];
|
|
key.copy_from_slice(&remote_static);
|
|
Some(key)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Verify initiator's static public key matches the expected key from route info
|
|
let expected_pubkey = self.get_remote_static_pubkey(remote_peer_id).await?;
|
|
if remote_static != expected_pubkey {
|
|
return Err(Error::RouteError(Some(format!(
|
|
"responder: initiator static pubkey mismatch for peer {}, expected {} bytes, got {} bytes",
|
|
remote_peer_id,
|
|
expected_pubkey.len(),
|
|
remote_static.len()
|
|
))));
|
|
}
|
|
|
|
let server_network_name = self.global_ctx.get_network_name();
|
|
let algo = self.global_ctx.get_flags().encryption_algorithm.clone();
|
|
let key = SessionKey::new(server_network_name.clone(), remote_peer_id);
|
|
let upsert = self
|
|
.peer_session_store
|
|
.upsert_responder_session(
|
|
&key,
|
|
msg1_pb.a_session_generation,
|
|
algo.clone(),
|
|
msg1_pb.client_encryption_algorithm.clone(),
|
|
remote_static_key,
|
|
)
|
|
.map_err(|e| Error::RouteError(Some(format!("{e:?}"))))?;
|
|
let msg2_pb = RelayNoiseMsg2Pb {
|
|
action: match upsert.action {
|
|
PeerSessionAction::Join => PeerConnSessionActionPb::Join as i32,
|
|
PeerSessionAction::Sync => PeerConnSessionActionPb::Sync as i32,
|
|
PeerSessionAction::Create => PeerConnSessionActionPb::Create as i32,
|
|
},
|
|
b_session_generation: upsert.session_generation,
|
|
root_key_32: upsert.root_key.map(|k| k.to_vec()),
|
|
initial_epoch: upsert.initial_epoch,
|
|
b_conn_id: Some(uuid::Uuid::new_v4().into()),
|
|
a_conn_id_echo: msg1_pb.a_conn_id,
|
|
server_encryption_algorithm: algo,
|
|
};
|
|
let payload = msg2_pb.encode_to_vec();
|
|
let mut out = vec![0u8; 4096];
|
|
let out_len = hs
|
|
.write_message(&payload, &mut out)
|
|
.map_err(|e| Error::RouteError(Some(format!("noise write msg2 failed: {e:?}"))))?;
|
|
|
|
self.register_handshake_success(remote_peer_id);
|
|
|
|
self.send_handshake_packet(
|
|
out[..out_len].to_vec(),
|
|
PacketType::RelayHandshakeAck,
|
|
remote_peer_id,
|
|
NextHopPolicy::LeastHop,
|
|
)
|
|
.await?;
|
|
|
|
// Flush any packets buffered while waiting for the handshake to complete
|
|
self.flush_pending_packets(remote_peer_id, upsert.session)
|
|
.await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn decrypt_if_needed(self: &Arc<Self>, packet: &mut ZCPacket) -> Result<bool, Error> {
|
|
if !self.is_secure_mode_enabled() {
|
|
return Ok(false);
|
|
}
|
|
let hdr = packet
|
|
.peer_manager_header()
|
|
.ok_or_else(|| Error::RouteError(Some("packet without header".to_string())))?;
|
|
let from_peer_id = hdr.from_peer_id.get();
|
|
let network = self.global_ctx.get_network_identity();
|
|
let key = SessionKey::new(network.network_name.clone(), from_peer_id);
|
|
let Some(session) = self.peer_session_store.get(&key) else {
|
|
tracing::debug!(
|
|
"relay session not found for peer {}, try handshake",
|
|
from_peer_id
|
|
);
|
|
self.ensure_session(from_peer_id, NextHopPolicy::LeastHop)
|
|
.await?;
|
|
return Ok(false);
|
|
};
|
|
let now = Instant::now();
|
|
let mut entry = self.states.entry(from_peer_id).or_default();
|
|
entry.last_active_at = now;
|
|
session.decrypt_payload(from_peer_id, self.my_peer_id, packet)?;
|
|
Ok(true)
|
|
}
|
|
|
|
pub fn evict_idle_sessions(&self, idle: Duration) {
|
|
let now = Instant::now();
|
|
let mut to_remove = Vec::new();
|
|
for entry in self.states.iter() {
|
|
if now.duration_since(entry.last_active_at) > idle {
|
|
to_remove.push(*entry.key());
|
|
}
|
|
}
|
|
for peer_id in to_remove {
|
|
self.states.remove(&peer_id);
|
|
self.pending_handshakes.remove(&peer_id);
|
|
self.handshake_locks.remove(&peer_id);
|
|
self.pending_packets.remove(&peer_id);
|
|
}
|
|
shrink_dashmap(&self.states, None);
|
|
shrink_dashmap(&self.pending_handshakes, None);
|
|
shrink_dashmap(&self.handshake_locks, None);
|
|
shrink_dashmap(&self.pending_packets, None);
|
|
}
|
|
|
|
pub fn has_state(&self, peer_id: PeerId) -> bool {
|
|
self.states.contains_key(&peer_id)
|
|
}
|
|
|
|
pub fn failure_count(&self, peer_id: PeerId) -> Option<u32> {
|
|
self.states.get(&peer_id).map(|v| v.failure_count)
|
|
}
|
|
|
|
pub fn is_backoff_active(&self, peer_id: PeerId) -> bool {
|
|
self.states
|
|
.get(&peer_id)
|
|
.and_then(|v| v.next_retry_at)
|
|
.is_some_and(|ts| Instant::now() < ts)
|
|
}
|
|
|
|
/// Remove relay-specific state for a specific peer.
|
|
/// This does NOT remove the session from PeerSessionStore, because the
|
|
/// session lifecycle is independent of any particular connection type
|
|
/// (relay or direct). The session may still be used by direct connections
|
|
/// or for fast reconnection (Join instead of Create).
|
|
pub fn remove_peer(&self, peer_id: PeerId) {
|
|
self.states.remove(&peer_id);
|
|
self.pending_handshakes.remove(&peer_id);
|
|
self.handshake_locks.remove(&peer_id);
|
|
self.pending_packets.remove(&peer_id);
|
|
shrink_dashmap(&self.states, None);
|
|
shrink_dashmap(&self.pending_handshakes, None);
|
|
shrink_dashmap(&self.handshake_locks, None);
|
|
shrink_dashmap(&self.pending_packets, None);
|
|
|
|
tracing::debug!(?peer_id, "RelayPeerMap removed peer relay state");
|
|
}
|
|
}
|