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