1mod batch;
34pub mod behavior;
35#[cfg_attr(not(feature = "cortex"), allow(dead_code))]
42mod cancel_registry;
43pub mod channel;
44pub mod compute;
45mod config;
46pub mod contested;
47pub mod continuity;
48#[cfg(feature = "cortex")]
49pub mod cortex;
50mod crypto;
51mod failure;
52pub mod identity;
53mod mesh;
54#[cfg(feature = "dataforts")]
61pub mod dataforts;
62#[cfg(feature = "cortex")]
63pub mod mesh_rpc;
64#[cfg(feature = "cortex")]
65pub mod mesh_rpc_metrics;
66#[cfg(feature = "netdb")]
70pub mod netdb;
71#[cfg(feature = "cortex")]
72pub mod org_admission_gate;
73mod pool;
74mod protocol;
75mod proxy;
76#[cfg(feature = "redex")]
77pub mod redex;
78mod reliability;
79mod reroute;
80mod route;
81mod router;
82mod session;
83pub mod state;
84mod stream;
85pub mod subnet;
86pub mod subprotocol;
87mod swarm;
88mod transport;
89#[cfg(feature = "nat-traversal")]
90pub mod traversal;
91
92#[cfg(target_os = "linux")]
93mod linux;
94
95pub use batch::AdaptiveBatcher;
96pub use channel::{
97 queue_group_hash, AckReason, AclPrincipal, AuthGuard, AuthVerdict, ChannelConfig,
98 ChannelConfigRegistry, ChannelError, ChannelHash, ChannelId, ChannelName, ChannelPublisher,
99 ChannelRegistry, MembershipMsg, OnFailure, OriginBinding, PublishConfig, PublishReport,
100 QueueGroupPolicy, ResolvedConfig, SubscriberRoster, Visibility, SUBPROTOCOL_CHANNEL_MEMBERSHIP,
101};
102pub use compute::{
103 DaemonError, DaemonFactoryRegistry, DaemonHost, DaemonHostConfig, DaemonRegistry, DaemonStats,
104 FactoryEntry, MeshDaemon, MigrationError, MigrationMessage, MigrationOrchestrator,
105 MigrationPhase, MigrationSourceHandler, MigrationState, MigrationTargetHandler,
106 PlacementDecision, Scheduler, SchedulerError, SUBPROTOCOL_MIGRATION,
107};
108pub use config::{ConnectionRole, NetAdapterConfig, ReliabilityConfig};
109pub use contested::{
110 CorrelatedFailureConfig, CorrelatedFailureDetector, CorrelationVerdict, FailureCause,
111 PartitionDetector, PartitionPhase, PartitionRecord, ReconcileOutcome, Side,
112 SUBPROTOCOL_PARTITION,
113};
114pub use continuity::{
115 assess_continuity, CausalCone, Causality, ContinuityProof, ContinuityStatus, Discontinuity,
116 DiscontinuityReason, ForkRecord, HorizonDivergence, ObservationWindow, ProofError,
117 PropagationModel, SuperpositionPhase, SuperpositionState, SUBPROTOCOL_CONTINUITY,
118};
119#[cfg(feature = "cortex")]
120pub use cortex::{
121 CortexAdapter, CortexAdapterConfig, CortexAdapterError, EventEnvelope, EventMeta,
122 FoldErrorPolicy, IntoRedexPayload, StartPosition, EVENT_META_SIZE,
123};
124pub use crypto::{CryptoError, SessionKeys, StaticKeypair};
125pub use failure::{
126 CircuitBreaker, CircuitState, FailureDetector, FailureDetectorConfig, FailureStats,
127 LossSimulator, NodeStatus, PeerFailureEvent, RecoveryAction, RecoveryManager, RecoveryStats,
128 VerdictStatus,
129};
130pub use identity::{
131 EntityError, EntityId, EntityKeypair, OriginStamp, PermissionToken, TokenCache, TokenError,
132 TokenScope,
133};
134#[cfg(feature = "nat-traversal")]
137pub use mesh::UpgradeAttemptGuard;
138pub use mesh::{
139 ControlPlaneStats, MeshNode, MeshNodeConfig, PartitionFilter, SensingReadinessOverlay,
140 SensingRegistrationError, UnregisteredChannelPolicy, ACK_RANGES_CAPABILITY_TAG,
141};
142#[cfg(feature = "netdb")]
143pub use netdb::{MemoriesFilter, NetDb, NetDbBuilder, NetDbError, NetDbSnapshot, TasksFilter};
144pub use pool::{PacketBuilder, PacketPool, SharedLocalPool, ThreadLocalPool};
150pub use protocol::{
151 EventFrame, NackPayload, NetHeader, PacketFlags, HEADER_SIZE, NONCE_SIZE, TAG_SIZE,
152};
153pub use proxy::{
154 ForwardResult, HopStats, MultiHopPacketBuilder, NetProxy, ProxyConfig, ProxyError, ProxyStats,
155};
156#[cfg(feature = "redex")]
157pub use redex::{
158 FsyncPolicy, IndexOp, IndexStart, OrderedAppender, Redex, RedexEntry, RedexError, RedexEvent,
159 RedexFile, RedexFileConfig, RedexFlags, RedexFold, RedexIndex, TypedRedexFile,
160};
161pub use reliability::{FireAndForget, ReliabilityMode, ReliableStream, RetransmitDescriptor};
162pub use reroute::ReroutePolicy;
163pub use route::{
164 AggregateStats, AlternateProvenance, RouteCandidateView, RouteEntry, RouteFlags,
165 RouteObservation, RoutingHeader, RoutingTable, SchedulerStreamStats, TransitionOutcome,
166 ROUTING_HEADER_SIZE,
167};
168pub use router::{FairScheduler, NetRouter, RouteAction, RouterConfig, RouterError, RouterStats};
169#[doc(hidden)]
173pub use router::{
174 arm_send_drain_histo, send_batch_stats, send_drain_histo_snapshot, send_drain_max,
175};
176pub use session::{NetSession, SessionManager, StreamState, TxAdmit, TxSlotGuard};
177pub use state::{
178 CausalChainBuilder, CausalEvent, CausalLink, ChainError, EntityLog, HorizonEncoder, LogError,
179 LogIndex, ObservedHorizon, SnapshotStore, StateSnapshot, CAUSAL_LINK_SIZE, SUBPROTOCOL_CAUSAL,
180 SUBPROTOCOL_SNAPSHOT,
181};
182pub use stream::{
183 CloseBehavior, Reliability, Stream, StreamConfig, StreamError, StreamStats,
184 DEFAULT_STREAM_WINDOW_BYTES,
185};
186pub use subnet::{DropReason, ForwardDecision, SubnetGateway, SubnetId, SubnetPolicy, SubnetRule};
187pub use subprotocol::{
188 negotiate, MigrationSubprotocolHandler, NegotiatedSet, OutboundMigrationMessage,
189 SubprotocolDescriptor, SubprotocolManifest, SubprotocolRegistry, SubprotocolVersion,
190 SUBPROTOCOL_NEGOTIATION,
191};
192pub use swarm::{
193 Capabilities, CapabilityAd, EdgeInfo, GraphStats, LocalGraph, NodeInfo, Pingwave,
194 MAX_GRAPH_NODES, MAX_SEEN_PINGWAVES, PINGWAVE_SIZE,
195};
196pub use transport::{NetSocket, PacketReceiver, PacketSender, ParsedPacket, SocketBufferConfig};
197#[cfg(feature = "batched-ingress")]
203#[doc(hidden)]
204pub use transport::{
205 arm_recv_drain_histo, recv_batch_stats, recv_drain_histo_snapshot, recv_drain_max,
206 RECV_DRAIN_BUCKETS,
207};
208
209use async_trait::async_trait;
210use bytes::Bytes;
211use crossbeam_queue::SegQueue;
212use dashmap::DashMap;
213use std::sync::atomic::{AtomicBool, Ordering};
214use std::sync::Arc;
215use tokio::sync::Mutex as TokioMutex;
216use tokio::sync::Notify;
217use tokio::task::JoinHandle;
218
219use crate::adapter::{Adapter, ShardPollResult};
220use crate::error::AdapterError;
221use crate::event::{Batch, StoredEvent};
222
223use crypto::NoiseHandshake;
224use session::SessionManager as SessionMgr;
225use transport::NetSocket as Socket;
226
227pub use routing::{route_to_shard, stream_id_from_bytes, stream_id_from_key};
229
230const COARSE_CLOCK_REFRESH_NS: u64 = 1_000_000; #[inline]
261pub(crate) fn current_timestamp() -> u64 {
262 thread_local! {
283 static COARSE_CLOCK: std::cell::Cell<Option<(std::time::Instant, u64)>>
284 = const { std::cell::Cell::new(None) };
285 }
286 COARSE_CLOCK.with(|cell| {
287 let now_inst = std::time::Instant::now();
288 let (store, ns) = coarse_clock_advance(cell.get(), now_inst, || {
289 let elapsed = std::time::SystemTime::now()
290 .duration_since(std::time::UNIX_EPOCH)
291 .unwrap_or_default();
292 u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX)
293 });
294 if let Some(pair) = store {
295 cell.set(Some(pair));
296 }
297 ns
298 })
299}
300
301#[inline]
318fn coarse_clock_advance(
319 cached: Option<(std::time::Instant, u64)>,
320 now_inst: std::time::Instant,
321 read_wall: impl FnOnce() -> u64,
322) -> (Option<(std::time::Instant, u64)>, u64) {
323 if let Some((last_inst, last_ns)) = cached {
324 if now_inst.duration_since(last_inst).as_nanos() < COARSE_CLOCK_REFRESH_NS as u128 {
325 return (None, last_ns);
326 }
327 }
328 let ns = read_wall();
329 (Some((now_inst, ns)), ns)
330}
331
332#[inline]
338pub(crate) fn current_timestamp_micros() -> u64 {
339 std::time::SystemTime::now()
340 .duration_since(std::time::UNIX_EPOCH)
341 .map(|d| d.as_micros() as u64)
342 .unwrap_or(0)
343}
344
345mod routing {
349 use xxhash_rust::xxh3::xxh3_64;
350
351 #[inline]
355 pub fn stream_id_from_bytes(data: &[u8]) -> u64 {
356 xxh3_64(data)
357 }
358
359 #[inline]
363 pub fn stream_id_from_key(key: &str) -> u64 {
364 xxh3_64(key.as_bytes())
365 }
366
367 #[inline]
375 pub fn route_to_shard(data: &[u8], num_shards: u16) -> u16 {
376 assert!(num_shards > 0, "num_shards must be > 0");
377 (xxh3_64(data) % num_shards as u64) as u16
378 }
379
380 #[cfg(test)]
381 mod tests {
382 use super::*;
383
384 #[test]
385 fn test_stream_id_deterministic() {
386 let data = b"test event data";
387 let id1 = stream_id_from_bytes(data);
388 let id2 = stream_id_from_bytes(data);
389 assert_eq!(id1, id2);
390 }
391
392 #[test]
393 fn test_stream_id_different_for_different_data() {
394 let id1 = stream_id_from_bytes(b"event1");
395 let id2 = stream_id_from_bytes(b"event2");
396 assert_ne!(id1, id2);
397 }
398
399 #[test]
400 fn test_stream_id_from_key() {
401 let id = stream_id_from_key("user:12345");
402 assert_ne!(id, 0);
403 }
404
405 #[test]
406 fn test_route_to_shard_range() {
407 let num_shards = 16u16;
408 for i in 0..1000 {
409 let data = format!("event_{}", i);
410 let shard = route_to_shard(data.as_bytes(), num_shards);
411 assert!(shard < num_shards);
412 }
413 }
414
415 #[test]
416 #[should_panic(expected = "num_shards must be > 0")]
417 fn test_route_to_shard_zero_shards_panics() {
418 route_to_shard(b"test", 0);
421 }
422
423 #[test]
424 fn test_route_to_shard_distribution() {
425 let num_shards = 8u16;
426 let mut counts = [0u32; 8];
427
428 for i in 0..8000 {
429 let data = format!("event_{}", i);
430 let shard = route_to_shard(data.as_bytes(), num_shards);
431 counts[shard as usize] += 1;
432 }
433
434 let expected = 1000;
436 for count in counts {
437 assert!(count > expected / 2, "shard count {} too low", count);
438 assert!(count < expected * 2, "shard count {} too high", count);
439 }
440 }
441 }
442}
443
444type InboundQueues = Arc<DashMap<u16, SegQueue<StoredEvent>>>;
446
447pub(crate) struct HandshakePacer {
461 entries: std::collections::HashMap<std::net::SocketAddr, (u32, std::time::Instant)>,
463 max_per_window: u32,
465 window: std::time::Duration,
467 last_gc: std::time::Instant,
469 gc_size_threshold: usize,
473}
474
475impl HandshakePacer {
476 pub(crate) fn new(max_per_window: u32, window: std::time::Duration) -> Self {
477 Self {
478 entries: std::collections::HashMap::new(),
479 max_per_window,
480 window,
481 last_gc: std::time::Instant::now(),
482 gc_size_threshold: 4096,
486 }
487 }
488
489 pub(crate) fn check_and_record(&mut self, source: std::net::SocketAddr) -> bool {
493 let now = std::time::Instant::now();
494 if now.duration_since(self.last_gc) >= self.window
503 || self.entries.len() >= self.gc_size_threshold
504 {
505 let cutoff = self.window.saturating_mul(2);
506 self.entries
507 .retain(|_, (_, start)| now.duration_since(*start) < cutoff);
508 self.last_gc = now;
509 }
510
511 let entry = self.entries.entry(source).or_insert((0, now));
512 if now.duration_since(entry.1) > self.window {
513 entry.0 = 0;
515 entry.1 = now;
516 }
517 entry.0 = entry.0.saturating_add(1);
518 entry.0 <= self.max_per_window
519 }
520}
521
522pub struct NetAdapter {
524 config: NetAdapterConfig,
526 socket: Option<Arc<Socket>>,
528 session: Option<Arc<NetSession>>,
530 session_manager: SessionMgr,
532 inbound: InboundQueues,
534 tasks: TokioMutex<Vec<JoinHandle<()>>>,
536 shutdown: Arc<AtomicBool>,
538 shutdown_notify: Arc<Notify>,
540 initialized: AtomicBool,
542 handshake_pacer: parking_lot::Mutex<HandshakePacer>,
547}
548
549impl NetAdapter {
550 pub fn new(config: NetAdapterConfig) -> Result<Self, AdapterError> {
552 config
553 .validate()
554 .map_err(|e| AdapterError::Fatal(format!("invalid config: {}", e)))?;
555
556 Ok(Self {
557 session_manager: SessionMgr::new(config.session_timeout),
558 config,
559 socket: None,
560 session: None,
561 inbound: Arc::new(DashMap::new()),
562 tasks: TokioMutex::new(Vec::new()),
563 shutdown: Arc::new(AtomicBool::new(false)),
564 shutdown_notify: Arc::new(Notify::new()),
565 initialized: AtomicBool::new(false),
566 handshake_pacer: parking_lot::Mutex::new(HandshakePacer::new(
570 5,
571 std::time::Duration::from_secs(1),
572 )),
573 })
574 }
575
576 async fn perform_handshake(
579 &self,
580 socket: &Socket,
581 ) -> Result<(SessionKeys, std::net::SocketAddr), AdapterError> {
582 let mut attempt = 0;
583 let max_attempts = self.config.handshake_retries;
584
585 const HANDSHAKE_RETRY_SLEEP_CAP_MS: u64 = 5_000;
593
594 loop {
595 attempt += 1;
596 match self.try_handshake(socket).await {
597 Ok(result) => return Ok(result),
598 Err(e) if attempt < max_attempts => {
599 tracing::warn!(
600 attempt = attempt,
601 max = max_attempts,
602 error = %e,
603 "handshake failed, retrying"
604 );
605 let backoff_ms =
606 (100u64.saturating_mul(attempt as u64)).min(HANDSHAKE_RETRY_SLEEP_CAP_MS);
607 tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
608 }
609 Err(e) => return Err(e),
610 }
611 }
612 }
613
614 async fn try_handshake(
617 &self,
618 socket: &Socket,
619 ) -> Result<(SessionKeys, std::net::SocketAddr), AdapterError> {
620 let timeout = self.config.handshake_timeout;
621 let socket_arc = socket.socket_arc();
622
623 if self.config.is_initiator() {
624 let peer_pubkey = self
626 .config
627 .peer_static_pubkey
628 .as_ref()
629 .ok_or_else(|| AdapterError::Fatal("missing peer public key".into()))?;
630
631 let mut handshake = NoiseHandshake::initiator(&self.config.psk, peer_pubkey)
632 .map_err(|e| AdapterError::Fatal(format!("handshake init failed: {}", e)))?;
633
634 let msg1 = handshake
636 .write_message(&[])
637 .map_err(|e| AdapterError::Connection(format!("write_message failed: {}", e)))?;
638
639 let mut builder = PacketBuilder::new(&[0u8; 32], 0);
640 let packet = builder.build_handshake(&msg1);
641
642 socket
643 .send_to(&packet, self.config.peer_addr)
644 .await
645 .map_err(|e| AdapterError::Connection(format!("send failed: {}", e)))?;
646
647 let (parsed, _source) = tokio::time::timeout(timeout, async {
651 let mut recv_buf = [0u8; protocol::MAX_PACKET_SIZE];
664 loop {
665 let (n, source) = socket_arc
666 .recv_from(&mut recv_buf)
667 .await
668 .map_err(|e| AdapterError::Connection(format!("recv failed: {}", e)))?;
669
670 if source != self.config.peer_addr {
672 continue;
673 }
674
675 let data = bytes::Bytes::copy_from_slice(&recv_buf[..n]);
676
677 if let Some(p) = ParsedPacket::parse(data, source) {
678 if p.header.flags.is_handshake() {
679 return Ok::<_, AdapterError>((p, source));
680 }
681 }
682 }
684 })
685 .await
686 .map_err(|_| AdapterError::Connection("handshake timeout".into()))??;
687
688 handshake
690 .read_message(&parsed.payload)
691 .map_err(|e| AdapterError::Connection(format!("read_message failed: {}", e)))?;
692
693 let keys = handshake
695 .into_session_keys()
696 .map_err(|e| AdapterError::Fatal(format!("key extraction failed: {}", e)))?;
697 Ok((keys, self.config.peer_addr))
698 } else {
699 let keypair = self
701 .config
702 .static_keypair
703 .as_ref()
704 .ok_or_else(|| AdapterError::Fatal("missing static keypair".into()))?;
705
706 let (parsed, source) = tokio::time::timeout(timeout, async {
713 loop {
714 let mut recv_buf = bytes::BytesMut::with_capacity(protocol::MAX_PACKET_SIZE);
715 recv_buf.resize(protocol::MAX_PACKET_SIZE, 0);
716
717 let (n, source) = socket_arc
718 .recv_from(&mut recv_buf)
719 .await
720 .map_err(|e| AdapterError::Connection(format!("recv failed: {}", e)))?;
721
722 recv_buf.truncate(n);
723 let data = recv_buf.freeze();
724
725 if let Some(p) = ParsedPacket::parse(data, source) {
726 if p.header.flags.is_handshake() {
727 let allowed = self.handshake_pacer.lock().check_and_record(source);
730 if !allowed {
731 tracing::debug!(
732 %source,
733 "handshake responder: dropping packet from \
734 rate-limited source"
735 );
736 continue;
737 }
738 return Ok::<_, AdapterError>((p, source));
739 }
740 }
741 }
743 })
744 .await
745 .map_err(|_| AdapterError::Connection("handshake timeout".into()))??;
746
747 let mut handshake = NoiseHandshake::responder(&self.config.psk, keypair)
748 .map_err(|e| AdapterError::Fatal(format!("handshake init failed: {}", e)))?;
749
750 handshake
752 .read_message(&parsed.payload)
753 .map_err(|e| AdapterError::Connection(format!("read_message failed: {}", e)))?;
754
755 let msg2 = handshake
757 .write_message(&[])
758 .map_err(|e| AdapterError::Connection(format!("write_message failed: {}", e)))?;
759
760 let mut builder = PacketBuilder::new(&[0u8; 32], 0);
761 let packet = builder.build_handshake(&msg2);
762
763 socket
766 .send_to(&packet, source)
767 .await
768 .map_err(|e| AdapterError::Connection(format!("send failed: {}", e)))?;
769
770 let keys = handshake
772 .into_session_keys()
773 .map_err(|e| AdapterError::Fatal(format!("key extraction failed: {}", e)))?;
774 Ok((keys, source))
775 }
776 }
777
778 fn process_packet(
780 data: Bytes,
781 source: std::net::SocketAddr,
782 session: &NetSession,
783 inbound: &InboundQueues,
784 num_shards: u16,
785 ) {
786 let mut parsed = match ParsedPacket::parse(data, source) {
788 Some(p) => p,
789 None => return,
790 };
791
792 if !parsed.header.flags.is_handshake()
796 && !parsed.header.flags.is_heartbeat()
797 && !parsed.is_valid_length()
798 {
799 return;
800 }
801
802 if parsed.header.flags.is_handshake() {
804 return;
805 }
806
807 if parsed.header.session_id != session.session_id() {
809 return;
810 }
811
812 if parsed.header.flags.is_heartbeat() {
829 if source == session.peer_addr() {
830 session.verify_and_touch_heartbeat(&parsed);
831 }
832 return;
833 }
834
835 let aad = parsed.header.aad();
840 let counter = u64::from_le_bytes(parsed.header.nonce[4..12].try_into().unwrap_or([0u8; 8]));
841 let rx_cipher = session.rx_cipher();
842 let payload = std::mem::take(&mut parsed.payload);
843 let decrypted = match rx_cipher.decrypt_to_bytes(counter, &aad, payload) {
851 Ok(d) => {
852 if !rx_cipher.try_admit_rx_counter(counter) {
853 return;
854 }
855 d
856 }
857 Err(_) => return,
858 };
859
860 let events = EventFrame::read_events(decrypted, parsed.header.event_count);
862
863 let stream_id = parsed.header.stream_id;
865 let shard_id = if num_shards > 0 {
866 (stream_id % num_shards as u64) as u16
867 } else {
868 0
869 };
870
871 let is_fresh = {
884 let stream = session.get_or_create_stream(stream_id);
885 let fresh = stream.with_reliability(|r| r.on_receive(parsed.header.sequence));
891 stream.update_rx_seq(parsed.header.sequence);
892 fresh
893 };
894
895 if is_fresh {
896 let queue = inbound.entry(shard_id).or_default();
898 let seq = parsed.header.sequence;
899 for (i, event_data) in events.into_iter().enumerate() {
900 use std::fmt::Write;
901 let mut event_id = String::with_capacity(24);
902 let _ = write!(event_id, "{}:{}", seq, i);
903 queue.push(StoredEvent::new(event_id, event_data, seq, shard_id));
904 }
905 } else {
906 tracing::debug!(
907 seq = parsed.header.sequence,
908 stream_id,
909 "Dropping duplicate packet"
910 );
911 }
912
913 session.touch();
914 }
915
916 #[cfg(target_os = "linux")]
927 fn spawn_receiver(
928 shutdown: Arc<AtomicBool>,
929 shutdown_notify: Arc<Notify>,
930 socket: Arc<Socket>,
931 session: Arc<NetSession>,
932 inbound: InboundQueues,
933 num_shards: u16,
934 ) -> JoinHandle<()> {
935 let mut receiver = transport::BatchedPacketReceiver::new(socket.socket_arc());
936
937 tokio::spawn(async move {
938 while !shutdown.load(Ordering::Acquire) {
939 tokio::select! {
940 result = receiver.recv() => {
941 match result {
942 Ok((data, source)) => {
943 Self::process_packet(data, source, &session, &inbound, num_shards);
944 }
945 Err(e) if e.kind() == std::io::ErrorKind::ConnectionReset => {
946 tracing::warn!("batch receiver thread exited, stopping receiver");
947 break;
948 }
949 Err(e) => {
950 if !shutdown.load(Ordering::Acquire) {
951 tracing::warn!(error = %e, "receive error");
952 }
953 }
954 }
955 }
956 _ = shutdown_notify.notified() => {
957 break;
958 }
959 }
960 }
961 })
962 }
963
964 #[cfg(not(target_os = "linux"))]
966 fn spawn_receiver(
967 shutdown: Arc<AtomicBool>,
968 shutdown_notify: Arc<Notify>,
969 socket: Arc<Socket>,
970 session: Arc<NetSession>,
971 inbound: InboundQueues,
972 num_shards: u16,
973 ) -> JoinHandle<()> {
974 tokio::spawn(async move {
975 let mut receiver = PacketReceiver::new(socket.socket_arc());
976
977 while !shutdown.load(Ordering::Acquire) {
978 tokio::select! {
982 result = receiver.recv() => {
983 match result {
984 Ok((data, source)) => {
985 Self::process_packet(data, source, &session, &inbound, num_shards);
986 }
987 Err(e) => {
988 if !shutdown.load(Ordering::Acquire) {
989 tracing::warn!(error = %e, "receive error");
990 }
991 }
992 }
993 }
994 _ = shutdown_notify.notified() => {
995 break;
996 }
997 }
998 }
999 })
1000 }
1001
1002 fn spawn_heartbeat(
1004 shutdown: Arc<AtomicBool>,
1005 shutdown_notify: Arc<Notify>,
1006 socket: Arc<Socket>,
1007 session: Arc<NetSession>,
1008 interval: std::time::Duration,
1009 peer_addr: std::net::SocketAddr,
1010 ) -> JoinHandle<()> {
1011 tokio::spawn(async move {
1012 let mut ticker = tokio::time::interval(interval);
1013
1014 loop {
1015 tokio::select! {
1016 _ = ticker.tick() => {
1017 if shutdown.load(Ordering::Acquire) || !session.is_active() {
1018 break;
1019 }
1020
1021 let packet = session.build_heartbeat();
1035
1036 if let Err(e) = socket.send_to(&packet, peer_addr).await {
1037 tracing::warn!(error = %e, "heartbeat send failed");
1038 }
1039 }
1040 _ = shutdown_notify.notified() => {
1041 break;
1042 }
1043 }
1044 }
1045 })
1046 }
1047}
1048
1049#[async_trait]
1050impl Adapter for NetAdapter {
1051 async fn init(&mut self) -> Result<(), AdapterError> {
1052 if self.initialized.load(Ordering::Acquire) {
1053 return Ok(());
1054 }
1055
1056 let socket_config = match (
1058 self.config.socket_recv_buffer,
1059 self.config.socket_send_buffer,
1060 ) {
1061 (Some(recv), Some(send)) => transport::SocketBufferConfig {
1062 recv_buffer_size: recv,
1063 send_buffer_size: send,
1064 },
1065 _ => transport::SocketBufferConfig::default(),
1066 };
1067 let socket = Socket::with_config(self.config.bind_addr, socket_config)
1068 .await
1069 .map_err(|e| AdapterError::Connection(format!("socket creation failed: {}", e)))?;
1070
1071 let socket = Arc::new(socket);
1072 self.socket = Some(socket.clone());
1073
1074 let (keys, actual_peer) = self.perform_handshake(&socket).await?;
1076
1077 let session = Arc::new(NetSession::new(
1081 keys,
1082 actual_peer,
1083 self.config.packet_pool_size,
1084 self.config.default_reliability.is_reliable(),
1085 ));
1086 self.session = Some(session.clone());
1087
1088 self.session_manager.set_session_arc(session.clone());
1090
1091 let recv_task = Self::spawn_receiver(
1093 self.shutdown.clone(),
1094 self.shutdown_notify.clone(),
1095 socket.clone(),
1096 session.clone(),
1097 self.inbound.clone(),
1098 self.config.num_shards,
1099 );
1100
1101 let heartbeat_task = Self::spawn_heartbeat(
1102 self.shutdown.clone(),
1103 self.shutdown_notify.clone(),
1104 socket,
1105 session,
1106 self.config.heartbeat_interval,
1107 actual_peer,
1108 );
1109
1110 {
1111 let mut tasks = self.tasks.lock().await;
1112 tasks.push(recv_task);
1113 tasks.push(heartbeat_task);
1114 }
1115
1116 self.initialized.store(true, Ordering::Release);
1117
1118 tracing::info!(
1119 bind_addr = %self.config.bind_addr,
1120 peer_addr = %self.config.peer_addr,
1121 role = ?self.config.role,
1122 "Net adapter initialized"
1123 );
1124
1125 Ok(())
1126 }
1127
1128 async fn on_batch(&self, batch: std::sync::Arc<Batch>) -> Result<(), AdapterError> {
1129 let session = self
1130 .session
1131 .as_ref()
1132 .ok_or_else(|| AdapterError::Connection("not connected".into()))?;
1133
1134 let socket = self
1135 .socket
1136 .as_ref()
1137 .ok_or_else(|| AdapterError::Connection("socket not initialized".into()))?;
1138
1139 let stream_id = batch.shard_id as u64;
1140 let peer_addr = session.peer_addr();
1141
1142 let reliable = {
1146 let stream = session.get_or_create_stream(stream_id);
1147 stream.with_reliability(|r| r.needs_ack())
1148 };
1150
1151 let mut current_batch: Vec<Bytes> = Vec::with_capacity(64);
1153 let mut current_size = 0usize;
1154
1155 let pool = session.thread_local_pool();
1157 let mut builder = pool.get();
1158
1159 for event in &batch.events {
1160 let event_bytes = event.raw.clone();
1161 let frame_size = EventFrame::LEN_SIZE + event_bytes.len();
1162
1163 if current_size + frame_size > protocol::MAX_PAYLOAD_SIZE && !current_batch.is_empty() {
1165 let seq;
1167 {
1168 let stream = session.get_or_create_stream(stream_id);
1169 seq = stream.next_tx_seq();
1170 }
1171
1172 let flags = if reliable {
1173 PacketFlags::RELIABLE
1174 } else {
1175 PacketFlags::NONE
1176 };
1177
1178 let packet = builder.build(stream_id, seq, ¤t_batch, flags);
1179
1180 socket
1182 .send_to(&packet, peer_addr)
1183 .await
1184 .map_err(|e| AdapterError::Connection(format!("send failed: {}", e)))?;
1185
1186 if reliable {
1193 let descriptor = std::sync::Arc::new(reliability::RetransmitDescriptor {
1198 seq,
1199 stream_id,
1200 events: current_batch.clone(),
1201 flags,
1202 });
1203 let stream = session.get_or_create_stream(stream_id);
1204 stream.with_reliability(|r| r.on_send(descriptor));
1205 }
1206
1207 current_batch.clear();
1208 current_size = 0;
1209 }
1210
1211 current_batch.push(event_bytes);
1212 current_size += frame_size;
1213 }
1214
1215 if !current_batch.is_empty() {
1217 let seq;
1218 {
1219 let stream = session.get_or_create_stream(stream_id);
1220 seq = stream.next_tx_seq();
1221 }
1222
1223 let flags = if reliable {
1224 PacketFlags::RELIABLE
1225 } else {
1226 PacketFlags::NONE
1227 };
1228
1229 let packet = builder.build(stream_id, seq, ¤t_batch, flags);
1230
1231 socket
1232 .send_to(&packet, peer_addr)
1233 .await
1234 .map_err(|e| AdapterError::Connection(format!("send failed: {}", e)))?;
1235
1236 if reliable {
1237 let descriptor = std::sync::Arc::new(reliability::RetransmitDescriptor {
1239 seq,
1240 stream_id,
1241 events: current_batch.clone(),
1242 flags,
1243 });
1244 let stream = session.get_or_create_stream(stream_id);
1245 stream.with_reliability(|r| r.on_send(descriptor));
1246 }
1247 }
1248
1249 session.touch();
1250
1251 Ok(())
1252 }
1253
1254 async fn poll_shard(
1255 &self,
1256 shard_id: u16,
1257 from_id: Option<&str>,
1258 limit: usize,
1259 ) -> Result<ShardPollResult, AdapterError> {
1260 let mut events = Vec::with_capacity(limit);
1261
1262 if let Some(queue) = self.inbound.get(&shard_id) {
1263 while events.len() < limit {
1264 if let Some(event) = queue.pop() {
1265 if from_id.is_none() || event_id_gt(&event.id, from_id.unwrap_or("")) {
1266 events.push(event);
1267 }
1268 } else {
1273 break;
1274 }
1275 }
1276 }
1277
1278 let has_more = self
1279 .inbound
1280 .get(&shard_id)
1281 .map(|q| !q.is_empty())
1282 .unwrap_or(false);
1283 let next_id = events.last().map(|e| e.id.clone());
1284
1285 Ok(ShardPollResult {
1286 events,
1287 next_id,
1288 has_more,
1289 })
1290 }
1291
1292 async fn flush(&self) -> Result<(), AdapterError> {
1293 Ok(())
1296 }
1297
1298 async fn shutdown(&self) -> Result<(), AdapterError> {
1299 self.shutdown.store(true, Ordering::Release);
1300
1301 self.shutdown_notify.notify_waiters();
1304
1305 self.session_manager.clear_session();
1307
1308 let mut tasks = self.tasks.lock().await;
1310 for task in tasks.drain(..) {
1311 let _ = task.await;
1312 }
1313
1314 self.initialized.store(false, Ordering::Release);
1315
1316 tracing::info!("Net adapter shutdown complete");
1317
1318 Ok(())
1319 }
1320
1321 fn name(&self) -> &'static str {
1322 "net"
1323 }
1324
1325 async fn is_healthy(&self) -> bool {
1326 self.initialized.load(Ordering::Acquire) && self.session_manager.check_session()
1327 }
1328}
1329
1330impl std::fmt::Debug for NetAdapter {
1331 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1332 f.debug_struct("NetAdapter")
1333 .field("config", &self.config)
1334 .field("initialized", &self.initialized.load(Ordering::Relaxed))
1335 .finish()
1336 }
1337}
1338
1339fn event_id_gt(a: &str, b: &str) -> bool {
1346 fn parse_id(id: &str) -> Option<(u64, u64)> {
1347 let (seq, idx) = id.split_once(':')?;
1348 Some((seq.parse().ok()?, idx.parse().ok()?))
1349 }
1350
1351 match (parse_id(a), parse_id(b)) {
1352 (Some(a), Some(b)) => a > b,
1353 _ => a > b, }
1355}
1356
1357#[cfg(test)]
1358mod tests {
1359 use super::*;
1360
1361 #[test]
1362 fn test_adapter_creation() {
1363 let psk = [0x42u8; 32];
1364 let peer_pubkey = [0x24u8; 32];
1365
1366 let config = NetAdapterConfig::initiator(
1367 "127.0.0.1:0".parse().unwrap(),
1368 "127.0.0.1:9999".parse().unwrap(),
1369 psk,
1370 peer_pubkey,
1371 );
1372
1373 let adapter = NetAdapter::new(config).unwrap();
1374 assert_eq!(adapter.name(), "net");
1375 }
1376
1377 #[test]
1384 fn coarse_clock_reuses_cache_within_refresh_window() {
1385 let t0 = std::time::Instant::now();
1386 let within = t0 + std::time::Duration::from_nanos(COARSE_CLOCK_REFRESH_NS - 1);
1387 let (store, ns) = coarse_clock_advance(Some((t0, 42)), within, || {
1388 panic!("cache hit must not read the wall clock")
1389 });
1390 assert_eq!(ns, 42, "hit must return the cached reading");
1391 assert!(store.is_none(), "hit must keep the hit path store-free");
1392 }
1393
1394 #[test]
1398 fn coarse_clock_refreshes_at_and_past_the_window() {
1399 let t0 = std::time::Instant::now();
1400 let at_boundary = t0 + std::time::Duration::from_nanos(COARSE_CLOCK_REFRESH_NS);
1401 let (store, ns) = coarse_clock_advance(Some((t0, 42)), at_boundary, || 100);
1402 assert_eq!(ns, 100, "boundary read must refresh");
1403 assert_eq!(
1404 store,
1405 Some((at_boundary, 100)),
1406 "refresh must rebase the window on the read instant"
1407 );
1408 }
1409
1410 #[test]
1413 fn coarse_clock_cold_start_reads_wall_clock() {
1414 let t0 = std::time::Instant::now();
1415 let (store, ns) = coarse_clock_advance(None, t0, || 7);
1416 assert_eq!(ns, 7);
1417 assert_eq!(store, Some((t0, 7)));
1418 }
1419
1420 #[test]
1421 fn current_timestamp_advances_after_refresh_interval() {
1422 let first = current_timestamp();
1427 std::thread::sleep(std::time::Duration::from_millis(5));
1428 let later = current_timestamp();
1429 assert!(
1430 later > first,
1431 "post-refresh reading must advance: first={}, later={}",
1432 first,
1433 later
1434 );
1435 }
1436
1437 #[test]
1438 fn test_shard_id_from_stream_id_uses_modulo() {
1439 let num_shards: u16 = 8;
1443
1444 let stream_a: u64 = 0xDEAD_BEEF_0000_0003;
1449 let stream_b: u64 = 0xCAFE_BABE_0000_0003;
1450
1451 let shard_a = (stream_a % num_shards as u64) as u16;
1452 let shard_b = (stream_b % num_shards as u64) as u16;
1453
1454 assert!(
1455 shard_a < num_shards,
1456 "shard must be in range [0, num_shards)"
1457 );
1458 assert!(
1459 shard_b < num_shards,
1460 "shard must be in range [0, num_shards)"
1461 );
1462
1463 let big_stream: u64 = 0xFFFF_FFFF_FFFF_FFFF;
1465 let shard_big = (big_stream % num_shards as u64) as u16;
1466 assert!(shard_big < num_shards);
1467
1468 assert_ne!(
1471 big_stream as u16, shard_big,
1472 "modulo must differ from truncation for large stream IDs"
1473 );
1474 }
1475
1476 #[test]
1477 fn test_invalid_config() {
1478 let psk = [0x42u8; 32];
1479 let peer_pubkey = [0x24u8; 32];
1480
1481 let mut config = NetAdapterConfig::initiator(
1482 "127.0.0.1:0".parse().unwrap(),
1483 "127.0.0.1:9999".parse().unwrap(),
1484 psk,
1485 peer_pubkey,
1486 );
1487 config.peer_static_pubkey = None;
1488
1489 let result = NetAdapter::new(config);
1490 assert!(result.is_err());
1491 }
1492
1493 #[test]
1496 fn test_event_id_gt_numeric_ordering() {
1497 assert!(event_id_gt("2:0", "1:0"));
1499 assert!(!event_id_gt("1:0", "2:0"));
1500 assert!(!event_id_gt("1:0", "1:0"));
1501
1502 assert!(event_id_gt("10:0", "9:0"));
1504 assert!(event_id_gt("100:0", "99:0"));
1505 assert!(!event_id_gt("9:0", "10:0"));
1506
1507 assert!(event_id_gt("5:2", "5:1"));
1509 assert!(!event_id_gt("5:1", "5:2"));
1510
1511 assert!(event_id_gt("1000000:0", "999999:0"));
1513 }
1514
1515 #[test]
1521 fn test_event_id_gt_edge_cases() {
1522 assert!(event_id_gt("1:0", ""));
1524 assert!(event_id_gt("b", "a"));
1526 assert!(!event_id_gt("a", "b"));
1527 }
1528
1529 #[test]
1534 fn test_build_then_process_packet_roundtrip() {
1535 use crate::adapter::net::crypto::{NoiseHandshake, StaticKeypair};
1536 use dashmap::DashMap;
1537 use std::sync::Arc;
1538
1539 let psk = [0x42u8; 32];
1541 let responder_kp = StaticKeypair::generate();
1542
1543 let mut initiator = NoiseHandshake::initiator(&psk, &responder_kp.public).unwrap();
1544 let mut responder = NoiseHandshake::responder(&psk, &responder_kp).unwrap();
1545
1546 let msg1 = initiator.write_message(&[]).unwrap();
1547 responder.read_message(&msg1).unwrap();
1548 let msg2 = responder.write_message(&[]).unwrap();
1549 initiator.read_message(&msg2).unwrap();
1550
1551 let init_keys = initiator.into_session_keys().unwrap();
1552 let resp_keys = responder.into_session_keys().unwrap();
1553
1554 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1556 let events = vec![
1557 Bytes::from(r#"{"token":"hello"}"#),
1558 Bytes::from(r#"{"token":"world"}"#),
1559 ];
1560 let packet = builder.build(0, 0, &events, PacketFlags::NONE);
1561
1562 let resp_session = Arc::new(NetSession::new(
1564 resp_keys,
1565 "127.0.0.1:5000".parse().unwrap(),
1566 4,
1567 false,
1568 ));
1569 let inbound: InboundQueues = Arc::new(DashMap::new());
1570 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1571
1572 NetAdapter::process_packet(packet, source, &resp_session, &inbound, 1);
1573
1574 let queue = inbound.get(&0).expect("shard 0 should have events");
1576 assert_eq!(queue.len(), 2, "expected 2 events, got {}", queue.len());
1577
1578 let e1 = queue.pop().unwrap();
1579 assert_eq!(&e1.raw[..], br#"{"token":"hello"}"#);
1580
1581 let e2 = queue.pop().unwrap();
1582 assert_eq!(&e2.raw[..], br#"{"token":"world"}"#);
1583 }
1584
1585 fn make_session_keys() -> (SessionKeys, SessionKeys) {
1587 use crate::adapter::net::crypto::{NoiseHandshake, StaticKeypair};
1588
1589 let psk = [0x42u8; 32];
1590 let responder_kp = StaticKeypair::generate();
1591
1592 let mut initiator = NoiseHandshake::initiator(&psk, &responder_kp.public).unwrap();
1593 let mut responder = NoiseHandshake::responder(&psk, &responder_kp).unwrap();
1594
1595 let msg1 = initiator.write_message(&[]).unwrap();
1596 responder.read_message(&msg1).unwrap();
1597 let msg2 = responder.write_message(&[]).unwrap();
1598 initiator.read_message(&msg2).unwrap();
1599
1600 (
1601 initiator.into_session_keys().unwrap(),
1602 responder.into_session_keys().unwrap(),
1603 )
1604 }
1605
1606 #[test]
1607 fn test_process_packet_rejects_truncated_packet() {
1608 use dashmap::DashMap;
1609 use std::sync::Arc;
1610
1611 let (init_keys, resp_keys) = make_session_keys();
1612
1613 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1615 let packet = builder.build(0, 0, &[Bytes::from_static(b"hello")], PacketFlags::NONE);
1616
1617 let resp_session = Arc::new(NetSession::new(
1618 resp_keys,
1619 "127.0.0.1:5000".parse().unwrap(),
1620 4,
1621 false,
1622 ));
1623 let inbound: InboundQueues = Arc::new(DashMap::new());
1624 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1625
1626 let truncated = packet.slice(..packet.len() - 10);
1628 NetAdapter::process_packet(truncated, source, &resp_session, &inbound, 1);
1629 assert!(
1630 inbound.get(&0).is_none() || inbound.get(&0).unwrap().is_empty(),
1631 "truncated packet must be silently dropped"
1632 );
1633 }
1634
1635 #[test]
1636 fn test_process_packet_rejects_tampered_payload() {
1637 use dashmap::DashMap;
1638 use std::sync::Arc;
1639
1640 let (init_keys, resp_keys) = make_session_keys();
1641
1642 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1643 let packet = builder.build(0, 0, &[Bytes::from_static(b"hello")], PacketFlags::NONE);
1644
1645 let resp_session = Arc::new(NetSession::new(
1646 resp_keys,
1647 "127.0.0.1:5000".parse().unwrap(),
1648 4,
1649 false,
1650 ));
1651 let inbound: InboundQueues = Arc::new(DashMap::new());
1652 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1653
1654 let mut tampered = bytes::BytesMut::from(&packet[..]);
1656 tampered[super::protocol::HEADER_SIZE + 2] ^= 0xFF;
1657 NetAdapter::process_packet(tampered.freeze(), source, &resp_session, &inbound, 1);
1658
1659 assert!(
1660 inbound.get(&0).is_none() || inbound.get(&0).unwrap().is_empty(),
1661 "tampered packet must be rejected by AEAD"
1662 );
1663 }
1664
1665 #[test]
1666 fn test_process_packet_rejects_wrong_session_id() {
1667 use dashmap::DashMap;
1668 use std::sync::Arc;
1669
1670 let (init_keys, resp_keys) = make_session_keys();
1671
1672 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1673 let packet = builder.build(0, 0, &[Bytes::from_static(b"hello")], PacketFlags::NONE);
1674
1675 let mut wrong_keys = resp_keys;
1677 wrong_keys.session_id = 0xDEAD;
1678 let resp_session = Arc::new(NetSession::new(
1679 wrong_keys,
1680 "127.0.0.1:5000".parse().unwrap(),
1681 4,
1682 false,
1683 ));
1684 let inbound: InboundQueues = Arc::new(DashMap::new());
1685 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1686
1687 NetAdapter::process_packet(packet, source, &resp_session, &inbound, 1);
1688
1689 assert!(
1690 inbound.get(&0).is_none() || inbound.get(&0).unwrap().is_empty(),
1691 "packet with wrong session_id must be dropped"
1692 );
1693 }
1694
1695 #[test]
1696 fn test_process_packet_multi_packet_batch_all_events_arrive() {
1697 use dashmap::DashMap;
1698 use std::sync::Arc;
1699
1700 let (init_keys, resp_keys) = make_session_keys();
1701
1702 let resp_session = Arc::new(NetSession::new(
1703 resp_keys,
1704 "127.0.0.1:5000".parse().unwrap(),
1705 4,
1706 false,
1707 ));
1708 let inbound: InboundQueues = Arc::new(DashMap::new());
1709 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1710
1711 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1715 let total_events = 200;
1716 let mut seq = 0u64;
1717
1718 let mut current_batch: Vec<Bytes> = Vec::new();
1720 let mut current_size = 0;
1721
1722 for i in 0..total_events {
1723 let data = format!("{{\"i\":{},\"pad\":\"{}\"}}", i, "x".repeat(150));
1724 let event_bytes = Bytes::from(data);
1725 let frame_size = EventFrame::LEN_SIZE + event_bytes.len();
1726
1727 if current_size + frame_size > protocol::MAX_PAYLOAD_SIZE && !current_batch.is_empty() {
1728 let packet = builder.build(0, seq, ¤t_batch, PacketFlags::NONE);
1729 NetAdapter::process_packet(packet, source, &resp_session, &inbound, 1);
1730 seq += 1;
1731 current_batch.clear();
1732 current_size = 0;
1733 }
1734
1735 current_batch.push(event_bytes);
1736 current_size += frame_size;
1737 }
1738
1739 if !current_batch.is_empty() {
1740 let packet = builder.build(0, seq, ¤t_batch, PacketFlags::NONE);
1741 NetAdapter::process_packet(packet, source, &resp_session, &inbound, 1);
1742 }
1743
1744 let queue = inbound.get(&0).expect("shard 0 should have events");
1746 assert_eq!(
1747 queue.len(),
1748 total_events,
1749 "all {} events must arrive across multiple packets",
1750 total_events
1751 );
1752 }
1753
1754 #[test]
1755 fn test_build_then_process_packet_both_directions() {
1756 use dashmap::DashMap;
1757 use std::sync::Arc;
1758
1759 let (init_keys, resp_keys) = make_session_keys();
1760 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1761
1762 {
1764 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1765 let packet = builder.build(0, 0, &[Bytes::from_static(b"i2r")], PacketFlags::NONE);
1766
1767 let session = Arc::new(NetSession::new(resp_keys.clone(), source, 4, false));
1768 let inbound: InboundQueues = Arc::new(DashMap::new());
1769 NetAdapter::process_packet(packet, source, &session, &inbound, 1);
1770
1771 let queue = inbound.get(&0).expect("i2r: shard 0 should have events");
1772 assert_eq!(queue.len(), 1, "i2r: expected 1 event");
1773 assert_eq!(&queue.pop().unwrap().raw[..], b"i2r");
1774 }
1775
1776 {
1778 let mut builder = PacketBuilder::new(&resp_keys.tx_key, resp_keys.session_id);
1779 let packet = builder.build(0, 0, &[Bytes::from_static(b"r2i")], PacketFlags::NONE);
1780
1781 let session = Arc::new(NetSession::new(init_keys.clone(), source, 4, false));
1782 let inbound: InboundQueues = Arc::new(DashMap::new());
1783 NetAdapter::process_packet(packet, source, &session, &inbound, 1);
1784
1785 let queue = inbound.get(&0).expect("r2i: shard 0 should have events");
1786 assert_eq!(queue.len(), 1, "r2i: expected 1 event");
1787 assert_eq!(&queue.pop().unwrap().raw[..], b"r2i");
1788 }
1789 }
1790
1791 #[test]
1792 fn test_poll_shard_cursor_drops_consumed_events() {
1793 use std::sync::Arc;
1798
1799 let (init_keys, resp_keys) = make_session_keys();
1800
1801 let resp_session = Arc::new(NetSession::new(
1802 resp_keys,
1803 "127.0.0.1:5000".parse().unwrap(),
1804 4,
1805 false,
1806 ));
1807 let inbound: InboundQueues = Arc::new(DashMap::new());
1808 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1809
1810 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1812 for seq in 0..3u64 {
1813 let events = vec![Bytes::from(format!("event-{}", seq))];
1814 let packet = builder.build(0, seq, &events, PacketFlags::NONE);
1815 NetAdapter::process_packet(packet, source, &resp_session, &inbound, 1);
1816 }
1817
1818 let queue = inbound.get(&0u16).unwrap();
1819 assert_eq!(queue.len(), 3);
1820
1821 let from_id = "0:0";
1824 let mut events = Vec::new();
1825 while events.len() < 10 {
1826 if let Some(event) = queue.pop() {
1827 if event_id_gt(&event.id, from_id) {
1828 events.push(event);
1829 }
1830 } else {
1832 break;
1833 }
1834 }
1835
1836 assert_eq!(events.len(), 2, "should get 2 events after cursor 0:0");
1837 assert_eq!(events[0].id, "1:0");
1838 assert_eq!(events[1].id, "2:0");
1839
1840 assert_eq!(queue.len(), 0, "queue should be empty after poll drains it");
1842 }
1843
1844 #[test]
1845 fn test_process_packet_old_counter_rejected() {
1846 use std::sync::Arc;
1849
1850 let (init_keys, resp_keys) = make_session_keys();
1851 let resp_session = Arc::new(NetSession::new(
1852 resp_keys,
1853 "127.0.0.1:5000".parse().unwrap(),
1854 4,
1855 false,
1856 ));
1857 let inbound: InboundQueues = Arc::new(DashMap::new());
1858 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1859
1860 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1862 for seq in 0..1100u64 {
1863 let packet = builder.build(0, seq, &[Bytes::from_static(b"x")], PacketFlags::NONE);
1864 NetAdapter::process_packet(packet, source, &resp_session, &inbound, 1);
1865 }
1866 assert_eq!(inbound.get(&0).unwrap().len(), 1100);
1867
1868 let mut stale_builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1872 let stale_packet =
1873 stale_builder.build(0, 9999, &[Bytes::from_static(b"stale")], PacketFlags::NONE);
1874 NetAdapter::process_packet(stale_packet, source, &resp_session, &inbound, 1);
1875
1876 assert_eq!(
1878 inbound.get(&0).unwrap().len(),
1879 1100,
1880 "packet with stale counter must be rejected"
1881 );
1882 }
1883
1884 #[test]
1885 fn test_process_packet_far_future_counter_rejected() {
1886 use std::sync::Arc;
1890
1891 let (_init_keys, resp_keys) = make_session_keys();
1892
1893 let resp_session = Arc::new(NetSession::new(
1898 resp_keys,
1899 "127.0.0.1:5000".parse().unwrap(),
1900 4,
1901 false,
1902 ));
1903
1904 let rx_cipher = resp_session.rx_cipher();
1906 assert!(
1907 !rx_cipher.is_valid_rx_counter(u64::MAX),
1908 "counter at u64::MAX must be rejected (far beyond MAX_FORWARD)"
1909 );
1910 assert!(
1911 rx_cipher.is_valid_rx_counter(0),
1912 "counter 0 should be valid initially"
1913 );
1914 }
1915
1916 #[test]
1933 fn process_packet_drops_duplicates_per_reliability_decision() {
1934 use dashmap::DashMap;
1935 use std::sync::Arc;
1936
1937 let (init_keys, resp_keys) = make_session_keys();
1938
1939 let resp_session = Arc::new(NetSession::new(
1943 resp_keys,
1944 "127.0.0.1:5000".parse().unwrap(),
1945 4,
1946 true, ));
1948 let inbound: InboundQueues = Arc::new(DashMap::new());
1949 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
1950
1951 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
1957 let packet0 = builder.build(7, 0, &[Bytes::from(r#"{"first":0}"#)], PacketFlags::NONE);
1958 let packet1 = builder.build(7, 1, &[Bytes::from(r#"{"first":1}"#)], PacketFlags::NONE);
1959 let packet0_dup = builder.build(
1963 7,
1964 0,
1965 &[Bytes::from(r#"{"dup":"should_not_appear"}"#)],
1966 PacketFlags::NONE,
1967 );
1968
1969 NetAdapter::process_packet(packet0, source, &resp_session, &inbound, 1);
1970 NetAdapter::process_packet(packet1, source, &resp_session, &inbound, 1);
1971 NetAdapter::process_packet(packet0_dup, source, &resp_session, &inbound, 1);
1972
1973 let queue = inbound.get(&0).expect("shard 0 should exist");
1974 assert_eq!(
1975 queue.len(),
1976 2,
1977 "duplicate packet must NOT enqueue (BUG_REPORT.md #5); \
1978 got {} events, expected exactly 2 (seq=0 and seq=1, no dup)",
1979 queue.len()
1980 );
1981
1982 let e0 = queue.pop().unwrap();
1985 assert_eq!(&e0.raw[..], br#"{"first":0}"#);
1986 let e1 = queue.pop().unwrap();
1987 assert_eq!(&e1.raw[..], br#"{"first":1}"#);
1988 assert!(queue.is_empty());
1989 }
1990
1991 #[test]
1999 fn heartbeat_is_aead_authenticated() {
2000 use crate::adapter::net::pool::PacketBuilder;
2001 use dashmap::DashMap;
2002 use std::sync::Arc;
2003
2004 let (init_keys, resp_keys) = make_session_keys();
2005
2006 let resp_session = Arc::new(NetSession::new(
2007 resp_keys,
2008 "127.0.0.1:5000".parse().unwrap(),
2009 4,
2010 false,
2011 ));
2012 let inbound: InboundQueues = Arc::new(DashMap::new());
2013 let source: std::net::SocketAddr = "127.0.0.1:5000".parse().unwrap();
2014
2015 let mut builder = PacketBuilder::new(&init_keys.tx_key, init_keys.session_id);
2018 let heartbeat = builder.build_heartbeat();
2019 let last_activity_before = resp_session.last_activity_ns();
2020 std::thread::sleep(std::time::Duration::from_millis(2));
2021
2022 NetAdapter::process_packet(heartbeat, source, &resp_session, &inbound, 1);
2024 let last_activity_after = resp_session.last_activity_ns();
2025 assert!(
2026 last_activity_after > last_activity_before,
2027 "legitimate AEAD-tagged heartbeat must call session.touch()"
2028 );
2029
2030 let mut forged = bytes::BytesMut::new();
2034 let header = NetHeader::heartbeat(resp_session.session_id());
2035 forged.extend_from_slice(&header.to_bytes());
2036 let forged = forged.freeze();
2037 let last_activity_before = resp_session.last_activity_ns();
2038 std::thread::sleep(std::time::Duration::from_millis(2));
2039 NetAdapter::process_packet(forged, source, &resp_session, &inbound, 1);
2040 let last_activity_after = resp_session.last_activity_ns();
2041 assert_eq!(
2042 last_activity_before, last_activity_after,
2043 "unauthenticated heartbeat (no AEAD tag) must NOT touch the session"
2044 );
2045
2046 let mut forged_tag = bytes::BytesMut::new();
2049 let mut header_bytes = NetHeader::heartbeat(resp_session.session_id()).to_bytes();
2050 header_bytes[12..16].copy_from_slice(&[0u8; 4]);
2053 header_bytes[16..24].copy_from_slice(&1u64.to_le_bytes());
2054 forged_tag.extend_from_slice(&header_bytes);
2055 forged_tag.extend_from_slice(&[0xAAu8; 16]); let forged_tag = forged_tag.freeze();
2057 let last_activity_before = resp_session.last_activity_ns();
2058 std::thread::sleep(std::time::Duration::from_millis(2));
2059 NetAdapter::process_packet(forged_tag, source, &resp_session, &inbound, 1);
2060 let last_activity_after = resp_session.last_activity_ns();
2061 assert_eq!(
2062 last_activity_before, last_activity_after,
2063 "heartbeat with garbage AEAD tag must NOT touch the session"
2064 );
2065 }
2066
2067 #[test]
2073 fn handshake_pacer_rejects_floods_per_source() {
2074 use std::time::Duration;
2075 let mut pacer = HandshakePacer::new(3, Duration::from_millis(50));
2076
2077 let attacker: std::net::SocketAddr = "10.0.0.1:9000".parse().unwrap();
2078 let legit: std::net::SocketAddr = "10.0.0.2:9000".parse().unwrap();
2079
2080 for _ in 0..3 {
2082 assert!(pacer.check_and_record(attacker));
2083 }
2084 for _ in 0..10 {
2086 assert!(
2087 !pacer.check_and_record(attacker),
2088 "attacker exceeding budget must be dropped"
2089 );
2090 }
2091
2092 assert!(
2095 pacer.check_and_record(legit),
2096 "legitimate source must still get through despite attacker flood"
2097 );
2098
2099 std::thread::sleep(Duration::from_millis(55));
2101 assert!(
2102 pacer.check_and_record(attacker),
2103 "attacker budget must refill after window"
2104 );
2105 }
2106}