Skip to main content

darkbio_wire/transport/
client.rs

1// wire-rs: encrypted protocol between Ark and host
2// Copyright 2026 Dark Bio AG. All rights reserved.
3//
4// Use of this source code is governed by a BSD-style
5// license that can be found in the LICENSE file.
6
7use crate::LogId;
8use crate::transport::DEFAULT_HANDSHAKE_TIMEOUT;
9use crate::transport::framing::FrameReader;
10use crate::transport::handshake;
11use crate::transport::io::check_deadline;
12use crate::transport::outbound::{Outbound, Side};
13use crate::transport::sealing;
14use crate::transport::sender::Sender;
15use crate::transport::server::Attestation;
16use crate::transport::{
17    CRYPTO_DOMAIN_WIRE, CRYPTO_DOMAIN_WIRE_ARK_TO_HOST, CRYPTO_DOMAIN_WIRE_HOST_TO_ARK, Closer,
18    Error, Read, Stream, Write,
19};
20use darkbio_crypto::{cbor, cose, xdsa, xhpke};
21use darkbio_trust as trust;
22use std::fmt;
23use std::sync::{Arc, Mutex};
24use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
25use tracing::{debug, info, trace, warn};
26
27/// Trust policy for the device attestation presented during a handshake.
28/// The caller decides which roots to trust and whether to allow self-signed
29/// attestations or recovery overrides. Transport enforces that decision.
30pub trait Verifier {
31    /// Session info extracted from an accepted attestation.
32    type Info;
33
34    /// Verifies the device attestation at `now`, the wall time of the stream's
35    /// clock. Returns the server's identity key along with any info extracted
36    /// from the attestation. Transport checks the handshake signature against
37    /// that key. Rejecting the attestation aborts the handshake.
38    fn verify(
39        &self,
40        attestation: &Attestation,
41        now: SystemTime,
42    ) -> Result<(xdsa::PublicKey, Self::Info), String>;
43}
44
45/// Authenticates the handshake against this pinned identity key. The presented
46/// attestation is returned unchanged, without checking who issued it.
47impl Verifier for xdsa::PublicKey {
48    type Info = Attestation;
49
50    fn verify(
51        &self,
52        attestation: &Attestation,
53        _: SystemTime,
54    ) -> Result<(xdsa::PublicKey, Self::Info), String> {
55        Ok((self.clone(), attestation.clone()))
56    }
57}
58
59/// Roots trusted to attest Arks. Hardware and emulator roots are checked
60/// separately, and attestations must be valid at the stream clock's wall
61/// time. Self-signed attestations from devices that have not been onboarded
62/// are rejected.
63#[derive(Debug)]
64pub struct Roots<'a> {
65    /// Roots attesting hardware Arks.
66    pub hardware: &'a [xdsa::PublicKey],
67    /// Roots attesting emulated Arks.
68    pub emulator: &'a [xdsa::PublicKey],
69}
70
71impl Verifier for Roots<'_> {
72    type Info = trust::device::Device;
73
74    fn verify(
75        &self,
76        attestation: &Attestation,
77        now: SystemTime,
78    ) -> Result<(xdsa::PublicKey, Self::Info), String> {
79        // Convert the supplied wall time to the trust API's Unix seconds
80        let now = now
81            .duration_since(UNIX_EPOCH)
82            .map_err(|err| err.to_string())?
83            .as_secs();
84
85        // Verify the attestation and return its identity with the claims
86        let device = trust::device::verify(
87            attestation.as_bytes(),
88            self.hardware,
89            self.emulator,
90            Some(now),
91        )
92        .map_err(|err| err.to_string())?;
93        Ok((device.identity.clone(), device))
94    }
95}
96/// Client side of the wire, exchanging encrypted messages over a byte stream.
97/// [`Client::connect`] sends a reset and runs the handshake. It returns a
98/// [`Sender`] for outbound messages. [`Client::recv`] decrypts inbound messages.
99///
100/// An empty frame from the server means it has no session with the client anymore.
101/// The client ends its session and returns [`Error::SessionReset`]. The caller
102/// can then reconnect.
103///
104/// Transport checks the shape of the device attestation. A [`Verifier`] decides
105/// whether to trust the server presenting it.
106pub struct Client<R: Read, W: Write> {
107    handshake_timeout: Duration, // Budget for each new handshake attempt
108
109    reader: FrameReader<R>, // COBS framed transport for ingress data
110    receiver: Option<xhpke::Receiver>, // Receive context used exclusively by this client
111    sealer: Option<Arc<Mutex<xhpke::Sender>>>, // Send context shared with active sends
112    outbound: Arc<Outbound<W>>, // Outgoing transport, shared with the senders
113    log_id: LogId,          // Label of the current session in log lines, unset before the first
114}
115
116impl<R: Read, W: Write> Client<R, W> {
117    /// Creates a client owning the byte stream and its shutdown operation, without
118    /// an encrypted session. Call [`Client::connect`] to establish one.
119    /// Output uses the stream's configured write timeout; its adapter must
120    /// enforce deadlines and the shutdown cancellation contract.
121    pub fn new(stream: Stream<R, W>) -> Self {
122        let (reader, writer, close, timeout) = stream.into_parts();
123
124        let outbound = Arc::new(Outbound::new(writer, Side::Client, close.clone(), timeout));
125
126        Self {
127            handshake_timeout: DEFAULT_HANDSHAKE_TIMEOUT,
128            reader: FrameReader::new(reader, close),
129            receiver: None,
130            sealer: None,
131            outbound,
132            log_id: LogId::default(),
133        }
134    }
135
136    /// Sets the budget for each subsequent handshake, starting when connect is
137    /// called. Defaults to [`DEFAULT_HANDSHAKE_TIMEOUT`]. Output and peer replies
138    /// share one deadline; progress and stale frames do not refresh it. Each
139    /// outgoing frame is also limited by the stream's write timeout. Waiting for
140    /// locks and verifier callbacks can extend the call beyond the deadline.
141    ///
142    /// Zero expires attempts immediately. A duration too large to add to an
143    /// [`Instant`] panics when the next handshake's deadline is constructed.
144    pub fn set_handshake_timeout(mut self, timeout: Duration) -> Self {
145        self.handshake_timeout = timeout;
146        self
147    }
148
149    /// A handle that permanently closes the stream from another thread.
150    pub fn closer(&self) -> Closer {
151        self.outbound.closer()
152    }
153
154    /// Permanently closes the stream and waits for adapter shutdown. Senders
155    /// observe closure through write failure; buffered messages remain readable.
156    /// See [`Closer::close`].
157    pub fn close(&self) {
158        self.outbound.close();
159    }
160
161    /// Establishes an encrypted session over the supplied stream, ending any
162    /// previous session first. Sends a reset and drives the handshake:
163    ///
164    ///   1. Client -> Server: HostHello { host_signer, host_crypto }           (plain CBOR)
165    ///   2. Server -> Client: ArkHello  { ark_attest, ark_crypto, a2h_encap }  (cose::seal)
166    ///   3. Client -> Server: HostAck   { h2a_encap }                          (cose::seal)
167    ///
168    /// The verifier receives the server's device attestation. Its accepted info
169    /// is returned alongside the new sender. That sender belongs to this session
170    /// and cannot send into a replacement established by a later handshake.
171    ///
172    /// Sends reset and hello before draining old input. The adapter must allow
173    /// that output to finish without concurrent client reads for this attempt
174    /// to progress. Backpressure can instead fail an outgoing frame on timeout.
175    ///
176    /// The handshake uses one configured deadline, shared by output and peer
177    /// waits. Existing writer-lock cleanup may extend the call. If connecting
178    /// fails, the client has no session and previously issued senders are invalid.
179    pub fn connect<V: Verifier>(&mut self, verifier: &V) -> Result<(Sender<W>, V::Info), Error> {
180        // Compute the deadline by which the handshake must finish
181        let deadline = self.outbound.clock.now() + self.handshake_timeout;
182
183        // Generate ephemeral client keys for this session
184        let host_xdsa_sk = xdsa::SecretKey::generate();
185        let host_xhpke_sk = xhpke::SecretKey::generate();
186
187        self.handshake(verifier, host_xdsa_sk, host_xhpke_sk, None, deadline)
188    }
189
190    /// Ends any previous session, sends a reset and drives the handshake with the
191    /// given ephemeral keys. Returns the new session's sender and verified info.
192    /// The optional signing time makes the exchange deterministic for test vectors.
193    fn handshake<V: Verifier>(
194        &mut self,
195        verifier: &V,
196        host_xdsa_sk: xdsa::SecretKey,
197        host_xhpke_sk: xhpke::SecretKey,
198        timestamp: Option<i64>,
199        deadline: Instant,
200    ) -> Result<(Sender<W>, V::Info), Error> {
201        let host_xdsa_pk = host_xdsa_sk.public_key();
202        let host_xhpke_pk = host_xhpke_sk.public_key();
203        debug!("starting wire handshake");
204
205        // Message 1: Send HostHello (plain CBOR, COBS-framed)
206        let hello = cbor::encode(&handshake::HostHello {
207            host_signer: host_xdsa_pk.clone(),
208            host_crypto: host_xhpke_pk.clone(),
209        })
210        .map_err(|err| {
211            self.end_session();
212            Error::HandshakeFailed(format!("failed to encode client hello: {}", err))
213        })?;
214
215        // Retire the old binding and serialize reset/hello after admitted sends.
216        // The client starts reading only once this output has finished.
217        self.receiver = None;
218        self.sealer = None;
219        self.outbound.send_reset(deadline)?;
220        self.outbound.send_packet(&hello, Some(deadline))?;
221
222        // Message 2: Skip old replies, notifications and partial-frame leftovers
223        // until ArkHello names this attempt's fresh key. All draining shares the
224        // same deadline; authentication follows below.
225        let recipient = host_xhpke_pk.fingerprint();
226        let packet = loop {
227            let packet = match self.reader.next_packet(Some(deadline)) {
228                Ok(Some(packet)) => packet,
229                Ok(None) | Err(Error::FrameDecodingFailed(_) | Error::FrameTooLarge(_)) => &[],
230                Err(err) => return Err(err),
231            };
232            if cose::recipient(packet).is_ok_and(|fp| fp == recipient) {
233                break packet;
234            }
235            debug!("skipping stale frame during handshake");
236        };
237        let auth = handshake::ArkHelloAuth {
238            host_signer: host_xdsa_pk.clone(),
239            host_crypto: host_xhpke_pk.clone(),
240        };
241
242        // Step 2a: Decrypt the outer COSE_Encrypt0 layer
243        let sign1 =
244            cose::decrypt(packet, &auth, &host_xhpke_sk, CRYPTO_DOMAIN_WIRE).map_err(|err| {
245                Error::HandshakeFailed(format!("failed to decrypt server hello: {}", err))
246            })?;
247
248        // Step 2b: Peek at the unverified payload to discover the server's identity
249        let unverified: handshake::ArkHello = cose::peek(&sign1).map_err(|err| {
250            Error::HandshakeFailed(format!("invalid server hello payload: {}", err))
251        })?;
252
253        // Step 2c: Hand the attestation to the verifier to obtain the server's
254        // identity key and the caller's session info
255        let attestation = Attestation::new(unverified.ark_attest)?;
256        let (ark_identity, info) = verifier
257            .verify(&attestation, self.outbound.clock.system_time())
258            .map_err(Error::HandshakeFailed)?;
259
260        // Step 2d: Verify the COSE_Sign1 signature with the discovered identity
261        let ark_hello: handshake::ArkHello = cose::verify_at(
262            &sign1,
263            &auth,
264            &ark_identity,
265            CRYPTO_DOMAIN_WIRE,
266            None,
267            0, // unused without a drift check, so the clock is not read
268        )
269        .map_err(|err| {
270            Error::HandshakeFailed(format!("server hello signature invalid: {}", err))
271        })?;
272
273        // Set up the server->Client receiver context
274        let enc_a2h: [u8; xhpke::ENCAP_KEY_SIZE] = ark_hello
275            .a2h_encap
276            .try_into()
277            .map_err(|_| Error::HandshakeFailed("invalid a2h_encap size".into()))?;
278
279        let receiver = host_xhpke_sk
280            .new_receiver(&enc_a2h, CRYPTO_DOMAIN_WIRE_ARK_TO_HOST)
281            .map_err(|err| {
282                Error::HandshakeFailed(format!("client receiver setup failed: {}", err))
283            })?;
284
285        // Set up the Client->server sender context
286        let ark_xhpke_pk = ark_hello.ark_crypto;
287        let (sender, enc_h2a) = ark_xhpke_pk
288            .new_sender(CRYPTO_DOMAIN_WIRE_HOST_TO_ARK)
289            .map_err(|err| {
290                Error::HandshakeFailed(format!("client sender setup failed: {}", err))
291            })?;
292
293        // Message 3: Send HostAck (COSE seal'd, COBS-framed)
294        let ack = handshake::HostAck {
295            h2a_encap: enc_h2a.to_vec(),
296        };
297        let auth = handshake::HostAckAuth {
298            ark_signer: ark_identity,
299            ark_crypto: ark_xhpke_pk.clone(),
300        };
301        let ack = cose::seal_at(
302            &ack,
303            &auth,
304            &host_xdsa_sk,
305            &ark_xhpke_pk,
306            CRYPTO_DOMAIN_WIRE,
307            timestamp.unwrap_or_else(|| handshake::timestamp(&self.outbound.clock)),
308        )
309        .map_err(|err| Error::HandshakeFailed(format!("failed to seal client ack: {}", err)))?;
310
311        self.outbound.send_packet(&ack, Some(deadline))?;
312        check_deadline(&self.outbound.clock, deadline).map_err(Error::RecvFailed)?;
313
314        // Session established, the ack ahead of anything sealed into it
315        let sender = self.new_session(sender, receiver);
316        info!(
317            "wire session {} established with ark {}",
318            self.log_id,
319            hex(&auth.ark_signer.fingerprint())
320        );
321        Ok((sender, info))
322    }
323
324    /// Reads and decrypts the next ark-to-host message. Invalid or oversized
325    /// frames end the session because its encryption sequence may be lost.
326    /// Decryption failures also end the session.
327    /// An empty frame means the server dropped the session and ends it here too.
328    /// Adapter read failures and EOF also end the session. Idle read timeouts are
329    /// retried internally. Call [`Self::connect`] to establish a new session after failure.
330    /// Without a receive context, returns an error without reading the stream.
331    ///
332    /// After decryption, message acceptance is ordered with session ending
333    /// without waiting for the writer. A concurrent send failure can cause a
334    /// decrypted message to be discarded before acceptance. An accepted message
335    /// may reach this caller after another thread ends the session. Returning
336    /// a receive error does wait for outgoing writes to finish.
337    pub fn recv(&mut self) -> Result<Vec<u8>, Error> {
338        let receiver = self
339            .receiver
340            .as_mut()
341            .ok_or_else(|| Error::EncryptionFailed("no active session".into()))?;
342
343        // Retrieve the next COBS encoded packet. A skipped frame may have
344        // carried a sealed message, so the session cannot continue past it.
345        // An empty frame is the server telling us it has no session with us.
346        let packet = match self.reader.next_packet(None) {
347            Err(err) => {
348                if matches!(err, Error::FrameDecodingFailed(_) | Error::FrameTooLarge(_)) {
349                    warn!("ending session {}: {}", self.log_id, err);
350                }
351                self.end_session();
352                return Err(err);
353            }
354            Ok(None) => {
355                info!("wire session {} reset by ark", self.log_id);
356                self.end_session();
357                return Err(Error::SessionReset);
358            }
359            Ok(Some(packet)) => packet,
360        };
361        let sealer = self
362            .sealer
363            .as_ref()
364            .expect("receiver has a sending context");
365        let opened = sealing::open(receiver, packet);
366        let undecryptable = opened.is_err();
367        let message = match self.outbound.finish_receive(sealer, opened) {
368            Err(err) => {
369                if undecryptable {
370                    warn!("ending session {}: {}", self.log_id, err);
371                } else {
372                    debug!(
373                        "discarding message read after session {} ended",
374                        self.log_id
375                    );
376                }
377                self.end_session();
378                return Err(err);
379            }
380            Ok(message) => message,
381        };
382        trace!("received ark-to-host message ({} bytes)", packet.len());
383        Ok(message)
384    }
385
386    /// Stores the negotiated contexts and returns a sender for the new session.
387    /// The sending context's allocation identifies the session. The client owns
388    /// both contexts and shares the sending context with active sends. Idle
389    /// senders hold weak references and keep neither context nor stream alive.
390    ///
391    /// Takes the writer lock, then the binding lock. An old write that already
392    /// holds the writer lock may finish first. Once the binding is replaced,
393    /// old sends cannot write and old received messages cannot be accepted.
394    /// This method performs no handshake, crypto or stream I/O.
395    fn new_session(&mut self, sender: xhpke::Sender, receiver: xhpke::Receiver) -> Sender<W> {
396        let sealer = Arc::new(Mutex::new(sender));
397        let sender = self.outbound.bind(&sealer);
398        self.log_id = sender.log_id();
399
400        self.receiver = Some(receiver);
401        self.sealer = Some(sealer);
402
403        sender
404    }
405
406    /// Ends the current binding before releasing the client's crypto contexts.
407    /// Waits for the writer. After this returns, no write or flush for that
408    /// session is running or can start. A send that gets the writer first may
409    /// finish. A send still sealing after removal cannot write its packet.
410    /// This takes no encryption lock and does not wait for crypto work.
411    ///
412    /// This does not close the stream or send a notification. An active write
413    /// may delay ending until its frame deadline. Another thread can use the
414    /// Closer to cancel I/O without taking the writer lock.
415    fn end_session(&mut self) {
416        if let Some(sealer) = self.sealer.as_ref() {
417            self.outbound.end(sealer);
418        }
419        self.receiver = None;
420        self.sealer = None;
421    }
422
423    /// Observes a handshake reaching writer acquisition behind an already admitted send.
424    #[cfg(any(test, feature = "fuzz"))]
425    pub(crate) fn watch_writer(&self) -> impl FnOnce() + use<R, W> {
426        let outbound = self.outbound.clone();
427        move || outbound.wait_writers(1)
428    }
429
430    /// Runs a test handshake with fixed keys and signing time for vector replay.
431    /// Not part of the normal transport API.
432    #[doc(hidden)]
433    #[inline]
434    #[cfg(any(test, feature = "bench", feature = "fuzz"))]
435    #[cfg_attr(coverage_nightly, coverage(off))]
436    pub fn handshake_with_keys<V: Verifier>(
437        &mut self,
438        verifier: &V,
439        host_xdsa_sk: xdsa::SecretKey,
440        host_xhpke_sk: xhpke::SecretKey,
441        timestamp: i64,
442    ) -> Result<(Sender<W>, V::Info), Error> {
443        self.handshake(
444            verifier,
445            host_xdsa_sk,
446            host_xhpke_sk,
447            Some(timestamp),
448            self.outbound.clock.now() + self.handshake_timeout,
449        )
450    }
451
452    /// Reads a framed packet without decryption for tests and benchmarks.
453    #[doc(hidden)]
454    #[inline]
455    #[cfg(any(test, feature = "bench", feature = "fuzz"))]
456    #[cfg_attr(coverage_nightly, coverage(off))]
457    pub fn next_packet_blob(&mut self) -> Result<Option<&[u8]>, Error> {
458        self.reader.next_packet(None)
459    }
460
461    /// Writes a packet without encryption for tests and benchmarks.
462    #[doc(hidden)]
463    #[inline]
464    #[cfg(any(test, feature = "bench", feature = "fuzz"))]
465    #[cfg_attr(coverage_nightly, coverage(off))]
466    pub fn send_packet_blob(&mut self, packet: &[u8]) -> Result<(), Error> {
467        self.outbound.send_packet(packet, None)
468    }
469
470    /// Reads an encoded frame without its delimiter for tests and benchmarks.
471    #[doc(hidden)]
472    #[inline]
473    #[cfg(any(test, feature = "bench", feature = "fuzz"))]
474    #[cfg_attr(coverage_nightly, coverage(off))]
475    pub fn next_frame_blob(&mut self) -> Result<&[u8], Error> {
476        self.reader.next_frame_blob()
477    }
478
479    /// Writes an already encoded frame with a delimiter for tests and benchmarks.
480    #[doc(hidden)]
481    #[inline]
482    #[cfg(any(test, feature = "bench", feature = "fuzz"))]
483    #[cfg_attr(coverage_nightly, coverage(off))]
484    pub fn send_frame_blob(&mut self, frame: &[u8]) -> Result<(), Error> {
485        self.outbound.send_frame_blob(frame)
486    }
487}
488
489impl<R: Read, W: Write> Drop for Client<R, W> {
490    /// Closes the stream to cancel blocked I/O, then ends the binding before
491    /// releasing the contexts. Shutdown must precede waiting for the writer.
492    /// Idle senders hold weak references and cannot extend the stream's lifetime.
493    fn drop(&mut self) {
494        self.outbound.close();
495        self.end_session();
496    }
497}
498
499impl<R: Read, W: Write> fmt::Debug for Client<R, W> {
500    /// Shows the session label, whether a session is established and the
501    /// handshake budget, never the adapters or the encryption contexts.
502    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
503        f.debug_struct("Client")
504            .field("session", &self.log_id)
505            .field("connected", &self.sealer.is_some())
506            .field("handshake_timeout", &self.handshake_timeout)
507            .finish_non_exhaustive()
508    }
509}
510
511/// Hex encodes a fingerprint for the session log line.
512fn hex(fingerprint: &xdsa::Fingerprint) -> String {
513    fingerprint
514        .to_bytes()
515        .iter()
516        .map(|byte| format!("{byte:02x}"))
517        .collect()
518}
519
520#[cfg(test)]
521#[cfg_attr(coverage_nightly, coverage(off))]
522mod tests {
523    use super::*;
524    use crate::transport::DEFAULT_WRITE_TIMEOUT;
525    use crate::transport::framing::FrameWriter;
526    use crate::transport::mock::{payload, self_attestation};
527    use crate::transport::server::Server;
528    use crate::transport::testing::{Memory, test_clock};
529    use crate::{memory, testing};
530    use darkbio_clock::Clock;
531    use std::io::{self, Read as _};
532    use std::sync::mpsc;
533    use std::thread;
534    use std::time::Instant;
535
536    /// A pair of contexts standing in for an established session.
537    fn contexts() -> (xhpke::Sender, xhpke::Receiver) {
538        let secret = xhpke::SecretKey::generate();
539        let (sender, encap) = secret.public_key().new_sender(b"test").unwrap();
540        let receiver = secret.new_receiver(&encap, b"test").unwrap();
541        (sender, receiver)
542    }
543
544    /// Writer accepting bytes immediately but holding its first flush until
545    /// the test releases it. This distinguishes finished writes from a fully
546    /// completed send, which must also wait for its flush.
547    struct BlockedFlush {
548        /// Clock shared with the client and its release gate.
549        clock: Clock,
550        entered: Option<mpsc::Sender<()>>,
551        release: testing::Gate,
552        deadline: Option<Instant>,
553    }
554
555    impl Write for BlockedFlush {
556        fn clock(&self) -> Clock {
557            self.clock.clone()
558        }
559
560        fn set_write_deadline(&mut self, deadline: Instant) -> io::Result<()> {
561            self.deadline = Some(deadline);
562            Ok(())
563        }
564    }
565
566    impl io::Write for BlockedFlush {
567        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
568            testing::remaining(
569                &self.clock,
570                self.deadline.expect("write deadline installed"),
571            )?;
572            Ok(bytes.len())
573        }
574
575        fn flush(&mut self) -> io::Result<()> {
576            let deadline = self.deadline.expect("write deadline installed");
577            testing::remaining(&self.clock, deadline)?;
578            if let Some(entered) = self.entered.take() {
579                entered.send(()).unwrap();
580                self.release.wait(Some(deadline))?;
581            }
582            Ok(())
583        }
584    }
585
586    // Tests that ending waits through an active flush. After it returns the
587    // old sender is refused, and fresh contexts can use the still-open stream.
588    #[test]
589    fn test_end_waits_for_flush() {
590        // Hold the first send in a flush on the client's clock
591        testing::init_tracing();
592        let tester = test_clock();
593        let clock = tester.clock();
594        let (entered_tx, entered) = mpsc::channel();
595        let release = testing::Gate::new(&clock);
596        let mut client = Client::new(Stream::new(
597            Memory::new(io::empty(), &clock),
598            BlockedFlush {
599                clock: clock.clone(),
600                entered: Some(entered_tx),
601                release: release.clone(),
602                deadline: None,
603            },
604            || {},
605        ));
606        let (crypto, receiver) = contexts();
607        let sender = client.new_session(crypto, receiver);
608        let sending = {
609            let sender = sender.clone();
610            thread::spawn(move || sender.send(&payload(1)))
611        };
612        entered.recv().unwrap();
613
614        // Attempt session ending while the flush remains parked
615        let (started_tx, started) = mpsc::channel();
616        let (ended_tx, ended) = mpsc::channel();
617        let outbound = client.outbound.clone();
618        let ending = thread::spawn(move || {
619            started_tx.send(()).unwrap();
620            client.end_session();
621            ended_tx.send(()).unwrap();
622            client
623        });
624        started.recv().unwrap();
625        tester.wait_blocked(1);
626        outbound.wait_writers(1);
627        assert!(ended.try_recv().is_err());
628
629        // Release output and require the old session to refuse further sends
630        release.open();
631        sending.join().unwrap().unwrap();
632        let mut client = ending.join().unwrap();
633        ended.recv().unwrap();
634        assert!(matches!(
635            sender.send(&payload(2)),
636            Err(Error::EncryptionFailed(_))
637        ));
638
639        // Establish fresh contexts on the still-open stream
640        let (crypto, receiver) = contexts();
641        let fresh = client.new_session(crypto, receiver);
642        fresh.send(&payload(3)).unwrap();
643        assert!(matches!(
644            sender.send(&payload(4)),
645            Err(Error::EncryptionFailed(_))
646        ));
647    }
648
649    // Tests that a real receive reads, decrypts and returns a message while an
650    // outgoing flush is blocked. Waiting for the writer on the success path
651    // would time out before the test releases that flush.
652    #[test]
653    fn test_recv_during_blocked_flush() {
654        // Prepare an encrypted incoming frame on a paused clock
655        testing::init_tracing();
656        let tester = test_clock();
657        let clock = tester.clock();
658        let (mut peer, receiver) = contexts();
659        let packet = sealing::seal(&mut peer, &payload(1)).unwrap();
660        let mut bytes = Vec::new();
661        FrameWriter::new(Memory::new(&mut bytes, &clock), Closer::new(&clock, || {}))
662            .send_packet(&packet, clock.now() + DEFAULT_WRITE_TIMEOUT)
663            .unwrap();
664        let (entered_tx, entered) = mpsc::channel();
665        let release = testing::Gate::new(&clock);
666        let mut client = Client::new(Stream::new(
667            Memory::new(io::Cursor::new(bytes), &clock),
668            BlockedFlush {
669                clock: clock.clone(),
670                entered: Some(entered_tx),
671                release: release.clone(),
672                deadline: None,
673            },
674            || {},
675        ));
676        let sender = client.new_session(contexts().0, receiver);
677        let sending = thread::spawn(move || sender.send(&payload(2)));
678        entered.recv().unwrap();
679
680        // Receive the message while the outgoing flush is parked
681        tester.wait_blocked(1);
682        let (received_tx, received) = mpsc::channel();
683        let receiving = thread::spawn(move || {
684            received_tx.send(client.recv()).unwrap();
685            client
686        });
687        let result = received.recv();
688
689        // Release output only after receiving has completed
690        release.open();
691        sending.join().unwrap().unwrap();
692        let _client = receiving.join().unwrap();
693        assert_eq!(result.unwrap().unwrap(), payload(1));
694    }
695
696    // Tests that releasing an old session's contexts cannot invalidate its
697    // replacement, which can still receive and send. Dropping the client ends
698    // the replacement even while active operations retain its sending context
699    // and outbound transport, modeled here by retaining those references.
700    #[test]
701    fn test_owner_drop() {
702        // Prepare a replacement session's incoming packet on a paused clock
703        testing::init_tracing();
704        let tester = test_clock();
705        let clock = tester.clock();
706        let (mut peer, receiver) = contexts();
707        let packet = sealing::seal(&mut peer, &payload(2)).unwrap();
708        let mut bytes = Vec::new();
709        FrameWriter::new(Memory::new(&mut bytes, &clock), Closer::new(&clock, || {}))
710            .send_packet(&packet, clock.now() + DEFAULT_WRITE_TIMEOUT)
711            .unwrap();
712        let mut client = Client::new(Stream::new(
713            Memory::new(&bytes[..], &clock),
714            Memory::new(Vec::new(), &clock),
715            || {},
716        ));
717        let (crypto, old_receiver) = contexts();
718        let stale = client.new_session(crypto, old_receiver);
719        let old_sealer = client.sealer.as_ref().unwrap().clone();
720
721        // End the old contexts without disturbing their replacement
722        let fresh = client.new_session(contexts().0, receiver);
723        client.outbound.end(&old_sealer);
724        drop(old_sealer);
725        assert!(matches!(
726            stale.send(&payload(1)),
727            Err(Error::EncryptionFailed(_))
728        ));
729        assert_eq!(client.recv().unwrap(), payload(2));
730        fresh.send(&payload(3)).unwrap();
731
732        // Drop the owner while retaining internal references to its session
733        let outbound = client.outbound.clone();
734        let sealer = client.sealer.as_ref().unwrap().clone();
735        drop(client);
736        assert!(outbound.finish_receive(&sealer, Ok(Vec::new())).is_err());
737        assert!(matches!(
738            fresh.send(&payload(4)),
739            Err(Error::EncryptionFailed(_))
740        ));
741    }
742
743    // Tests sending from other threads while the client blocks in a read.
744    // The server must receive every message in encryption order to decrypt
745    // and echo it successfully.
746    #[test]
747    fn test_senders() {
748        testing::init_tracing();
749
750        // Echo every request over a bounded in-memory stream, then hang up
751        let tester = test_clock();
752        let (host, ark_stream) = memory::duplex(64 * 1024, &tester.clock());
753
754        let signer = xdsa::SecretKey::generate();
755        let identity = signer.public_key();
756        let attestation = self_attestation(&signer);
757        let ark = thread::spawn(move || {
758            let mut server = Server::new(ark_stream, signer, attestation);
759            let mut sender = None;
760            for _ in 0..100 {
761                let req = testing::served(&mut server, &mut sender).unwrap();
762                sender.as_ref().unwrap().send(&req).unwrap();
763            }
764        });
765        let mut client = Client::new(host);
766        let (sender, _) = client.connect(&identity).unwrap();
767
768        // Send from a few threads at once while reading the echoes on this one
769        let senders: Vec<_> = (0..4)
770            .map(|thread| {
771                let sender = sender.clone();
772                thread::spawn(move || {
773                    for i in 0..25 {
774                        sender.send(&payload(thread * 100 + i)).unwrap();
775                    }
776                })
777            })
778            .collect();
779        let mut echoes: Vec<Vec<u8>> = (0..100).map(|_| client.recv().unwrap()).collect();
780        for sender in senders {
781            sender.join().unwrap();
782        }
783        ark.join().unwrap();
784
785        echoes.sort_unstable();
786        let mut expected: Vec<Vec<u8>> = (0..4)
787            .flat_map(|thread| (0..25).map(move |i| payload(thread * 100 + i)))
788            .collect();
789        expected.sort_unstable();
790        assert_eq!(echoes, expected);
791    }
792
793    // Tests that dropping the client ends the session for its senders and
794    // releases the transport writer even while sender handles remain.
795    #[test]
796    fn test_sender_outlives_client() {
797        // Send one packet before dropping the client owner
798        testing::init_tracing();
799        let tester = test_clock();
800        let clock = tester.clock();
801        let (mut reader, writer) = testing::pipe(&clock);
802        let (sender, receiver) = contexts();
803        let mut client = Client::new(Stream::new(Memory::new(io::empty(), &clock), writer, || {}));
804        let sender = client.new_session(sender, receiver);
805        sender.send(&payload(1)).unwrap();
806        drop(client);
807
808        // Refuse later sends and drain the closed writer to EOF
809        let result = sender.send(&payload(2));
810        assert!(matches!(result, Err(Error::Terminated)), "{result:?}");
811        // The read only returns once the writer is gone
812        let mut bytes = Vec::new();
813        reader
814            .set_read_deadline(Some(clock.now() + DEFAULT_WRITE_TIMEOUT))
815            .unwrap();
816        reader.read_to_end(&mut bytes).unwrap();
817        assert!(!bytes.is_empty());
818    }
819}