mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-09-03 09:35:41 +00:00
feat: relay peer end-to-end encryption via Noise IK handshake (#1960)
Enable encryption for non-direct nodes requiring relay forwarding. When secure_mode is enabled, peers perform Noise IK handshake to establish an encrypted PeerSession. Relay packets are encrypted at the sender and decrypted at the receiver. Intermediate forwarding nodes cannot read plaintext data. --------- Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: KKRainbow <5665404+KKRainbow@users.noreply.github.com>
This commit is contained in:
KKRainbow
co-authored by
Copilot
KKRainbow
Copilot
KKRainbow
parent
22b4c4be2c
commit
59d4475743
@@ -0,0 +1,692 @@
|
||||
use std::{sync::Arc, time::Instant};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use hmac::Mac;
|
||||
use prost::Message;
|
||||
use snow::params::NoiseParams;
|
||||
use tokio::sync::{oneshot, Mutex, OwnedMutexGuard};
|
||||
use tokio::time::{timeout, Duration};
|
||||
|
||||
use crate::{
|
||||
common::error::Error,
|
||||
common::{global_ctx::ArcGlobalCtx, PeerId},
|
||||
peers::peer_map::PeerMap,
|
||||
peers::peer_session::{PeerSession, PeerSessionAction, PeerSessionStore, SessionKey},
|
||||
peers::route_trait::NextHopPolicy,
|
||||
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>,
|
||||
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)>>,
|
||||
}
|
||||
|
||||
impl RelayPeerMap {
|
||||
pub fn new(
|
||||
peer_map: Arc<PeerMap>,
|
||||
global_ctx: ArcGlobalCtx,
|
||||
my_peer_id: PeerId,
|
||||
peer_session_store: Arc<PeerSessionStore>,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
peer_map,
|
||||
global_ctx,
|
||||
my_peer_id,
|
||||
peer_session_store,
|
||||
states: DashMap::new(),
|
||||
pending_handshakes: DashMap::new(),
|
||||
handshake_locks: DashMap::new(),
|
||||
pending_packets: DashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_secure_mode_enabled(&self) -> bool {
|
||||
self.global_ctx
|
||||
.config
|
||||
.get_secure_mode()
|
||||
.map(|cfg| cfg.enabled)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
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);
|
||||
self.send_via_next_hop(pkt, dst_peer_id, policy).await
|
||||
}
|
||||
|
||||
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(None));
|
||||
};
|
||||
if !self.peer_map.has_peer(next_hop) {
|
||||
return Err(Error::RouteError(None));
|
||||
}
|
||||
self.peer_map.send_msg_directly(msg, next_hop).await
|
||||
}
|
||||
|
||||
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) {
|
||||
if 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_network_name: network.network_name.clone(),
|
||||
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 server_handshake_hash = hs.get_handshake_hash().to_vec();
|
||||
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(),
|
||||
)));
|
||||
}
|
||||
if msg2_pb.b_network_name == network.network_name {
|
||||
if msg2_pb.role_hint != 1 {
|
||||
return Err(Error::RouteError(Some(
|
||||
"role_hint must be 1 when network_name is same".to_string(),
|
||||
)));
|
||||
}
|
||||
let Some(secret_proof_32) = msg2_pb.secret_proof_32 else {
|
||||
return Err(Error::RouteError(Some(
|
||||
"secret_proof_32 must be present when role_hint is 1".to_string(),
|
||||
)));
|
||||
};
|
||||
let verify_result = self
|
||||
.global_ctx
|
||||
.get_secret_proof(&server_handshake_hash)
|
||||
.map(|mac| mac.verify_slice(&secret_proof_32).is_ok());
|
||||
if verify_result != Some(true) {
|
||||
return Err(Error::RouteError(Some(
|
||||
"secret_proof_32 verify failed".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();
|
||||
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_network_name = msg1_pb.a_network_name.clone();
|
||||
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 (role_hint, secret_proof_32) = if remote_network_name == server_network_name {
|
||||
let proof = self
|
||||
.global_ctx
|
||||
.get_secret_proof(hs.get_handshake_hash())
|
||||
.map(|mac| mac.finalize().into_bytes().to_vec());
|
||||
(1, proof)
|
||||
} else {
|
||||
(2, None)
|
||||
};
|
||||
|
||||
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 {
|
||||
b_network_name: server_network_name,
|
||||
role_hint,
|
||||
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,
|
||||
secret_proof_32,
|
||||
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 fn decrypt_if_needed(&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 {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
tracing::debug!(?peer_id, "RelayPeerMap removed peer relay state");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user