1use crate::socket::error::{EncryptSendError, EncryptSendErrorKind, Result, SocketError};
2use crate::transport::Transport;
3use async_channel;
4use bytes::BytesMut;
5use futures::channel::oneshot;
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU32, Ordering};
8use wacore::handshake::{NoiseCipher, NoiseError};
9use wacore::libsignal::crypto::GcmInPlaceBuffer;
10use wacore::runtime::{AbortHandle, Runtime};
11
12const INLINE_ENCRYPT_THRESHOLD: usize = 16 * 1024;
13
14const TAG_LEN: usize = 16;
18
19struct FrameBody<'a> {
29 out: &'a mut BytesMut,
30 base: usize,
34}
35
36impl GcmInPlaceBuffer for FrameBody<'_> {
37 fn as_mut_slice(&mut self) -> &mut [u8] {
38 &mut self.out[self.base..]
39 }
40
41 fn as_slice(&self) -> &[u8] {
42 &self.out[self.base..]
43 }
44
45 fn resize(&mut self, new_len: usize, value: u8) {
46 self.out.resize(self.base + new_len, value);
47 }
48
49 fn truncate(&mut self, len: usize) {
50 self.out.truncate(self.base + len);
51 }
52}
53
54const MAX_BATCH_FRAMES: usize = 16;
58const MAX_BATCH_WIRE_BYTES: usize = 64 * 1024;
59
60type SendResult = std::result::Result<(), EncryptSendError>;
62
63fn frame_wire_len(plaintext_len: usize) -> usize {
67 plaintext_len + TAG_LEN + wacore::framing::FRAME_LENGTH_SIZE
68}
69
70#[derive(Debug)]
75struct SharedSendFailure(Arc<EncryptSendError>);
76
77impl std::fmt::Display for SharedSendFailure {
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 write!(f, "{}", self.0)
80 }
81}
82
83impl std::error::Error for SharedSendFailure {
84 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
85 std::error::Error::source(self.0.as_ref())
86 }
87}
88
89struct SendJob {
91 plaintext: bytes::Bytes,
92 response_tx: oneshot::Sender<SendResult>,
93}
94
95pub struct NoiseSocket {
96 read_key: Arc<NoiseCipher>,
97 read_counter: Arc<AtomicU32>,
98 send_job_tx: async_channel::Sender<SendJob>,
103 _sender_task_handle: AbortHandle,
106}
107
108impl NoiseSocket {
109 pub fn new(
110 runtime: Arc<dyn Runtime>,
111 transport: Arc<dyn Transport>,
112 write_key: NoiseCipher,
113 read_key: NoiseCipher,
114 ) -> Self {
115 Self::with_stats(runtime, transport, write_key, read_key, None)
116 }
117
118 pub fn with_stats(
122 runtime: Arc<dyn Runtime>,
123 transport: Arc<dyn Transport>,
124 write_key: NoiseCipher,
125 read_key: NoiseCipher,
126 stats: Option<Arc<wacore::stats::SessionStats>>,
127 ) -> Self {
128 let write_key = Arc::new(write_key);
129 let read_key = Arc::new(read_key);
130
131 let (send_job_tx, send_job_rx) = async_channel::bounded::<SendJob>(8);
135
136 let transport_clone = transport.clone();
138 let write_key_clone = write_key.clone();
139 let rt_clone = runtime.clone();
140 let sender_task_handle = runtime.spawn(Box::pin(Self::sender_task(
141 rt_clone,
142 transport_clone,
143 write_key_clone,
144 send_job_rx,
145 stats,
146 )));
147
148 Self {
149 read_key,
150 read_counter: Arc::new(AtomicU32::new(0)),
151 send_job_tx,
152 _sender_task_handle: sender_task_handle,
153 }
154 }
155
156 async fn sender_task(
160 runtime: Arc<dyn Runtime>,
161 transport: Arc<dyn Transport>,
162 write_key: Arc<NoiseCipher>,
163 send_job_rx: async_channel::Receiver<SendJob>,
164 stats: Option<Arc<wacore::stats::SessionStats>>,
165 ) {
166 let mut write_counter: u32 = 0;
167 let mut out_buf = BytesMut::with_capacity(4096);
170 let mut poisoned = false;
177 let mut waiters: Vec<(oneshot::Sender<SendResult>, usize)> = Vec::new();
180 let mut carry_over: Option<SendJob> = None;
185
186 loop {
187 let job = match carry_over.take() {
188 Some(job) => job,
189 None => match send_job_rx.recv().await {
190 Ok(job) => job,
191 Err(_) => break,
192 },
193 };
194 if poisoned {
195 let _ = job.response_tx.send(Err(EncryptSendError::poisoned()));
196 continue;
197 }
198
199 waiters.clear();
206 let mut encrypt_failure: Option<(oneshot::Sender<SendResult>, EncryptSendError)> = None;
207 let mut job = job;
208 loop {
209 let response_tx = job.response_tx;
210 match Self::encrypt_frame_into(
211 &runtime,
212 &write_key,
213 &mut write_counter,
214 job.plaintext,
215 &mut out_buf,
216 )
217 .await
218 {
219 Ok(wire_bytes) => waiters.push((response_tx, wire_bytes)),
220 Err(e) => {
221 encrypt_failure = Some((response_tx, e));
225 break;
226 }
227 }
228 if out_buf.len() >= MAX_BATCH_WIRE_BYTES || waiters.len() >= MAX_BATCH_FRAMES {
229 break;
230 }
231 match send_job_rx.try_recv() {
232 Ok(next) => {
233 let projected = out_buf.len() + frame_wire_len(next.plaintext.len());
238 if projected > MAX_BATCH_WIRE_BYTES {
239 carry_over = Some(next);
240 break;
241 }
242 job = next;
243 }
244 Err(_) => break,
245 }
246 }
247
248 let outcome = if out_buf.is_empty() {
249 Ok(())
250 } else {
251 let wire = out_buf.split().freeze();
254 if waiters.len() > 1 {
255 log::debug!(
259 "noise: coalesced {} frames into one {}-byte write",
260 waiters.len(),
261 wire.len()
262 );
263 }
264 match transport.send(wire).await {
265 Ok(()) => {
266 if let Some(stats) = stats.as_deref() {
267 for (_, wire_bytes) in &waiters {
268 stats.record_frame_sent(*wire_bytes);
269 }
270 }
271 Ok(())
272 }
273 Err(e) => Err(EncryptSendError::transport(e)),
274 }
275 };
276
277 {
278 if let Err(err) = &outcome
283 && matches!(err.kind, EncryptSendErrorKind::Transport)
284 {
285 poisoned = true;
286 transport.disconnect().await;
296 }
297 }
298
299 match outcome {
301 Ok(()) => {
302 for (response_tx, _) in waiters.drain(..) {
303 let _ = response_tx.send(Ok(()));
304 }
305 }
306 Err(err) if waiters.len() == 1 => {
313 let (response_tx, _) = waiters.drain(..).next().expect("length checked");
314 let _ = response_tx.send(Err(err));
315 }
316 Err(err) => {
321 let shared = Arc::new(err);
322 for (response_tx, _) in waiters.drain(..) {
323 let _ = response_tx.send(Err(EncryptSendError::transport(
324 SharedSendFailure(shared.clone()),
325 )));
326 }
327 }
328 }
329 if let Some((response_tx, err)) = encrypt_failure {
330 let _ = response_tx.send(Err(err));
331 }
332 }
333 }
334
335 async fn encrypt_frame_into(
343 runtime: &Arc<dyn Runtime>,
344 write_key: &Arc<NoiseCipher>,
345 write_counter: &mut u32,
346 plaintext: bytes::Bytes,
347 out_buf: &mut BytesMut,
348 ) -> std::result::Result<usize, EncryptSendError> {
349 let counter = *write_counter;
350 if counter == u32::MAX {
355 return Err(EncryptSendError::crypto(NoiseError::CounterExhausted));
356 }
357 let before = out_buf.len();
358
359 if plaintext.len() <= INLINE_ENCRYPT_THRESHOLD {
360 let body_len = plaintext.len() + TAG_LEN;
364 if let Err(e) = wacore::framing::append_frame_header_into(body_len, None, out_buf) {
365 return Err(EncryptSendError::framing(e));
366 }
367 let base = out_buf.len();
368 out_buf.extend_from_slice(&plaintext);
369 if let Err(e) = write_key
370 .encrypt_in_place_with_counter(counter, &mut FrameBody { out: out_buf, base })
371 {
372 out_buf.truncate(before);
380 return Err(EncryptSendError::crypto(e));
381 }
382 if out_buf.len() - base != body_len {
390 out_buf.truncate(before);
391 return Err(EncryptSendError::crypto(NoiseError::Encrypt(
392 wacore::libsignal::crypto::CryptoProviderError::BackendFailed,
393 )));
394 }
395 } else {
396 let write_key = write_key.clone();
397 let encrypt_result = wacore::runtime::blocking(&**runtime, move || {
400 write_key.encrypt_with_counter(counter, &plaintext)
401 })
402 .await;
403 let ciphertext = match encrypt_result {
404 Ok(c) => c,
405 Err(e) => return Err(EncryptSendError::crypto(e)),
406 };
407 if let Err(e) = wacore::framing::append_frame_into(&ciphertext, None, out_buf) {
408 return Err(EncryptSendError::framing(e));
409 }
410 }
411
412 *write_counter = counter + 1;
413 Ok(out_buf.len() - before)
414 }
415
416 pub(crate) async fn enqueue_send(
428 &self,
429 plaintext: bytes::Bytes,
430 ) -> std::result::Result<oneshot::Receiver<SendResult>, EncryptSendError> {
431 let (response_tx, response_rx) = oneshot::channel();
432
433 let job = SendJob {
434 plaintext,
435 response_tx,
436 };
437
438 if let Err(_send_err) = self.send_job_tx.send(job).await {
440 return Err(EncryptSendError::channel_closed());
441 }
442
443 Ok(response_rx)
444 }
445
446 pub(crate) async fn await_send(receiver: oneshot::Receiver<SendResult>) -> SendResult {
448 match receiver.await {
449 Ok(result) => result,
450 Err(_) => {
451 Err(EncryptSendError::channel_closed())
453 }
454 }
455 }
456
457 pub async fn encrypt_and_send(&self, plaintext: bytes::Bytes) -> SendResult {
458 let receiver = self.enqueue_send(plaintext).await?;
459 Self::await_send(receiver).await
460 }
461
462 pub fn decrypt_frame(&self, mut ciphertext: BytesMut) -> Result<BytesMut> {
463 let counter = self
467 .read_counter
468 .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |c| c.checked_add(1))
469 .map_err(|_| SocketError::Cipher(NoiseError::CounterExhausted))?;
470 self.read_key
471 .decrypt_in_place_with_counter(counter, &mut ciphertext)
472 .map_err(SocketError::Cipher)?;
473 Ok(ciphertext)
474 }
475}
476
477#[cfg(test)]
481mod tests {
482 use super::*;
483 use std::future::Future;
484 use std::sync::atomic::{AtomicBool, Ordering};
485 use wacore::framing::FRAME_LENGTH_SIZE;
486
487 #[tokio::test]
488 async fn test_encrypt_and_send_succeeds() {
489 let transport = Arc::new(crate::transport::mock::MockTransport);
490
491 let key = [0u8; 32];
492 let write_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
493 let read_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
494
495 let socket = NoiseSocket::new(
496 Arc::new(crate::runtime_impl::TokioRuntime),
497 transport,
498 write_key,
499 read_key,
500 );
501
502 let result = socket.encrypt_and_send(bytes::Bytes::new()).await;
503 assert!(result.is_ok(), "encrypt_and_send should succeed");
504 }
505
506 #[tokio::test]
507 async fn decrypt_frame_errors_on_counter_exhaustion() {
508 let key = [0u8; 32];
509 let socket = NoiseSocket::new(
510 Arc::new(crate::runtime_impl::TokioRuntime),
511 Arc::new(crate::transport::mock::MockTransport),
512 NoiseCipher::new(&key).expect("32-byte key"),
513 NoiseCipher::new(&key).expect("32-byte key"),
514 );
515 socket.read_counter.store(u32::MAX, Ordering::SeqCst);
518 let err = socket
519 .decrypt_frame(BytesMut::from(&b"ignored"[..]))
520 .expect_err("exhausted read counter must error, not wrap");
521 assert!(matches!(
522 err,
523 SocketError::Cipher(NoiseError::CounterExhausted)
524 ));
525 }
526
527 #[tokio::test]
532 async fn test_large_frame_round_trips_via_bytes_path() {
533 use async_lock::Mutex;
534 use async_trait::async_trait;
535 use std::sync::Arc;
536 use std::sync::atomic::{AtomicU32, Ordering};
537
538 struct CapturingTransport {
539 captured: Arc<Mutex<Vec<Vec<u8>>>>,
540 read_key: NoiseCipher,
541 counter: AtomicU32,
542 }
543
544 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
545 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
546 impl Transport for CapturingTransport {
547 async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
548 let mut data = data.to_vec();
549 data.drain(..3); let counter = self.counter.fetch_add(1, Ordering::SeqCst);
551 self.read_key
552 .decrypt_in_place_with_counter(counter, &mut data)
553 .expect("frame should decrypt");
554 self.captured.lock().await.push(data);
555 Ok(())
556 }
557 async fn disconnect(&self) {}
558 }
559
560 let captured = Arc::new(Mutex::new(Vec::new()));
561 let key = [7u8; 32];
562 let transport = Arc::new(CapturingTransport {
563 captured: captured.clone(),
564 read_key: NoiseCipher::new(&key).expect("32-byte key"),
565 counter: AtomicU32::new(0),
566 });
567 let socket = NoiseSocket::new(
568 Arc::new(crate::runtime_impl::TokioRuntime),
569 transport,
570 NoiseCipher::new(&key).expect("32-byte key"),
571 NoiseCipher::new(&key).expect("32-byte key"),
572 );
573
574 let small: Vec<u8> = (0..1_000u32).map(|i| i as u8).collect();
575 let large: Vec<u8> = (0..40_000u32).map(|i| (i % 251) as u8).collect();
576 assert!(small.len() <= INLINE_ENCRYPT_THRESHOLD);
577 assert!(large.len() > INLINE_ENCRYPT_THRESHOLD);
578
579 socket
580 .encrypt_and_send(bytes::Bytes::from(small.clone()))
581 .await
582 .expect("small frame send");
583 socket
584 .encrypt_and_send(bytes::Bytes::from(large.clone()))
585 .await
586 .expect("large frame send");
587
588 let got = captured.lock().await;
589 assert_eq!(got.len(), 2);
590 assert_eq!(got[0], small, "inline (<=16KB) frame must round-trip");
591 assert_eq!(
592 got[1], large,
593 "large (>16KB) frame must round-trip via the moved-Bytes path"
594 );
595 }
596
597 struct AcceptThenFailTransport {
601 sent: std::sync::Mutex<Vec<bytes::Bytes>>,
602 fail_from: usize,
603 disconnected: AtomicBool,
604 }
605
606 impl AcceptThenFailTransport {
607 fn new(fail_from: usize) -> Self {
608 Self {
609 sent: std::sync::Mutex::new(Vec::new()),
610 fail_from,
611 disconnected: AtomicBool::new(false),
612 }
613 }
614
615 fn sent(&self) -> Vec<bytes::Bytes> {
616 self.sent.lock().expect("send mutex").clone()
617 }
618
619 fn disconnected(&self) -> bool {
620 self.disconnected.load(Ordering::SeqCst)
621 }
622 }
623
624 #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
625 #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
626 impl Transport for AcceptThenFailTransport {
627 async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
628 let mut sent = self.sent.lock().expect("send mutex");
629 sent.push(data);
630 if sent.len() > self.fail_from {
631 return Err(anyhow::anyhow!(
632 "injected failure after accepting the frame"
633 ));
634 }
635 Ok(())
636 }
637 async fn disconnect(&self) {
638 self.disconnected.store(true, Ordering::SeqCst);
639 }
640 }
641
642 fn test_socket(transport: Arc<dyn Transport>) -> NoiseSocket {
643 let key = [0x11u8; 32];
644 NoiseSocket::new(
645 Arc::new(crate::runtime_impl::TokioRuntime),
646 transport,
647 NoiseCipher::new(&key).expect("32-byte key"),
648 NoiseCipher::new(&key).expect("32-byte key"),
649 )
650 }
651
652 #[tokio::test]
656 async fn send_error_before_write_poisons_the_sender() {
657 let transport = Arc::new(crate::transport::mock::CapturingMockTransport::new());
658 transport.fail_next_sends(1);
659 let socket = test_socket(transport.clone());
660
661 let first = socket
662 .encrypt_and_send(bytes::Bytes::from_static(b"first"))
663 .await
664 .expect_err("injected transport failure");
665 assert!(matches!(first.kind, EncryptSendErrorKind::Transport));
666
667 for attempt in 0..3 {
668 let err = socket
669 .encrypt_and_send(bytes::Bytes::from_static(b"later"))
670 .await
671 .expect_err("sends after a transport failure must be refused");
672 assert!(
673 matches!(err.kind, EncryptSendErrorKind::Poisoned),
674 "attempt {attempt} should be rejected as poisoned, got {err:?}"
675 );
676 assert!(err.is_transport_unavailable(), "must force a reconnect");
677 }
678
679 assert_eq!(
680 transport.sent_count(),
681 0,
682 "no frame may reach the wire after the sender is poisoned"
683 );
684 assert_eq!(transport.failed_sends(), 1, "only the first send was tried");
685 }
686
687 #[tokio::test]
691 async fn ambiguous_send_error_poisons_the_sender() {
692 let transport = Arc::new(AcceptThenFailTransport::new(0));
693 let socket = test_socket(transport.clone());
694
695 let first = socket
696 .encrypt_and_send(bytes::Bytes::from_static(b"first"))
697 .await
698 .expect_err("transport reported failure after accepting the frame");
699 assert!(matches!(first.kind, EncryptSendErrorKind::Transport));
700
701 let second = socket
702 .encrypt_and_send(bytes::Bytes::from_static(b"second"))
703 .await
704 .expect_err("sends after an ambiguous failure must be refused");
705 assert!(matches!(second.kind, EncryptSendErrorKind::Poisoned));
706
707 assert_eq!(
708 transport.sent().len(),
709 1,
710 "exactly the one ambiguous frame reached the transport"
711 );
712 }
713
714 #[tokio::test]
720 async fn poisoning_the_sender_closes_the_transport() {
721 let transport: Arc<AcceptThenFailTransport> = Arc::new(AcceptThenFailTransport::new(0));
722 let socket = test_socket(transport.clone());
723
724 let first = socket
725 .encrypt_and_send(bytes::Bytes::from_static(b"first"))
726 .await;
727 assert!(first.is_err(), "the injected failure must surface");
728 assert!(
729 transport.disconnected(),
730 "the first transport error must close the transport so the client reconnects"
731 );
732 }
733
734 #[tokio::test]
739 async fn every_encrypted_frame_burns_its_own_counter() {
740 let key = [0x33u8; 32];
741 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
742 let write_key = Arc::new(NoiseCipher::new(&key).expect("32-byte key"));
743
744 let mut write_counter: u32 = 0;
745 let mut out_buf = BytesMut::new();
746
747 for expected_counter in 0..2u32 {
748 assert_eq!(write_counter, expected_counter);
749 NoiseSocket::encrypt_frame_into(
750 &runtime,
751 &write_key,
752 &mut write_counter,
753 bytes::Bytes::from(vec![expected_counter as u8; 32]),
754 &mut out_buf,
755 )
756 .await
757 .expect("encrypt must succeed");
758 assert_eq!(
759 write_counter,
760 expected_counter + 1,
761 "each frame must consume its counter at encrypt time"
762 );
763 }
764
765 let read_key = NoiseCipher::new(&key).expect("32-byte key");
766 for (counter, frame) in split_frames(&out_buf).into_iter().enumerate() {
767 let mut body = frame;
768 read_key
769 .decrypt_in_place_with_counter(counter as u32, &mut body)
770 .expect("each frame must decrypt under its own distinct counter");
771 assert_eq!(body, vec![counter as u8; 32]);
772 }
773 }
774
775 #[tokio::test]
782 async fn a_frame_is_sealed_in_place_behind_the_one_before_it() {
783 let key = [0x21u8; 32];
784 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
785 let write_key = Arc::new(NoiseCipher::new(&key).expect("32-byte key"));
786 let mut write_counter: u32 = 0;
787 let mut out_buf = BytesMut::new();
788
789 let first = bytes::Bytes::from(vec![0xA1u8; 40]);
790 NoiseSocket::encrypt_frame_into(
791 &runtime,
792 &write_key,
793 &mut write_counter,
794 first,
795 &mut out_buf,
796 )
797 .await
798 .expect("first frame");
799 let first_frame = out_buf.to_vec();
800
801 let second_plain = vec![0xB2u8; 77];
802 let wire_len = NoiseSocket::encrypt_frame_into(
803 &runtime,
804 &write_key,
805 &mut write_counter,
806 bytes::Bytes::from(second_plain.clone()),
807 &mut out_buf,
808 )
809 .await
810 .expect("second frame");
811
812 assert_eq!(
813 &out_buf[..first_frame.len()],
814 &first_frame[..],
815 "sealing the second frame must not reach back into the first"
816 );
817 assert_eq!(wire_len, frame_wire_len(second_plain.len()));
818 assert_eq!(out_buf.len(), first_frame.len() + wire_len);
819
820 let second = &out_buf[first_frame.len()..];
821 let declared =
822 ((second[0] as usize) << 16) | ((second[1] as usize) << 8) | second[2] as usize;
823 assert_eq!(
824 declared,
825 second_plain.len() + TAG_LEN,
826 "the header must count the ciphertext that was sealed after it"
827 );
828
829 let read_key = NoiseCipher::new(&key).expect("32-byte key");
830 let mut body = BytesMut::from(&second[FRAME_LENGTH_SIZE..]);
831 read_key
832 .decrypt_in_place_with_counter(1, &mut body)
833 .expect("the sealed body must authenticate under its own counter");
834 assert_eq!(&body[..], &second_plain[..]);
835 }
836
837 #[test]
842 fn the_frame_body_view_never_reaches_before_its_own_frame() {
843 let mut out = BytesMut::from(&b"earlier-frame"[..]);
844 let base = out.len();
845 out.extend_from_slice(b"body");
846
847 let mut view = FrameBody {
848 out: &mut out,
849 base,
850 };
851 assert_eq!(view.as_slice(), b"body");
852 assert_eq!(view.len(), 4);
853 view.as_mut_slice()[0] = b'B';
854
855 view.resize(4 + TAG_LEN, 0);
857 assert_eq!(view.len(), 4 + TAG_LEN);
858 view.truncate(4);
859 assert_eq!(view.as_slice(), b"Body");
860
861 assert_eq!(
862 &out[..base],
863 &b"earlier-frame"[..],
864 "no view operation may touch the bytes staged before this frame"
865 );
866 }
867
868 #[tokio::test]
874 async fn a_failed_frame_leaves_the_batch_buffer_byte_identical() {
875 let key = [0x22u8; 32];
876 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
877 let write_key = Arc::new(NoiseCipher::new(&key).expect("32-byte key"));
878 let mut out_buf = BytesMut::new();
879
880 let mut write_counter: u32 = u32::MAX - 1;
882 NoiseSocket::encrypt_frame_into(
883 &runtime,
884 &write_key,
885 &mut write_counter,
886 bytes::Bytes::from(vec![0xC3u8; 24]),
887 &mut out_buf,
888 )
889 .await
890 .expect("the last usable counter must still encrypt");
891 let staged = out_buf.to_vec();
892 assert_eq!(write_counter, u32::MAX);
893
894 let err = NoiseSocket::encrypt_frame_into(
895 &runtime,
896 &write_key,
897 &mut write_counter,
898 bytes::Bytes::from(vec![0xD4u8; 24]),
899 &mut out_buf,
900 )
901 .await
902 .expect_err("an exhausted counter must not wrap");
903 assert!(matches!(err.kind, EncryptSendErrorKind::Crypto));
904 assert_eq!(
905 out_buf.to_vec(),
906 staged,
907 "the rejected frame must not leave a header or a plaintext behind"
908 );
909 assert_eq!(write_counter, u32::MAX, "a rejected frame burns no counter");
910
911 let read_key = NoiseCipher::new(&key).expect("32-byte key");
913 let mut body = BytesMut::from(&staged[FRAME_LENGTH_SIZE..]);
914 read_key
915 .decrypt_in_place_with_counter(u32::MAX - 1, &mut body)
916 .expect("the frame staged before the failure must still be sendable");
917 assert_eq!(&body[..], &[0xC3u8; 24][..]);
918 }
919
920 #[tokio::test]
928 async fn order_survives_a_full_job_channel() {
929 let key = [0x88u8; 32];
930 let transport = GatedTransport::closed();
931 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
932 let socket = Arc::new(NoiseSocket::new(
933 runtime,
934 transport.clone(),
935 NoiseCipher::new(&key).expect("32-byte key"),
936 NoiseCipher::new(&key).expect("32-byte key"),
937 ));
938
939 const FRAMES: usize = 20;
941 let sends: Vec<BoxSend> = (0..FRAMES)
942 .map(|i| {
943 let socket = socket.clone();
944 Box::pin(async move {
945 socket
946 .encrypt_and_send(bytes::Bytes::from(vec![i as u8; 32]))
947 .await
948 }) as BoxSend
949 })
950 .collect();
951 let mut joined = futures::future::join_all(sends);
952 assert!(
953 futures::FutureExt::now_or_never(&mut joined).is_none(),
954 "the gate is closed, so nothing can have completed"
955 );
956
957 transport.gate.add_permits(FRAMES);
958 for result in joined.await {
959 result.expect("send must succeed");
960 }
961
962 let read_key = NoiseCipher::new(&key).expect("32-byte key");
963 let bodies: Vec<Vec<u8>> = transport
964 .writes()
965 .iter()
966 .flat_map(|w| split_frames(w))
967 .collect();
968 assert_eq!(bodies.len(), FRAMES, "every frame must reach the wire");
969 for (counter, mut body) in bodies.into_iter().enumerate() {
970 read_key
971 .decrypt_in_place_with_counter(counter as u32, &mut body)
972 .expect("a frame written out of counter order cannot authenticate");
973 assert_eq!(
979 body,
980 vec![counter as u8; 32],
981 "frame {counter} must carry the payload submitted at position {counter}"
982 );
983 }
984 }
985
986 #[tokio::test]
992 async fn a_lone_waiter_gets_the_transport_error_untouched() {
993 #[derive(Debug)]
994 struct TypedTransportError;
995 impl std::fmt::Display for TypedTransportError {
996 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
997 write!(f, "typed transport error")
998 }
999 }
1000 impl std::error::Error for TypedTransportError {}
1001
1002 struct TypedFailTransport;
1003
1004 #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
1005 #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
1006 impl Transport for TypedFailTransport {
1007 async fn send(&self, _data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
1008 Err(anyhow::Error::new(TypedTransportError))
1009 }
1010 async fn disconnect(&self) {}
1011 }
1012
1013 let key = [0x77u8; 32];
1014 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
1015 let socket = NoiseSocket::new(
1016 runtime,
1017 Arc::new(TypedFailTransport),
1018 NoiseCipher::new(&key).expect("32-byte key"),
1019 NoiseCipher::new(&key).expect("32-byte key"),
1020 );
1021
1022 let err = socket
1023 .encrypt_and_send(bytes::Bytes::from(vec![9u8; 32]))
1024 .await
1025 .expect_err("the transport always fails");
1026
1027 assert!(matches!(err.kind, EncryptSendErrorKind::Transport));
1028 assert!(
1029 err.source.downcast_ref::<TypedTransportError>().is_some(),
1030 "a lone waiter must receive the transport's own error type, got: {:?}",
1031 err.source
1032 );
1033 }
1034
1035 #[tokio::test]
1040 async fn a_batch_never_overshoots_the_byte_ceiling() {
1041 let key = [0x66u8; 32];
1042 let transport = GatedTransport::closed();
1043 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
1044 let socket = Arc::new(NoiseSocket::new(
1045 runtime,
1046 transport.clone(),
1047 NoiseCipher::new(&key).expect("32-byte key"),
1048 NoiseCipher::new(&key).expect("32-byte key"),
1049 ));
1050
1051 const FRAME_BYTES: usize = 20 * 1024;
1054 const FRAMES: usize = 5;
1055 let mut sends = queue_all(
1056 &socket,
1057 (0..FRAMES).map(|i| bytes::Bytes::from(vec![i as u8; FRAME_BYTES])),
1058 );
1059 transport.gate.add_permits(FRAMES);
1060 for result in (&mut sends).await {
1061 result.expect("send must succeed");
1062 }
1063
1064 let writes = transport.writes();
1065 for write in &writes {
1066 let frames = split_frames(write);
1067 assert!(
1068 frames.len() == 1 || write.len() <= MAX_BATCH_WIRE_BYTES,
1069 "a multi-frame write must respect the ceiling: {} bytes in {} frames",
1070 write.len(),
1071 frames.len()
1072 );
1073 }
1074 assert!(
1075 writes.iter().any(|w| split_frames(w).len() > 1),
1076 "the burst must still coalesce, otherwise this proves nothing"
1077 );
1078
1079 let read_key = NoiseCipher::new(&key).expect("32-byte key");
1080 let bodies: Vec<Vec<u8>> = writes.iter().flat_map(|w| split_frames(w)).collect();
1081 assert_eq!(bodies.len(), FRAMES, "a held-over frame must still be sent");
1082 for (counter, mut body) in bodies.into_iter().enumerate() {
1083 read_key
1084 .decrypt_in_place_with_counter(counter as u32, &mut body)
1085 .expect("holding a frame over must not disturb counter order");
1086 }
1087 }
1088
1089 #[tokio::test]
1094 async fn the_transport_cause_reaches_the_caller() {
1095 let key = [0x55u8; 32];
1096 let transport: Arc<AcceptThenFailTransport> = Arc::new(AcceptThenFailTransport::new(0));
1097 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
1098 let socket = NoiseSocket::new(
1099 runtime,
1100 transport,
1101 NoiseCipher::new(&key).expect("32-byte key"),
1102 NoiseCipher::new(&key).expect("32-byte key"),
1103 );
1104
1105 let err = socket
1106 .encrypt_and_send(bytes::Bytes::from(vec![7u8; 32]))
1107 .await
1108 .expect_err("the transport always fails");
1109
1110 assert!(matches!(err.kind, EncryptSendErrorKind::Transport));
1111 let chain = format!("{:#}", err.source);
1113 assert!(
1114 chain.contains("injected failure after accepting the frame"),
1115 "the transport's cause was lost on the way to the caller: {chain}"
1116 );
1117 }
1118
1119 struct GatedTransport {
1123 writes: std::sync::Mutex<Vec<bytes::Bytes>>,
1124 gate: tokio::sync::Semaphore,
1125 }
1126
1127 impl GatedTransport {
1128 fn closed() -> Arc<Self> {
1129 Arc::new(Self {
1130 writes: std::sync::Mutex::new(Vec::new()),
1131 gate: tokio::sync::Semaphore::new(0),
1132 })
1133 }
1134
1135 fn writes(&self) -> Vec<bytes::Bytes> {
1136 self.writes.lock().expect("writes mutex").clone()
1137 }
1138 }
1139
1140 #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
1141 #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
1142 impl Transport for GatedTransport {
1143 async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
1144 let permit = self.gate.acquire().await.expect("gate open");
1145 permit.forget();
1146 self.writes.lock().expect("writes mutex").push(data);
1147 Ok(())
1148 }
1149 async fn disconnect(&self) {}
1150 }
1151
1152 fn queue_all(
1161 socket: &Arc<NoiseSocket>,
1162 payloads: impl Iterator<Item = bytes::Bytes>,
1163 ) -> futures::future::JoinAll<BoxSend> {
1164 let sends: Vec<BoxSend> = payloads
1165 .map(|payload| {
1166 let socket = socket.clone();
1167 Box::pin(async move { socket.encrypt_and_send(payload).await }) as BoxSend
1168 })
1169 .collect();
1170 let mut joined = futures::future::join_all(sends);
1171 let queued = futures::FutureExt::now_or_never(&mut joined);
1172 assert!(
1173 queued.is_none(),
1174 "the sends must still be in flight: the transport gate is closed"
1175 );
1176 joined
1177 }
1178
1179 type BoxSend = std::pin::Pin<Box<dyn Future<Output = SendResult> + Send>>;
1180
1181 fn split_frames(mut wire: &[u8]) -> Vec<Vec<u8>> {
1183 let mut frames = Vec::new();
1184 while !wire.is_empty() {
1185 let mut len = 0usize;
1186 for byte in &wire[..FRAME_LENGTH_SIZE] {
1187 len = (len << 8) | *byte as usize;
1188 }
1189 let body = &wire[FRAME_LENGTH_SIZE..FRAME_LENGTH_SIZE + len];
1190 frames.push(body.to_vec());
1191 wire = &wire[FRAME_LENGTH_SIZE + len..];
1192 }
1193 frames
1194 }
1195
1196 #[tokio::test]
1201 async fn queued_frames_leave_in_one_write_in_counter_order() {
1202 let key = [0x44u8; 32];
1203 let transport = GatedTransport::closed();
1204 let runtime: Arc<dyn Runtime> = Arc::new(crate::runtime_impl::TokioRuntime);
1205 let socket = Arc::new(NoiseSocket::new(
1206 runtime,
1207 transport.clone(),
1208 NoiseCipher::new(&key).expect("32-byte key"),
1209 NoiseCipher::new(&key).expect("32-byte key"),
1210 ));
1211
1212 const FRAMES: usize = 5;
1213 let mut sends = queue_all(
1214 &socket,
1215 (0..FRAMES).map(|i| bytes::Bytes::from(vec![i as u8; 32])),
1216 );
1217 transport.gate.add_permits(FRAMES);
1218 for result in (&mut sends).await {
1219 result.expect("send must succeed");
1220 }
1221
1222 let writes = transport.writes();
1223 assert!(
1224 writes.len() < FRAMES,
1225 "queued frames must coalesce, got {} writes for {FRAMES} frames",
1226 writes.len()
1227 );
1228
1229 let read_key = NoiseCipher::new(&key).expect("32-byte key");
1230 let bodies: Vec<Vec<u8>> = writes.iter().flat_map(|w| split_frames(w)).collect();
1231 assert_eq!(bodies.len(), FRAMES, "every frame must reach the wire");
1232
1233 let mut payloads = Vec::new();
1237 for (counter, body) in bodies.into_iter().enumerate() {
1238 let mut body = body;
1239 read_key
1240 .decrypt_in_place_with_counter(counter as u32, &mut body)
1241 .expect("frames must be written in counter order");
1242 assert_eq!(body, vec![body[0]; 32], "frame body must survive intact");
1243 payloads.push(body[0]);
1244 }
1245
1246 payloads.sort_unstable();
1250 let expected: Vec<u8> = (0..FRAMES as u8).collect();
1251 assert_eq!(
1252 payloads, expected,
1253 "each producer's payload must appear exactly once"
1254 );
1255 }
1256
1257 #[tokio::test]
1260 async fn framing_error_does_not_poison_the_sender() {
1261 let transport = Arc::new(crate::transport::mock::CapturingMockTransport::new());
1262 let socket = test_socket(transport.clone());
1263
1264 let oversize = bytes::Bytes::from(vec![0u8; wacore::framing::FRAME_MAX_SIZE - 16]);
1267 let err = socket
1268 .encrypt_and_send(oversize)
1269 .await
1270 .expect_err("frame exceeds the 24-bit length prefix");
1271 assert!(matches!(err.kind, EncryptSendErrorKind::Framing));
1272
1273 socket
1274 .encrypt_and_send(bytes::Bytes::from_static(b"still usable"))
1275 .await
1276 .expect("a rejected oversize frame must leave the connection usable");
1277 assert_eq!(transport.sent_count(), 1);
1278 }
1279
1280 #[tokio::test]
1281 async fn test_concurrent_sends_maintain_order() {
1282 use async_lock::Mutex;
1283 use async_trait::async_trait;
1284 use std::sync::Arc;
1285
1286 struct RecordingTransport {
1289 recorded_order: Arc<Mutex<Vec<u8>>>,
1290 read_key: NoiseCipher,
1291 counter: AtomicU32,
1292 }
1293
1294 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
1295 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
1296 impl Transport for RecordingTransport {
1297 async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
1298 for mut frame in split_frames(&data) {
1302 let counter = self.counter.fetch_add(1, Ordering::SeqCst);
1303
1304 if self
1305 .read_key
1306 .decrypt_in_place_with_counter(counter, &mut frame)
1307 .is_ok()
1308 && !frame.is_empty()
1309 {
1310 let index = frame[0];
1311 let mut order = self.recorded_order.lock().await;
1312 order.push(index);
1313 }
1314 }
1315 Ok(())
1316 }
1317
1318 async fn disconnect(&self) {}
1319 }
1320
1321 let recorded_order = Arc::new(Mutex::new(Vec::new()));
1322 let key = [0u8; 32];
1323 let write_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
1324 let read_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
1325
1326 let transport = Arc::new(RecordingTransport {
1327 recorded_order: recorded_order.clone(),
1328 read_key: NoiseCipher::new(&key).expect("32-byte key should be valid"),
1329 counter: AtomicU32::new(0),
1330 });
1331
1332 let socket = Arc::new(NoiseSocket::new(
1333 Arc::new(crate::runtime_impl::TokioRuntime),
1334 transport,
1335 write_key,
1336 read_key,
1337 ));
1338
1339 let mut handles = Vec::new();
1341 for i in 0..10 {
1342 let socket = socket.clone();
1343 handles.push(tokio::spawn(async move {
1344 let mut plaintext = vec![i as u8];
1346 plaintext.extend_from_slice(&[0u8; 99]);
1347 socket.encrypt_and_send(bytes::Bytes::from(plaintext)).await
1348 }));
1349 }
1350
1351 for handle in handles {
1353 let result = handle.await.expect("task should complete");
1354 assert!(result.is_ok(), "All sends should succeed");
1355 }
1356
1357 let order = recorded_order.lock().await;
1359 let expected: Vec<u8> = (0..10).collect();
1360 assert_eq!(*order, expected, "Sends should maintain FIFO order");
1361 }
1362
1363 #[tokio::test]
1366 async fn test_encrypted_buffer_sizing_is_sufficient() {
1367 use async_trait::async_trait;
1368 use std::sync::Arc;
1369 use std::sync::atomic::{AtomicUsize, Ordering};
1370
1371 struct SizeRecordingTransport {
1373 last_size: Arc<AtomicUsize>,
1374 }
1375
1376 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
1377 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
1378 impl Transport for SizeRecordingTransport {
1379 async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
1380 self.last_size.store(data.len(), Ordering::SeqCst);
1381 Ok(())
1382 }
1383 async fn disconnect(&self) {}
1384 }
1385
1386 let last_size = Arc::new(AtomicUsize::new(0));
1387 let transport = Arc::new(SizeRecordingTransport {
1388 last_size: last_size.clone(),
1389 });
1390
1391 let key = [0u8; 32];
1392 let write_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
1393 let read_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
1394
1395 let socket = NoiseSocket::new(
1396 Arc::new(crate::runtime_impl::TokioRuntime),
1397 transport,
1398 write_key,
1399 read_key,
1400 );
1401
1402 let test_sizes = [0, 1, 50, 100, 500, 1000, 1024, 2000, 5000, 16384, 20000];
1404
1405 for size in test_sizes {
1406 let plaintext = vec![0xABu8; size];
1407 let result = socket
1408 .encrypt_and_send(bytes::Bytes::from(plaintext.clone()))
1409 .await;
1410
1411 assert!(
1412 result.is_ok(),
1413 "encrypt_and_send should succeed for payload size {}",
1414 size
1415 );
1416
1417 let actual_encrypted_size = last_size.load(Ordering::SeqCst);
1418
1419 let expected_max = size + 19;
1422 assert_eq!(
1423 actual_encrypted_size, expected_max,
1424 "Encrypted size for {} byte payload should be {} (got {})",
1425 size, expected_max, actual_encrypted_size
1426 );
1427 }
1428 }
1429
1430 #[tokio::test]
1433 async fn session_stats_match_transport_bytes() {
1434 let factory = crate::transport::mock::CapturingMockTransportFactory::new();
1435 let transport = factory.transport();
1436 let key = [0u8; 32];
1437 let stats = Arc::new(wacore::stats::SessionStats::new());
1438
1439 let socket = NoiseSocket::with_stats(
1440 Arc::new(crate::runtime_impl::TokioRuntime),
1441 transport.clone(),
1442 NoiseCipher::new(&key).expect("32-byte key"),
1443 NoiseCipher::new(&key).expect("32-byte key"),
1444 Some(stats.clone()),
1445 );
1446
1447 for size in [0usize, 100, 5000] {
1448 socket
1449 .encrypt_and_send(bytes::Bytes::from(vec![0u8; size]))
1450 .await
1451 .expect("send");
1452 }
1453
1454 let sent = transport.sent();
1455 let wire_total: usize = sent.iter().map(|f| f.len()).sum();
1456 let snap = stats.snapshot();
1457 assert_eq!(snap.frames_sent, sent.len() as u64);
1458 assert_eq!(snap.bytes_sent, wire_total as u64);
1459 assert!(stats.first_send_since_recv_ms() > 0);
1460 }
1461
1462 #[tokio::test]
1464 async fn test_encrypted_buffer_sizing_edge_cases() {
1465 use async_trait::async_trait;
1466 use std::sync::Arc;
1467
1468 struct NoOpTransport;
1469
1470 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
1471 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
1472 impl Transport for NoOpTransport {
1473 async fn send(&self, _data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> {
1474 Ok(())
1475 }
1476 async fn disconnect(&self) {}
1477 }
1478
1479 let transport = Arc::new(NoOpTransport);
1480 let key = [0u8; 32];
1481 let write_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
1482 let read_key = NoiseCipher::new(&key).expect("32-byte key should be valid");
1483
1484 let socket = NoiseSocket::new(
1485 Arc::new(crate::runtime_impl::TokioRuntime),
1486 transport,
1487 write_key,
1488 read_key,
1489 );
1490
1491 let result = socket.encrypt_and_send(bytes::Bytes::new()).await;
1493 assert!(result.is_ok(), "Empty payload should encrypt successfully");
1494
1495 let at_threshold = bytes::Bytes::from(vec![0u8; 16 * 1024]);
1497 let result = socket.encrypt_and_send(at_threshold).await;
1498 assert!(
1499 result.is_ok(),
1500 "Payload at inline threshold should encrypt successfully"
1501 );
1502
1503 let above_threshold = bytes::Bytes::from(vec![0u8; 16 * 1024 + 1]);
1505 let result = socket.encrypt_and_send(above_threshold).await;
1506 assert!(
1507 result.is_ok(),
1508 "Payload above inline threshold should encrypt successfully"
1509 );
1510 }
1511}