Skip to main content

whatsapp_rust/socket/
noise_socket.rs

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
14/// AES-GCM tag length. A frame's wire size is a fixed function of its plaintext
15/// length, which is what lets the length prefix be written before the ciphertext
16/// exists.
17const TAG_LEN: usize = 16;
18
19/// The region of the batch buffer one frame's ciphertext occupies, exposed to
20/// AES-GCM as if it were a buffer of its own.
21///
22/// Sealing through this view puts the ciphertext and its tag straight where the
23/// transport will read them. The alternative, sealing into scratch space and
24/// copying the result in, costs a second full pass over every byte sent, which
25/// is the copy comparable stacks are built to avoid: quinn seals with
26/// `PacketKey::encrypt(&self, packet, buf, header_len)` directly in the datagram
27/// buffer, and rustls encrypts each fragment into the record it will send.
28struct FrameBody<'a> {
29    out: &'a mut BytesMut,
30    /// Offset in `out` where this frame's ciphertext starts, i.e. just past its
31    /// length prefix. Held as an offset rather than a slice so the AEAD can grow
32    /// the buffer by the tag through the same view.
33    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
54/// Ceilings on one batched write. They bound how much is buffered before the
55/// first frame reaches the socket; the batch never waits for work, so these
56/// only matter when a burst is already queued.
57const MAX_BATCH_FRAMES: usize = 16;
58const MAX_BATCH_WIRE_BYTES: usize = 64 * 1024;
59
60/// Result type for send operations.
61type SendResult = std::result::Result<(), EncryptSendError>;
62
63/// Wire size a plaintext will occupy once encrypted and framed: the AES-GCM tag
64/// plus the length prefix. Used to test a queued frame against the batch ceiling
65/// before paying to encrypt it.
66fn frame_wire_len(plaintext_len: usize) -> usize {
67    plaintext_len + TAG_LEN + wacore::framing::FRAME_LENGTH_SIZE
68}
69
70/// One batched write's failure, handed to every waiter in that batch.
71///
72/// `anyhow::Error` is not `Clone`, so a shared reference is what lets all the
73/// callers see the real cause instead of a re-worded copy.
74#[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
89/// A job sent to the dedicated sender task.
90struct 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    /// Channel to send jobs to the dedicated sender task.
99    /// Using a channel instead of a mutex avoids blocking callers while
100    /// the current send is in progress - they can enqueue their work and
101    /// await the result without holding a lock.
102    send_job_tx: async_channel::Sender<SendJob>,
103    /// Handle to the sender task. Aborted on drop to prevent resource leaks
104    /// if the task is stuck on a slow/hanging network operation.
105    _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    /// Like [`Self::new`], recording sent frames into `stats` (the main WA
119    /// session socket passes the client's [`SessionStats`](wacore::stats::SessionStats); VoIP relay
120    /// sockets and tests pass `None`).
121    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        // Small buffer matched to typical steady-state throughput; the sender
132        // task is network-bound (awaits `transport.send`), so a transient
133        // WebSocket stall will backpressure producers here rather than queue.
134        let (send_job_tx, send_job_rx) = async_channel::bounded::<SendJob>(8);
135
136        // Spawn the dedicated sender task
137        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    /// Dedicated sender task that processes send jobs sequentially.
157    /// This ensures frames are sent in counter order without requiring a mutex.
158    /// The task owns the write counter and processes jobs one at a time.
159    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        // BytesMut: split().freeze() yields a zero-copy Bytes while retaining
168        // the underlying allocation for the next frame.
169        let mut out_buf = BytesMut::with_capacity(4096);
170        // A failed transport write says nothing about how much of the frame the
171        // peer received, so the counter that frame consumed can neither be
172        // reused (nonce reuse under the same write key) nor confidently skipped
173        // (the peer's read counter would desync). Both outcomes are unrecoverable
174        // in-band, so the whole sender goes out of service and the connection
175        // must be re-established with a fresh handshake key.
176        let mut poisoned = false;
177        // Reused across batches: one allocation for the life of the connection
178        // instead of one per batch.
179        let mut waiters: Vec<(oneshot::Sender<SendResult>, usize)> = Vec::new();
180        // A job pulled off the channel that would have overflowed the byte
181        // ceiling, held over to open the next batch. Dropping it (on shutdown)
182        // drops its response channel, which the caller sees as a closed sender:
183        // a held-over job can be lost, but it can never hang its caller.
184        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            // Encrypt everything already queued into one buffer and write it
200            // once. Three independent producers answer a single inbound message
201            // (the reply, the delivery receipt and the stanza ack), so a write
202            // per frame turned into a syscall, a TLS record and a WebSocket
203            // message per frame. Only frames that are ALREADY waiting are taken:
204            // never block for more, or this trades syscalls for latency.
205            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                        // The counter is untouched on this frame, and every
222                        // frame already in the buffer must still go out so the
223                        // peer's counters stay contiguous.
224                        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                        // Check the ceiling before appending, not after, or a
234                        // nearly-full batch overshoots it by a whole frame. A
235                        // frame that cannot fit any batch still goes alone
236                        // rather than deadlocking against the ceiling.
237                        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                // Zero-copy: split() hands the written bytes over and out_buf
252                // keeps its capacity for the next batch.
253                let wire = out_buf.split().freeze();
254                if waiters.len() > 1 {
255                    // The only externally visible sign that a batch happened.
256                    // Without it, "does the peer accept several frames in one
257                    // WebSocket message?" cannot be answered from a live run.
258                    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                // Crypto and framing failures are rejected before any byte
279                // reaches the wire and leave the counter untouched, so they do
280                // not compromise the keystream. Only a transport failure is
281                // ambiguous.
282                if let Err(err) = &outcome
283                    && matches!(err.kind, EncryptSendErrorKind::Transport)
284                {
285                    poisoned = true;
286                    // Poisoning only stops this half. A write can fail while
287                    // the read half stays open (half-open socket, or a
288                    // Transport that reports Err without emitting
289                    // Disconnected), and then nothing else would notice: the
290                    // read loop keeps running and the client reports itself
291                    // connected while every send fails forever. Closing the
292                    // transport makes the existing disconnect path observe the
293                    // drop and reconnect with a fresh handshake key, which is
294                    // the only way this sender becomes usable again.
295                    transport.disconnect().await;
296                }
297            }
298
299            // Every frame in this batch shares the fate of the single write.
300            match outcome {
301                Ok(()) => {
302                    for (response_tx, _) in waiters.drain(..) {
303                        let _ = response_tx.send(Ok(()));
304                    }
305                }
306                // One waiter owns the failure outright. This is the overwhelmingly
307                // common case, and handing over the error untouched is what keeps
308                // `err.source.downcast_ref::<MyTransportError>()` working for a
309                // caller with its own Transport: wrapping would bury the typed
310                // error one level down for no benefit, since there is nobody to
311                // share it with.
312                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                // Several waiters, and EncryptSendError is not Clone: they share
317                // one Arc. Display renders only the kind, so re-wording per waiter
318                // would hand each caller "transport error" with the cause gone;
319                // sharing keeps the whole chain reachable for a refcount bump each.
320                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    /// Encrypt one plaintext and append the framed result to `out_buf`,
336    /// returning its wire size. The counter is burned once the framed ciphertext
337    /// is committed to `out_buf`, whether or not the write that carries it
338    /// succeeds. Every error path leaves `out_buf` exactly as it found it, which
339    /// is the only reason leaving the counter unburned there is sound: a change
340    /// that keeps partial output must burn the counter too, or the next frame
341    /// reuses its nonce.
342    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        // Refuse to wrap the per-direction frame counter: reusing an AES-GCM
351        // nonce under the same key is catastrophic. 2^32 frames per connection
352        // is unreachable in practice, so erroring here forces a reconnect
353        // rather than a silent nonce reuse.
354        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            // Ciphertext is exactly the plaintext plus the tag, so the length
361            // prefix is known before the bytes it counts exist and the frame can
362            // be sealed where it already sits in the batch.
363            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                // Unlike the paths above, this one has already appended the
373                // prefix and the plaintext. Rolling both back is what keeps the
374                // rest of the batch, which still has to go out, contiguous, and
375                // what keeps this counter safe to hand to the next frame. The
376                // default AEAD cannot fail on a fixed-size key and nonce, so
377                // only a `set_crypto_provider` backend reaches this: it is the
378                // contract for those, not dead code.
379                out_buf.truncate(before);
380                return Err(EncryptSendError::crypto(e));
381            }
382            // The length prefix was written from `plaintext.len() + TAG_LEN`
383            // before the ciphertext existed, which is sound only because
384            // `TransportAead` is AES-256-GCM by contract. A `set_crypto_provider`
385            // backend that grows the buffer by anything else would put a frame
386            // on the wire whose prefix disagrees with its body, desyncing the
387            // peer's parser for the rest of the connection. Checking costs one
388            // comparison and turns that into a refused send.
389            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            // `Bytes` is Send + 'static: move it into the blocking task (a
398            // refcount bump) instead of copying the whole >16KB plaintext.
399            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    /// Hands `plaintext` to the sender task and returns the channel its result
417    /// will arrive on, without waiting for it.
418    ///
419    /// Split out of [`Self::encrypt_and_send`] so a burst can enqueue every
420    /// frame before awaiting any of them. The sender coalesces whatever is
421    /// already queued into one transport write, so a caller that awaited each
422    /// frame before enqueueing the next would hand them over one completion
423    /// apart and get one write per frame.
424    ///
425    /// The returned receiver must be awaited, or the result is dropped and the
426    /// caller cannot tell a delivered frame from a failed one.
427    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        // Send job to the sender task. If channel is closed, sender task has stopped.
439        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    /// Awaits a receiver handed out by [`Self::enqueue_send`].
447    pub(crate) async fn await_send(receiver: oneshot::Receiver<SendResult>) -> SendResult {
448        match receiver.await {
449            Ok(result) => result,
450            Err(_) => {
451                // Sender task dropped without sending a response
452                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        // Checked increment: error instead of wrapping the read counter (AES-GCM
464        // nonce reuse). fetch_update returns the pre-increment counter to use, or
465        // Err when it would overflow u32. Mirrors the write side.
466        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// AbortHandle aborts the sender task on drop automatically, so no manual
478// Drop impl is needed — the `sender_task_handle` field's own Drop does the work.
479
480#[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        // At u32::MAX the next read would wrap the counter to 0 and reuse a nonce;
516        // the counter check fires before decryption, so the bytes don't matter.
517        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    /// Frames above INLINE_ENCRYPT_THRESHOLD take the blocking path that now moves
528    /// the `Bytes` plaintext (refcount) instead of `to_vec()`-copying it. Verify
529    /// both a small (inline) and a large (>16KB) frame still encrypt to ciphertext
530    /// that decrypts back to the exact original.
531    #[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); // strip the 3-byte frame length prefix
550                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    /// A transport that accepts (and records) the frame and *then* reports
598    /// failure: the ambiguous case where the peer may well have consumed the
599    /// frame, so its read counter has already advanced.
600    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    /// Transport failure *before* anything is written: every later send on the
653    /// same connection must be refused, so no second frame can be encrypted
654    /// under the counter the failed frame consumed.
655    #[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    /// Transport failure *after* the frame was accepted (the ambiguous case:
688    /// the peer may have decrypted it and advanced its read counter). The
689    /// sender must still refuse everything that follows.
690    #[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    /// Poisoning the sender is only half a recovery: a write can fail while the
715    /// read half stays open, and then nothing tears the connection down. The
716    /// sender must close the transport so the existing disconnect path
717    /// reconnects, instead of leaving a client that looks connected and cannot
718    /// send.
719    #[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    /// Drives the per-frame primitive directly: every frame burns exactly one
735    /// counter at encrypt time, whether or not the write that carries it ever
736    /// succeeds. Proven by decrypting the two frames with counters 0 and 1 -
737    /// reuse would make the second decrypt fail.
738    #[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    /// The frame is sealed straight into the batch buffer, so the offset it is
776    /// sealed at is load-bearing in a way a staging copy never was: too low and
777    /// AES-GCM overwrites the length prefix or the frame before it, too high and
778    /// the plaintext leaks past the ciphertext. Pinned by encrypting a second
779    /// frame behind a first and checking the first is untouched, the header
780    /// counts exactly the ciphertext, and the body decrypts under its counter.
781    #[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    /// The AEAD grows and shrinks the buffer through this view, so every one of
838    /// its operations has to be relative to the frame's own start. An absolute
839    /// `resize` or `truncate` here would silently eat the frames already staged
840    /// for the same write.
841    #[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        // Growing by a tag-sized amount, the way sealing does.
856        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    /// A frame that cannot be encrypted must leave the batch buffer byte for byte
869    /// as it found it: the frames already in it still have to reach the wire, and
870    /// the counter it declined to burn is handed to whoever comes next. Counter
871    /// exhaustion is the failure that is reachable without swapping the process
872    /// wide crypto provider.
873    #[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        // One frame already staged, then the counter runs out mid-batch.
881        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        // The staged frame is intact and complete, not just the right length.
912        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    /// Order must survive a full job channel, not just an empty one.
921    ///
922    /// A burst larger than the channel leaves some sends parked waiting for a
923    /// slot, and the whole ordering guarantee (`send_raw_bytes_burst` promises
924    /// arrival order, and the ack worker relies on it) then rests on those
925    /// parked senders being woken in the order they queued. Frame N decrypts
926    /// only under counter N, so any reordering fails here.
927    #[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        // Comfortably past the channel's capacity, so later sends must park.
940        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            // Decrypting alone would not catch a reorder: jobs that woke out of
974            // FIFO order would be encrypted in that order too, so their
975            // counters would still line up. The payload is what pins it -
976            // unlike the concurrent-producer test, these sends are polled in
977            // order by one joined future, so submission order is deterministic.
978            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    /// A single-frame send hands its caller the transport's own error, not a
987    /// wrapper. Callers with a custom `Transport` downcast to their own error
988    /// type to decide whether a failure is retryable, and `downcast_ref` looks
989    /// at the concrete type rather than walking the chain, so wrapping the
990    /// common case would silently break that.
991    #[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    /// The byte ceiling must hold across a burst. Checking it after appending
1036    /// would let a nearly-full batch overshoot by a whole frame, which for a
1037    /// large stanza is the difference between a bounded buffer and an unbounded
1038    /// one.
1039    #[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        // Sized so three fit under the ceiling and the fourth cannot: the batch
1052        // has to stop and hold it over rather than append it.
1053        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    /// The transport's own error must survive the hop to every caller. It cannot
1090    /// be cloned, and `EncryptSendError`'s Display renders only the kind, so a
1091    /// naive rebuild silently degrades "connection reset by peer" into
1092    /// "transport error" and the caller loses the only diagnostic there was.
1093    #[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        // `{:#}` walks the anyhow chain; the injected message must still be in it.
1112        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    /// A transport whose writes block until permits are handed out, so a test
1120    /// can pile jobs into the sender's channel and then release them: the state
1121    /// batching exists for.
1122    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    /// Queues every payload on `socket` and returns the joined sends, still
1153    /// pending.
1154    ///
1155    /// This is the batching tests' precondition: all N frames sitting in the
1156    /// sender's channel at once. It holds by construction rather than by
1157    /// waiting - `send` on a channel with room resolves on its first poll, so
1158    /// polling the joined future once has queued every job - which is why these
1159    /// tests do not spin on `yield_now` and hope the scheduler cooperated.
1160    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    /// Splits a concatenated run of length-prefixed frames into their bodies.
1182    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    /// Frames queued while a write is in flight leave together in one write, in
1197    /// counter order, and every caller is answered. Batching is only sound if
1198    /// all three hold: a lost waiter hangs a send forever, and reordering would
1199    /// desync the peer's read counter.
1200    #[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        // Decrypting frame N under counter N is the order proof: the counter is
1234        // the AES-GCM nonce, so a frame written out of order fails to
1235        // authenticate here.
1236        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        // Which producer wins which counter is not fixed - the senders race into
1247        // the channel - so the invariant is that each one's payload is on the
1248        // wire exactly once, none dropped and none duplicated.
1249        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    /// A framing failure is detected before any byte reaches the wire, so it
1258    /// must not disable the connection the way a transport failure does.
1259    #[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        // Ciphertext = payload + 16-byte tag, so this is the smallest payload
1265        // whose frame no longer fits the 24-bit length prefix.
1266        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        // Create a mock transport that records the order of sends by decrypting
1287        // the first byte (which contains the task index)
1288        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                // One write can carry several frames: the sender coalesces
1299                // whatever is already queued, so each write is unpacked frame by
1300                // frame before decrypting.
1301                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        // Spawn multiple concurrent sends with their indices
1340        let mut handles = Vec::new();
1341        for i in 0..10 {
1342            let socket = socket.clone();
1343            handles.push(tokio::spawn(async move {
1344                // Use index as the first byte of plaintext to identify this send
1345                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        // Wait for all sends to complete
1352        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        // Verify all sends completed in FIFO order (0, 1, 2, ..., 9)
1358        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    /// Tests that the encrypted buffer sizing formula (plaintext.len() + 32) is sufficient.
1364    /// This verifies the optimization in client.rs that sizes the buffer based on payload.
1365    #[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        // Transport that records the actual encrypted data size
1372        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        // Test various payload sizes: tiny, small, medium, large, very large
1403        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            // Verify the actual encrypted size fits within our allocated capacity
1420            // Encrypted size = plaintext + 16 (AES-GCM tag) + 3 (frame header) = plaintext + 19
1421            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    /// Locks the SessionStats wire accounting to the transport truth: bytes
1431    /// counted must equal the frames the transport actually saw.
1432    #[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    /// Tests edge cases for buffer sizing
1463    #[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        // Test empty payload
1492        let result = socket.encrypt_and_send(bytes::Bytes::new()).await;
1493        assert!(result.is_ok(), "Empty payload should encrypt successfully");
1494
1495        // Test payload at inline threshold boundary (16KB)
1496        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        // Test payload just above inline threshold
1504        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}