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