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