Skip to main content

darkbio_wire/protocol/
session.rs

1// wire-rs: encrypted protocol between Ark and host
2// Copyright 2026 Dark Bio AG. All rights reserved.
3
4//! Session state, request queues, and the reader, writer, and deadline workers.
5
6use super::envelope::{Header, IncomingEnvelope, MessageKind, Parity, Side};
7use super::operation::{
8    OperationHandle, OperationKey, OutgoingBody, OutgoingMessage, PendingOperation,
9};
10use super::promise::PromiseResult;
11use super::worker;
12use super::{
13    Closer, DEFAULT_ABANDONMENT_TIMEOUT, DEFAULT_MAX_INBOUND_BYTES, DEFAULT_MAX_INBOUND_REQUESTS,
14    Error, Message, Promise, Requester, Responder, schema,
15};
16use crate::LogId;
17use crate::transport::{self, Read, Stream, Verifier, Write};
18use prost::bytes::Bytes;
19use std::collections::{HashMap, HashSet, VecDeque};
20use std::fmt;
21use std::sync::atomic::{AtomicUsize, Ordering};
22use std::sync::{Arc, Condvar, Mutex};
23use std::time::{Duration, Instant};
24
25/// Establishes a client session, verifies the peer and returns the verifier's info.
26/// Takes ownership of the stream and constructs the transport internally. Failure
27/// closes the client stream; the application can open another stream and reconnect.
28/// The verifier is only borrowed during this blocking call.
29/// Failure to start a required worker or an escaping worker panic aborts the process.
30pub fn connect<R, W, V>(stream: Stream<R, W>, verifier: &V) -> Result<(Session, V::Info), Error>
31where
32    R: Read + Send + 'static,
33    W: Write + Send + 'static,
34    V: Verifier,
35{
36    let mut client = transport::Client::new(stream);
37
38    let (sender, info) = client.connect(verifier)?;
39    #[cfg(any(test, feature = "fuzz"))]
40    let workers = Arc::new(worker::Tracker::default());
41    let session = Session::start(
42        Side::Client,
43        sender,
44        Some(client.closer()),
45        #[cfg(any(test, feature = "fuzz"))]
46        workers.clone(),
47    );
48    let inner = session.inner.clone();
49
50    worker::spawn(
51        "wire-client-reader",
52        #[cfg(any(test, feature = "fuzz"))]
53        &workers,
54        move || run_reader(client, inner),
55    );
56    Ok((session, info))
57}
58
59/// Receives and handles messages for one client session. The reader retains its
60/// session state until it exits; closing the session shuts down the stream and
61/// wakes a blocked transport read.
62fn run_reader<R: Read, W: Write>(mut client: transport::Client<R, W>, session: Arc<SessionInner>) {
63    loop {
64        let result = client
65            .recv()
66            .map_err(Error::from)
67            .and_then(|bytes| session.handle_message(bytes));
68        if let Err(error) = result {
69            // Receive and decode failures close this client session and its
70            // stream, waking any other blocked transport I/O.
71            session.close(error);
72            break;
73        }
74    }
75}
76
77/// Owner of one session and its incoming request queue.
78///
79/// Incoming requests use [`Message`], paired with a common responder. The
80/// application selects the reply type; there is no static request/response pairing
81/// table. The host/server role and envelope direction are handled internally.
82///
83/// Closing or dropping the session fails pending promises, discards queued requests,
84/// and wakes blocked `recv()` calls. Completed promises keep their results.
85/// Application jobs keep running, but their handles still refer to the closed
86/// session and cannot send messages through a replacement session.
87///
88/// A client session also closes its stream. A server session leaves the server's
89/// stream available for another handshake. Handles do not keep the session open.
90/// The owner cannot be cloned; obtain requesters or closers for other threads:
91///
92/// ```compile_fail,E0599
93/// use darkbio_wire::protocol::Session;
94/// fn duplicate(session: Session) { let _ = session.clone(); }
95/// ```
96pub struct Session {
97    /// Queues and pending operations shared with this session's workers.
98    pub(super) inner: Arc<SessionInner>,
99}
100
101impl Session {
102    /// Sets the lifetime of automatic `UNANSWERED` replies when responders are
103    /// subsequently dropped. Defaults to [`DEFAULT_ABANDONMENT_TIMEOUT`]. Replies
104    /// already queued keep their deadlines.
105    ///
106    /// The budget starts when the responder is dropped and includes queueing.
107    /// Expiry discards a queued reply; a write already started still runs under the
108    /// transport's independent timeout. Explicit request/reply deadlines are
109    /// unaffected. Zero or an unrepresentable deadline expires immediately.
110    /// Use [`super::Server::set_abandonment_timeout`] to also set the timeout for
111    /// future server sessions.
112    pub fn set_abandonment_timeout(self, timeout: Duration) -> Self {
113        self.inner.set_abandonment_timeout(timeout);
114        self
115    }
116
117    /// Sets the maximum accepted peer requests and buffered incoming bytes together.
118    /// Defaults to [`DEFAULT_MAX_INBOUND_REQUESTS`] and [`DEFAULT_MAX_INBOUND_BYTES`].
119    ///
120    /// `requests` counts queued requests, held responders and queued replies.
121    /// A slot is freed when the writer takes the reply or the reply is discarded.
122    /// Zero refuses all peer requests but still allows responses to our requests.
123    ///
124    /// `bytes` counts the full encoded envelopes of queued requests and unread
125    /// responses. `recv()`, `wait()` or dropping a response promise releases that
126    /// space in the budget. Zero allows no buffered envelopes. Decoded application
127    /// data, outgoing messages and transport buffers are excluded.
128    ///
129    /// Exceeding either limit closes this session with
130    /// [`Error::InboundRequestLimitExceeded`] or [`Error::InboundByteLimitExceeded`].
131    /// The reader never waits for space. Lowering a limit below usage also closes
132    /// the session. Completed promises keep their results and bytes until read or
133    /// dropped. Raising limits does not reopen a closed session.
134    pub fn set_inbound_limits(self, requests: usize, bytes: usize) -> Self {
135        self.inner.set_inbound_limits(requests, bytes);
136        self
137    }
138
139    /// Returns a clonable requester bound to this session.
140    pub fn requester(&self) -> Requester {
141        Requester::new(Arc::downgrade(&self.inner))
142    }
143
144    /// Blocks for the next peer request and its [`Responder`]. Closing the session
145    /// wakes this call with the error that closed it. The caller decides how to
146    /// handle each request; this method does not run application callbacks.
147    ///
148    /// Taking a request removes its bytes from the inbound byte count before
149    /// decoding it. The request still counts toward the inbound request limit
150    /// while its responder is held.
151    /// Invalid protobuf returns [`Error::Malformed`] and closes this session.
152    ///
153    /// If another thread closes the session after `recv()` takes a request from
154    /// the queue, `recv()` can still return it. Replying after closure returns an error.
155    pub fn recv(&mut self) -> Result<(Message, Responder), Error> {
156        self.inner.recv()
157    }
158
159    /// Returns a clonable handle for closing this session from another thread,
160    /// including while its owner is blocked in [`Self::recv`].
161    pub fn closer(&self) -> Closer {
162        Closer::session(Arc::downgrade(&self.inner))
163    }
164
165    /// Closes this session. Repeated calls have no further effect. This does not
166    /// wait for application jobs or guarantee the peer has observed closure.
167    /// Discards queued messages and fails pending promises. A transport write
168    /// already started may still finish, but cannot change a completed promise's
169    /// result or affect a replacement session.
170    pub fn close(&self) {
171        self.inner.close(Error::Closed);
172    }
173}
174
175impl Drop for Session {
176    /// Closes the session even when requesters, responders or closers remain.
177    fn drop(&mut self) {
178        self.close();
179    }
180}
181
182impl fmt::Debug for Session {
183    /// Shows the session label and whether the session is still open. A state
184    /// lock held elsewhere leaves the state out.
185    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
186        let mut session = f.debug_struct("Session");
187        session.field("id", &self.inner.log_id);
188        if let Ok(state) = self.inner.state.try_lock() {
189            session.field("open", &matches!(*state, State::Open { .. }));
190        }
191        session.finish_non_exhaustive()
192    }
193}
194
195/// Queues and pending operations for one session. Requesters, responders, and
196/// closers hold weak references to this object, even after a new session connects.
197pub(super) struct SessionInner {
198    /// Protects the queues, pending operations, and transition to `State::Closed`.
199    state: Mutex<State>,
200    /// Wakes `recv()`, the writer, and the deadline worker when `state` changes.
201    changed: Condvar,
202    /// Incoming bytes held by queued requests and unread responses. Accepting
203    /// messages and changing limits hold `state`. Consumers can release bytes
204    /// without that lock. Unread promises keep this counter alive after closure.
205    retained_bytes: Arc<AtomicUsize>,
206    /// Envelope direction and request parity, fixed for the whole session.
207    side: Side,
208    /// Label of the transport session in log lines, unset without a transport.
209    pub(super) log_id: LogId,
210    /// Closes the client's stream. Server sessions and tests without a stream
211    /// leave this empty; a server's stream is closed by `Server`.
212    stream_closer: Option<transport::Closer>,
213    /// Lets tests wait for worker threads to exit.
214    #[cfg(any(test, feature = "fuzz"))]
215    pub(super) workers: Arc<worker::Tracker>,
216    /// Controlled protocol time for scenarios; production always uses Instant::now.
217    #[cfg(any(test, feature = "fuzz"))]
218    time: Mutex<Option<Instant>>,
219    /// Notifies tests when the last `Arc<SessionInner>` is dropped.
220    #[cfg(any(test, feature = "fuzz"))]
221    drop_hook: Mutex<Option<std::sync::mpsc::Sender<()>>>,
222    /// Pauses the writer before `sender.disconnect()` in replacement tests.
223    #[cfg(any(test, feature = "fuzz"))]
224    disconnect_hook: Mutex<Option<(std::sync::mpsc::Sender<()>, std::sync::mpsc::Receiver<()>)>>,
225}
226
227/// An open session's queues, or the error that closed the session.
228// Keep the same inline state layout in tests; the wait hook crosses Clippy's
229// size threshold for the difference between variants.
230#[cfg_attr(any(test, feature = "fuzz"), allow(clippy::large_enum_variant))]
231enum State {
232    /// Holds queued messages and pending operations. `close()` replaces this
233    /// with `Closed`, then drops the queues after releasing the state lock.
234    Open {
235        /// Ceiling for accepted requests, including application-held responders.
236        max_inbound_requests: usize,
237        /// Ceiling for encoded requests and unread response promises.
238        max_inbound_bytes: usize,
239        /// Automatic reply lifetime selected when a responder is dropped.
240        abandonment: Duration,
241
242        /// Peer requests awaiting application receipt, paired with their request IDs.
243        incoming: VecDeque<(u64, IncomingEnvelope)>,
244        /// Rejects duplicate incoming IDs while the receive queue, responder,
245        /// or queued reply holds them. Taking a reply into the writer or
246        /// discarding an expired reply releases its ID. Release happens before
247        /// writing: the peer can receive the reply and reuse its ID before our
248        /// local flush returns.
249        reserved_ids: HashSet<u64>,
250
251        /// Requests and replies waiting for the writer to take them.
252        outgoing: VecDeque<OutgoingMessage>,
253        /// Next locally allocated ID, or exhaustion. Never wraps or reuses an ID.
254        next_id: Option<u64>,
255        /// Maps outgoing request IDs to operation keys. Entries remain after
256        /// a promise times out, until the peer responds or the session closes.
257        outstanding: HashMap<u64, OperationKey>,
258        /// Operations waiting to send a result to their promise. Each is removed
259        /// before sending that result, so a promise is completed only once.
260        operations: HashMap<OperationKey, PendingOperation>,
261
262        /// One-shot test notification sent under the state lock before waiting.
263        #[cfg(any(test, feature = "fuzz"))]
264        wait_hook: Option<std::sync::mpsc::Sender<()>>,
265    },
266    /// Keeps the first closing reason for every subsequent operation.
267    Closed(Error),
268}
269
270impl SessionInner {
271    /// Creates empty queues and pending-operation maps before starting workers.
272    fn new(
273        side: Side,
274        log_id: LogId,
275        stream_closer: Option<transport::Closer>,
276        #[cfg(any(test, feature = "fuzz"))] workers: Arc<worker::Tracker>,
277    ) -> Self {
278        Self {
279            state: Mutex::new(State::Open {
280                max_inbound_requests: DEFAULT_MAX_INBOUND_REQUESTS,
281                max_inbound_bytes: DEFAULT_MAX_INBOUND_BYTES,
282                abandonment: DEFAULT_ABANDONMENT_TIMEOUT,
283
284                incoming: VecDeque::new(),
285                reserved_ids: HashSet::new(),
286
287                outgoing: VecDeque::new(),
288                next_id: Some(Parity::from(side).first()),
289                outstanding: HashMap::new(),
290                operations: HashMap::new(),
291
292                #[cfg(any(test, feature = "fuzz"))]
293                wait_hook: None,
294            }),
295            changed: Condvar::new(),
296            retained_bytes: Arc::new(AtomicUsize::new(0)),
297            side,
298            log_id,
299            stream_closer,
300            #[cfg(any(test, feature = "fuzz"))]
301            workers,
302            #[cfg(any(test, feature = "fuzz"))]
303            time: Mutex::new(None),
304            #[cfg(any(test, feature = "fuzz"))]
305            drop_hook: Mutex::new(None),
306            #[cfg(any(test, feature = "fuzz"))]
307            disconnect_hook: Mutex::new(None),
308        }
309    }
310
311    /// Changes the timeout used when responders are dropped. Queued replies keep
312    /// their original deadlines.
313    pub(super) fn set_abandonment_timeout(&self, timeout: Duration) {
314        let mut state = self.state.lock().expect("session state not poisoned");
315        if let State::Open { abandonment, .. } = &mut *state {
316            *abandonment = timeout;
317        }
318    }
319
320    /// Updates both limits and closes the session if usage is too high. Holds the
321    /// same lock used to accept incoming messages.
322    pub(super) fn set_inbound_limits(&self, requests: usize, bytes: usize) {
323        let removed = {
324            let mut state = self.state.lock().expect("session state not poisoned");
325            let State::Open {
326                max_inbound_requests,
327                max_inbound_bytes,
328                reserved_ids,
329                ..
330            } = &mut *state
331            else {
332                return;
333            };
334            *max_inbound_requests = requests;
335            *max_inbound_bytes = bytes;
336            let error = if reserved_ids.len() > *max_inbound_requests {
337                Some(Error::InboundRequestLimitExceeded(*max_inbound_requests))
338            } else if self.retained_bytes.load(Ordering::Relaxed) > *max_inbound_bytes {
339                Some(Error::InboundByteLimitExceeded(*max_inbound_bytes))
340            } else {
341                None
342            };
343            error.and_then(|error| {
344                let (pending, queued) = state.workload();
345                let removed = state.close(error.clone(), self.now());
346                if removed.is_some() {
347                    self.log_closed(&error, pending, queued);
348                }
349                removed
350            })
351        };
352        if removed.is_some() {
353            self.finish_close(removed);
354        }
355    }
356
357    /// Counts the envelope's original bytes under the session lock. Decoding or
358    /// dropping it later releases those bytes to this session's counter.
359    fn retain_incoming(
360        self: &Arc<Self>,
361        bytes: Bytes,
362        header: Header,
363        limit: usize,
364    ) -> Result<IncomingEnvelope, Error> {
365        IncomingEnvelope::new(
366            bytes,
367            header,
368            &self.retained_bytes,
369            limit,
370            self.side,
371            Arc::downgrade(self),
372        )
373    }
374
375    /// Takes the next request from `incoming` and creates its responder. If the
376    /// queue is empty, waits on `changed`. Closing the session wakes the wait and
377    /// returns the error stored in `State::Closed`.
378    /// Decodes after releasing the lock so the reader can keep receiving messages.
379    fn recv(self: &Arc<Self>) -> Result<(Message, Responder), Error> {
380        let mut state = self.state.lock().expect("session state not poisoned");
381        let (message, responder) = loop {
382            match &mut *state {
383                State::Closed(error) => return Err(error.clone()),
384                State::Open {
385                    incoming,
386                    #[cfg(any(test, feature = "fuzz"))]
387                    wait_hook,
388                    ..
389                } => {
390                    if let Some((id, message)) = incoming.pop_front() {
391                        break (message, Responder::new(Arc::downgrade(self), id));
392                    }
393                    #[cfg(any(test, feature = "fuzz"))]
394                    if let Some(wait_hook) = wait_hook.take() {
395                        let _ = wait_hook.send(());
396                    }
397                    state = self
398                        .changed
399                        .wait(state)
400                        .expect("session state not poisoned");
401                }
402            }
403        };
404        drop(state);
405        Ok((message.decode()?, responder))
406    }
407
408    /// Replaces `Open` with `Closed` and fails pending promises under the state lock.
409    /// Expired operations receive `Timeout`; the rest receive the closing error.
410    /// Every call wakes `changed` and finishes any required stream shutdown.
411    pub(super) fn close(&self, error: Error) {
412        let (removed, pending, queued) = {
413            let mut state = self.state.lock().expect("session state not poisoned");
414            let (pending, queued) = state.workload();
415            (state.close(error.clone(), self.now()), pending, queued)
416        };
417        if removed.is_some() {
418            self.log_closed(&error, pending, queued);
419        }
420        self.finish_close(removed);
421    }
422
423    /// Logs the end of the session with its reason, the operations it failed
424    /// and the messages it dropped. Orderly endings are informational.
425    fn log_closed(&self, error: &Error, pending: usize, queued: usize) {
426        match error {
427            Error::Closed => tracing::info!(
428                "wire session {} closed locally (pending {}, queued {})",
429                self.log_id,
430                pending,
431                queued
432            ),
433            _ if error.orderly() => tracing::info!(
434                "wire session {} closed: {} (pending {}, queued {})",
435                self.log_id,
436                error.reason(),
437                pending,
438                queued
439            ),
440            _ => tracing::warn!(
441                "wire session {} failed: {} (pending {}, queued {})",
442                self.log_id,
443                error.reason(),
444                pending,
445                queued
446            ),
447        }
448    }
449
450    /// Drops queued work and wakes waiters outside the session lock.
451    fn finish_close(&self, removed: Option<State>) {
452        // Wake local waiters and drop queued work before adapter shutdown, which
453        // may wait for transport I/O already running to return.
454        self.changed.notify_all();
455        drop(removed);
456        if let Some(stream_closer) = &self.stream_closer {
457            stream_closer.close();
458        }
459    }
460
461    /// Queues an `UNANSWERED` reply when a responder is dropped. The
462    /// budget starts on entry, before acquiring the session lock, and includes
463    /// queueing. An unrepresentable deadline expires immediately instead of
464    /// panicking from `Drop`. The writer sends the reply later.
465    pub(super) fn reply_unanswered(self: &Arc<Self>, id: u64) {
466        let now = self.now();
467        let timeout = {
468            let state = self.state.lock().expect("session state not poisoned");
469            match &*state {
470                State::Open {
471                    abandonment: abandonment_timeout,
472                    ..
473                } => *abandonment_timeout,
474                State::Closed(_) => return,
475            }
476        };
477        tracing::debug!("answering request {} as unanswered", id);
478
479        // Fix this reply's deadline at drop. Later changes to the session's
480        // configuration do not retime already submitted work.
481        let deadline = now.checked_add(timeout).unwrap_or(now);
482        let _ = self.reply(
483            id,
484            Err(schema::Error::reserved(
485                schema::ReservedErrors::Unanswered,
486                "request left unanswered",
487            )),
488            deadline,
489        );
490    }
491
492    /// Creates a promise and queues a request through `enqueue()`. If the deadline
493    /// has already passed, the promise gets `Timeout` and nothing is queued.
494    pub(super) fn request(
495        self: &Arc<Self>,
496        request: Message,
497        deadline: Instant,
498    ) -> Result<Promise<Message>, Error> {
499        let (sender, promise) = Promise::pair(Arc::downgrade(self), deadline, true);
500        self.enqueue(
501            OutgoingBody::Request(request),
502            PendingOperation {
503                deadline,
504                sender,
505                log_id: None,
506            },
507        )?;
508        Ok(promise)
509    }
510
511    /// Queues a reply to `id` through `enqueue()`. Its promise waits for the write
512    /// and flush to finish, or fails if its deadline expires first.
513    pub(super) fn reply(
514        self: &Arc<Self>,
515        id: u64,
516        result: Result<Message, schema::Error>,
517        deadline: Instant,
518    ) -> Result<Promise<()>, Error> {
519        let (sender, promise) = Promise::pair(Arc::downgrade(self), deadline, false);
520        self.enqueue(
521            OutgoingBody::Reply { id, result },
522            PendingOperation {
523                deadline,
524                sender,
525                log_id: Some(LogId::from(id)),
526            },
527        )?;
528        Ok(promise)
529    }
530
531    /// Adds a `PendingOperation` and its `OutgoingMessage` under the state lock.
532    /// A closed session returns its error directly; an expired deadline fails
533    /// the promise.
534    fn enqueue(
535        self: &Arc<Self>,
536        body: OutgoingBody,
537        operation: PendingOperation,
538    ) -> Result<(), Error> {
539        {
540            let mut state = self.state.lock().expect("session state not poisoned");
541            let (operations, outgoing, reserved_ids) = match &mut *state {
542                State::Open {
543                    operations,
544                    outgoing,
545                    reserved_ids,
546                    ..
547                } => (operations, outgoing, reserved_ids),
548                State::Closed(error) => return Err(error.clone()),
549            };
550            let now = self.now();
551            if now >= operation.deadline {
552                if let OutgoingBody::Reply { id, .. } = body {
553                    reserved_ids.remove(&id);
554                }
555                operation.fail(Error::Timeout, now);
556                return Ok(());
557            }
558            let key = OperationKey::new();
559            outgoing.push_back(OutgoingMessage {
560                body,
561                operation: OperationHandle {
562                    session: Arc::downgrade(self),
563                    key: key.clone(),
564                },
565                #[cfg(any(test, feature = "fuzz"))]
566                deadline: operation.deadline,
567            });
568            operations.insert(key, operation);
569        }
570        self.changed.notify_all();
571        Ok(())
572    }
573
574    /// Fails expired operations and removes their queued messages. A promise waiter
575    /// can call this if the deadline worker has not yet processed its timeout.
576    /// Requests already sent remain in `outstanding` until answered or closed.
577    pub(super) fn expire(&self) {
578        let mut state = self.state.lock().expect("session state not poisoned");
579        state.expire(self.now());
580    }
581
582    /// Returns the earliest pending operation deadline for scenario assertions.
583    #[cfg(any(test, feature = "fuzz"))]
584    pub(super) fn next_deadline(&self) -> Option<Instant> {
585        let state = self.state.lock().expect("session state not poisoned");
586        match &*state {
587            State::Open { operations, .. } => operations
588                .values()
589                .map(|operation| operation.deadline)
590                .min(),
591            State::Closed(_) => None,
592        }
593    }
594
595    /// Takes the next unexpired `OutgoingMessage` for tests that drive writing themselves.
596    #[cfg(any(test, feature = "fuzz"))]
597    pub(super) fn take_outgoing(&self) -> Option<OutgoingMessage> {
598        let mut state = self.state.lock().expect("session state not poisoned");
599        state.expire(self.now());
600        match &mut *state {
601            State::Open {
602                outgoing,
603                reserved_ids,
604                ..
605            } => {
606                let message = outgoing.pop_front()?;
607                if let OutgoingBody::Reply { id, .. } = &message.body {
608                    reserved_ids.remove(id);
609                }
610                Some(message)
611            }
612            State::Closed(_) => None,
613        }
614    }
615
616    /// Records a write result under the state lock. A successful request write
617    /// leaves its operation waiting for an answer; a reply write completes it.
618    pub(super) fn record_write(&self, key: &OperationKey, result: Result<(), Error>) {
619        let mut state = self.state.lock().expect("session state not poisoned");
620        let State::Open { operations, .. } = &mut *state else {
621            return;
622        };
623        let Some(operation) = operations.get(key) else {
624            return;
625        };
626        let now = self.now();
627        if now >= operation.deadline || result.is_err() {
628            let operation = operations.remove(key).expect("operation held under lock");
629            operation.fail(result.err().unwrap_or(Error::Timeout), now);
630        } else if !operation.sender.response {
631            let operation = operations.remove(key).expect("operation held under lock");
632            let _ = operation.sender.send(Ok(PromiseResult::Written));
633        }
634    }
635
636    /// Supplies a response to a fixture operation without depending on wire IDs.
637    /// The fixture still encodes and retains bytes, matching real promise behavior.
638    #[cfg(any(test, feature = "fuzz"))]
639    pub(super) fn record_response(
640        self: &Arc<Self>,
641        key: &OperationKey,
642        result: Result<Message, Error>,
643    ) {
644        let mut state = self.state.lock().expect("session state not poisoned");
645        let State::Open {
646            operations,
647            max_inbound_bytes,
648            ..
649        } = &mut *state
650        else {
651            return;
652        };
653        let Some(operation) = operations.remove(key) else {
654            return;
655        };
656        let result = match result {
657            Ok(message) => Ok(message),
658            Err(Error::Remote(error)) => Err(error),
659            Err(error) => {
660                operation.fail(error, self.now());
661                return;
662            }
663        };
664        let peer = match self.side {
665            Side::Client => Side::Server,
666            Side::Server => Side::Client,
667        };
668        let bytes = peer
669            .encode(0, result)
670            .expect("fixture response belongs to the peer");
671        let bytes = Bytes::from(bytes.into_boxed_slice());
672        let header = self
673            .side
674            .decode_header(bytes.clone())
675            .expect("fixture response has a valid envelope");
676        let result = operation.complete_response(self.now(), || {
677            self.retain_incoming(bytes, header, *max_inbound_bytes)
678        });
679        drop(state);
680        if let Err(error) = result {
681            self.close(error);
682        }
683    }
684
685    /// Returns `Instant::now()` or the test clock. Deadline checks use this while
686    /// holding `state`; `reply_unanswered()` also calls it before waiting for that lock.
687    fn now(&self) -> Instant {
688        #[cfg(any(test, feature = "fuzz"))]
689        if let Some(now) = *self.time.lock().expect("scenario clock not poisoned") {
690            return now;
691        }
692        Instant::now()
693    }
694
695    /// Checks the outer envelope and routes its original bytes. `recv()` and
696    /// `wait()` decode nested payloads. Unknown and late responses are discarded
697    /// without decoding their bodies. The reader closes the session on error.
698    pub(super) fn handle_message(self: &Arc<Self>, bytes: Vec<u8>) -> Result<(), Error> {
699        let bytes = Bytes::from(bytes.into_boxed_slice());
700        let header = self.side.decode_header(bytes.clone())?;
701        {
702            let mut state = self.state.lock().expect("session state not poisoned");
703            match MessageKind::from_id(header.id, self.side.into()) {
704                MessageKind::Request => {
705                    if header.failed {
706                        return Err(self.side.malformed(
707                            Some(header),
708                            bytes.len(),
709                            "envelope",
710                            "request contains an error",
711                        ));
712                    }
713                    self.queue_request(&mut state, header, bytes)?;
714                }
715                MessageKind::Response => {
716                    let State::Open {
717                        operations,
718                        outstanding,
719                        max_inbound_bytes,
720                        ..
721                    } = &mut *state
722                    else {
723                        let State::Closed(error) = &*state else {
724                            unreachable!()
725                        };
726                        return Err(error.clone());
727                    };
728                    // A sent request keeps its ID here until answered, so an ID
729                    // without an operation belongs to a request that timed out
730                    if let Some(key) = outstanding.remove(&header.id) {
731                        if let Some(operation) = operations.remove(&key) {
732                            tracing::trace!(
733                                "received response {} ({})",
734                                header.id,
735                                header.payload.unwrap_or("none")
736                            );
737                            operation.complete_response(self.now(), || {
738                                self.retain_incoming(bytes, header, *max_inbound_bytes)
739                            })?;
740                        } else {
741                            tracing::debug!("discarding late response {}", header.id);
742                        }
743                    } else {
744                        tracing::warn!("discarding unmatched response {}", header.id);
745                    }
746                }
747            }
748        }
749        self.changed.notify_all();
750        Ok(())
751    }
752
753    /// Reserves a request slot and bytes under the session lock. If either limit
754    /// is exceeded, the queue stays untouched. Never waits for the application.
755    fn queue_request(
756        self: &Arc<Self>,
757        state: &mut State,
758        header: Header,
759        bytes: Bytes,
760    ) -> Result<(), Error> {
761        let State::Open {
762            incoming,
763            reserved_ids,
764            max_inbound_requests,
765            max_inbound_bytes,
766            ..
767        } = state
768        else {
769            let State::Closed(error) = state else {
770                unreachable!()
771            };
772            return Err(error.clone());
773        };
774        let id = header.id;
775        if reserved_ids.contains(&id) {
776            return Err(self.side.malformed(
777                Some(header),
778                bytes.len(),
779                "envelope",
780                "duplicate request ID",
781            ));
782        }
783        if reserved_ids.len() >= *max_inbound_requests {
784            tracing::warn!(
785                "inbound request limit exceeded (id: {}, used: {}, limit: {})",
786                id,
787                reserved_ids.len(),
788                max_inbound_requests
789            );
790            return Err(Error::InboundRequestLimitExceeded(*max_inbound_requests));
791        }
792        let payload = header.payload.unwrap_or("none");
793        let message = self.retain_incoming(bytes, header, *max_inbound_bytes)?;
794        reserved_ids.insert(id);
795        incoming.push_back((id, message));
796        tracing::trace!("received request {} ({})", id, payload);
797        Ok(())
798    }
799
800    /// Waits for and takes the next queued message, or returns `None` on closure.
801    /// Under the state lock, assigns each request an ID and records its operation
802    /// key in `outstanding`, so `handle_message()` can match the peer's response
803    /// even if it arrives before the write finishes.
804    pub(super) fn next_outgoing(&self) -> Option<(u64, OutgoingMessage)> {
805        let mut state = self.state.lock().expect("session state not poisoned");
806        loop {
807            // Remove expired messages before choosing the next one to send.
808            state.expire(self.now());
809            let State::Open {
810                outgoing,
811                next_id,
812                outstanding,
813                operations,
814                reserved_ids,
815                ..
816            } = &mut *state
817            else {
818                return None;
819            };
820            if let Some(outgoing) = outgoing.pop_front() {
821                let id = match &outgoing.body {
822                    OutgoingBody::Request(_) => {
823                        let id = next_id.expect("wire request IDs exhausted");
824                        *next_id = id.checked_add(2);
825                        // Store the ID before releasing the lock: a response can
826                        // arrive before the outgoing send finishes locally.
827                        outstanding.insert(id, outgoing.operation.key.clone());
828                        if let Some(operation) = operations.get_mut(&outgoing.operation.key) {
829                            operation.log_id = Some(LogId::from(id));
830                        }
831                        id
832                    }
833                    OutgoingBody::Reply { id, .. } => {
834                        // The peer may receive this reply and reuse the ID before
835                        // our flush returns. Finishing this write must not remove
836                        // a newer request that reuses the same ID.
837                        reserved_ids.remove(id);
838                        *id
839                    }
840                };
841                return Some((id, outgoing));
842            }
843            state = self
844                .changed
845                .wait(state)
846                .expect("session state not poisoned");
847        }
848    }
849
850    /// Sends queued messages through `sender`. Transport errors close the session;
851    /// messages that cannot be encoded fail only their own promise.
852    fn run_writer(&self, sender: transport::Sender<impl Write>) {
853        while let Some((id, outgoing)) = self.next_outgoing() {
854            // next_outgoing() released the state lock. The reader and deadline
855            // worker can continue while encoding or sending this message blocks.
856            let request = matches!(outgoing.body, OutgoingBody::Request(_));
857            let kind = if request { "request" } else { "reply" };
858            let body = match outgoing.body {
859                OutgoingBody::Request(body) => Ok(body),
860                OutgoingBody::Reply { result, .. } => result,
861            };
862            let payload = match &body {
863                Ok(message) => message.field_name(),
864                Err(_) => "err",
865            };
866            let result = self.side.encode(id, body).and_then(|bytes| {
867                // Announced before the transport confirms the write, so a
868                // message reads top down in the log
869                tracing::trace!("sending {} {} ({})", kind, id, payload);
870                sender.send(&bytes).map_err(Error::from)
871            });
872            match &result {
873                Ok(()) => {}
874                Err(Error::Transport(error)) => {
875                    // Wire failure ends the session even if this operation's promise
876                    // has already timed out while the transport write was blocked.
877                    self.close(Error::Transport(error.clone()));
878                    break;
879                }
880                Err(error) => tracing::debug!("not sending {} {}: {}", kind, id, error),
881            }
882            {
883                // Remaining failures are local encoding refusals. No request was
884                // sent, so there is no future answer to retain an ID for.
885                let mut state = self.state.lock().expect("session state not poisoned");
886                if let State::Open { outstanding, .. } = &mut *state
887                    && request
888                    && result.is_err()
889                {
890                    outstanding.remove(&id);
891                }
892            }
893            // Report this operation's write result, if it is still pending.
894            // A response or timeout may have completed it during the write.
895            outgoing.operation.record_write(result);
896        }
897        #[cfg(any(test, feature = "fuzz"))]
898        if let Some((entered, released)) = self.disconnect_hook.lock().unwrap().take() {
899            let _ = entered.send(());
900            let _ = released.recv();
901        }
902        // Disconnect the session this sender belongs to. The transport ignores
903        // this call if a new handshake has already replaced that session.
904        if let Err(error) = sender.disconnect() {
905            tracing::debug!(
906                "failed to signal dropped session {}: {}",
907                sender.log_id(),
908                error
909            );
910        }
911    }
912
913    /// Expires pending operations even when no caller is waiting on a promise.
914    /// Waits on `changed` until the next deadline or until new work arrives.
915    fn run_deadlines(&self) {
916        let mut state = self.state.lock().expect("session state not poisoned");
917        loop {
918            state.expire(self.now());
919            let State::Open { operations, .. } = &*state else {
920                return;
921            };
922            // Submitting an earlier deadline wakes this wait. Every wakeup
923            // recomputes the minimum under the same lock used by submission.
924            state = match operations
925                .values()
926                .map(|operation| operation.deadline)
927                .min()
928            {
929                Some(deadline) => {
930                    self.changed
931                        .wait_timeout(state, deadline.saturating_duration_since(self.now()))
932                        .expect("session state not poisoned")
933                        .0
934                }
935                None => self
936                    .changed
937                    .wait(state)
938                    .expect("session state not poisoned"),
939            };
940        }
941    }
942}
943
944impl Session {
945    /// Creates the session and starts its writer and deadline threads. Only
946    /// client sessions receive a stream closer; `Server` closes server streams.
947    pub(super) fn start<W: Write + Send + 'static>(
948        side: Side,
949        sender: transport::Sender<W>,
950        stream_closer: Option<transport::Closer>,
951        #[cfg(any(test, feature = "fuzz"))] workers: Arc<worker::Tracker>,
952    ) -> Self {
953        let session = Self {
954            inner: Arc::new(SessionInner::new(
955                side,
956                sender.log_id(),
957                stream_closer,
958                #[cfg(any(test, feature = "fuzz"))]
959                workers.clone(),
960            )),
961        };
962        let inner = session.inner.clone();
963        worker::spawn(
964            "wire-writer",
965            #[cfg(any(test, feature = "fuzz"))]
966            &workers,
967            move || inner.run_writer(sender),
968        );
969        let inner = session.inner.clone();
970        worker::spawn(
971            "wire-deadlines",
972            #[cfg(any(test, feature = "fuzz"))]
973            &workers,
974            move || inner.run_deadlines(),
975        );
976        session
977    }
978}
979
980impl State {
981    /// Counts the pending operations and the queued messages of an open session.
982    fn workload(&self) -> (usize, usize) {
983        match self {
984            Self::Open {
985                operations,
986                outgoing,
987                incoming,
988                ..
989            } => (operations.len(), outgoing.len() + incoming.len()),
990            Self::Closed(_) => (0, 0),
991        }
992    }
993
994    /// Stops accepting messages and fails pending operations under the session lock.
995    /// Returns the old queues to be dropped after releasing the lock.
996    fn close(&mut self, error: Error, now: Instant) -> Option<Self> {
997        if let Self::Closed(_) = self {
998            return None;
999        }
1000        let mut removed = std::mem::replace(self, Self::Closed(error.clone()));
1001        if let Self::Open { operations, .. } = &mut removed {
1002            for (_, operation) in operations.drain() {
1003                operation.fail(error.clone(), now);
1004            }
1005        }
1006        Some(removed)
1007    }
1008
1009    /// Removes expired entries from `operations`, sends `Timeout` to their
1010    /// promises, and discards any messages they still have in `outgoing`.
1011    fn expire(&mut self, now: Instant) {
1012        if let Self::Open {
1013            operations,
1014            outgoing,
1015            reserved_ids,
1016            ..
1017        } = self
1018        {
1019            let expired: Vec<_> = operations
1020                .iter()
1021                .filter(|(_, operation)| now >= operation.deadline)
1022                .map(|(key, _)| key.clone())
1023                .collect();
1024            for key in expired {
1025                operations
1026                    .remove(&key)
1027                    .expect("expired operation held under lock")
1028                    .fail(Error::Timeout, now);
1029            }
1030            // Only messages still in this queue can be discarded. Writes already
1031            // started keep running with their independent transport timeout.
1032            outgoing.retain(|outgoing| {
1033                let retained = operations.contains_key(&outgoing.operation.key);
1034                if !retained && let OutgoingBody::Reply { id, .. } = outgoing.body {
1035                    reserved_ids.remove(&id);
1036                }
1037                retained
1038            });
1039        }
1040    }
1041}
1042
1043// These fixtures let tests drive time, incoming requests, and write results.
1044#[cfg(any(test, feature = "fuzz"))]
1045impl Session {
1046    /// Creates a session without a stream or workers for lifecycle scenarios.
1047    pub(super) fn fixture() -> Self {
1048        Self::fixture_for(Side::Server)
1049    }
1050
1051    /// Creates either envelope direction without a stream or workers.
1052    pub(super) fn fixture_for(side: Side) -> Self {
1053        Self {
1054            inner: Arc::new(SessionInner::new(
1055                side,
1056                LogId::default(),
1057                None,
1058                Arc::new(worker::Tracker::default()),
1059            )),
1060        }
1061    }
1062}
1063
1064#[cfg(any(test, feature = "fuzz"))]
1065impl SessionInner {
1066    /// Pauses the writer before `sender.disconnect()`, so a test can connect a
1067    /// replacement session before letting the old writer finish.
1068    pub(super) fn pause_disconnect(
1069        &self,
1070    ) -> (std::sync::mpsc::Receiver<()>, std::sync::mpsc::Sender<()>) {
1071        let (entered, observed) = std::sync::mpsc::channel();
1072        let (release, released) = std::sync::mpsc::channel();
1073        *self.disconnect_hook.lock().unwrap() = Some((entered, released));
1074        (observed, release)
1075    }
1076
1077    /// Returns a receiver notified when the last `Arc<SessionInner>` is dropped.
1078    pub(super) fn watch_drop(&self) -> std::sync::mpsc::Receiver<()> {
1079        let (sender, receiver) = std::sync::mpsc::channel();
1080        *self.drop_hook.lock().unwrap() = Some(sender);
1081        receiver
1082    }
1083
1084    /// Moves a fresh scenario session to its last allocatable request ID.
1085    pub(super) fn use_last_request_id(&self) {
1086        let mut state = self.state.lock().unwrap();
1087        let State::Open {
1088            next_id,
1089            outstanding,
1090            ..
1091        } = &mut *state
1092        else {
1093            panic!("open session required")
1094        };
1095        assert!(outstanding.is_empty());
1096        *next_id = Some(if self.side == Side::Client {
1097            u64::MAX
1098        } else {
1099            u64::MAX - 1
1100        });
1101    }
1102
1103    /// Waits for the reader to remove a previously sent request's ID. Tests use
1104    /// this to leave a completed promise unread before changing limits or sessions.
1105    pub(super) fn wait_response(&self, id: u64) {
1106        let state = self.state.lock().unwrap();
1107        let (state, _) = self.changed.wait_timeout_while(state, Duration::from_secs(3), |state| {
1108            matches!(state, State::Open { outstanding, .. } if outstanding.contains_key(&id))
1109        }).unwrap();
1110        let State::Open { outstanding, .. } = &*state else {
1111            panic!("session closed before response fence");
1112        };
1113        assert!(
1114            !outstanding.contains_key(&id),
1115            "reader did not process response"
1116        );
1117    }
1118
1119    /// Returns the sorted request IDs still waiting for peer responses.
1120    pub(super) fn outstanding_ids(&self) -> Vec<u64> {
1121        let state = self.state.lock().unwrap();
1122        let State::Open { outstanding, .. } = &*state else {
1123            panic!("open session required")
1124        };
1125        let mut ids: Vec<_> = outstanding.keys().copied().collect();
1126        ids.sort_unstable();
1127        ids
1128    }
1129
1130    /// Supplies a request directly to the fixture's admission path, independently
1131    /// of parity (wire routing is covered by the connection scenarios).
1132    pub(super) fn inject_request(self: &Arc<Self>, id: u64, message: Message) -> Result<(), Error> {
1133        let peer = match self.side {
1134            Side::Client => Side::Server,
1135            Side::Server => Side::Client,
1136        };
1137        let bytes = peer
1138            .encode(id, Ok(message))
1139            .expect("fixture request belongs to peer");
1140        let bytes = Bytes::from(bytes.into_boxed_slice());
1141        let header = self.side.decode_header(bytes.clone())?;
1142        let result = {
1143            let mut state = self.state.lock().expect("session state not poisoned");
1144            self.queue_request(&mut state, header, bytes)
1145        };
1146        self.changed.notify_all();
1147        if let Err(error) = &result {
1148            self.close(error.clone());
1149        }
1150        result
1151    }
1152
1153    /// Returns the accepted request count and retained byte count for test checks.
1154    pub(super) fn inbound_usage(&self) -> (usize, usize) {
1155        let state = self.state.lock().expect("session state not poisoned");
1156        let requests = match &*state {
1157            State::Open { reserved_ids, .. } => reserved_ids.len(),
1158            State::Closed(_) => 0,
1159        };
1160        (requests, self.retained_bytes.load(Ordering::Relaxed))
1161    }
1162
1163    /// Advances the test clock without calling `expire()`, so tests can deliver
1164    /// results after a deadline but before the timeout has been processed.
1165    pub(super) fn set_time(&self, now: Instant) {
1166        let _state = self.state.lock().expect("session state not poisoned");
1167        let mut time = self.time.lock().expect("scenario clock not poisoned");
1168        assert!(
1169            time.is_none_or(|previous| now >= previous),
1170            "clock cannot go backwards"
1171        );
1172        *time = Some(now);
1173    }
1174
1175    /// Restores wall-clock time for scenarios that exercise the waiter's real timer.
1176    pub(super) fn use_realtime(&self) {
1177        let _state = self.state.lock().expect("session state not poisoned");
1178        *self.time.lock().expect("scenario clock not poisoned") = None;
1179    }
1180
1181    /// Arms a one-shot notification for the next receive waiting on an empty queue.
1182    /// The notification is sent while holding the lock, immediately before the
1183    /// condition-variable wait releases it, so a later close cannot run too early.
1184    ///
1185    /// # Panics
1186    /// The fixture must still be open and have no queued request.
1187    pub(super) fn watch_recv_wait(&self) -> std::sync::mpsc::Receiver<()> {
1188        let (sender, receiver) = std::sync::mpsc::channel();
1189        let mut state = self.state.lock().expect("session state not poisoned");
1190        let State::Open {
1191            incoming,
1192            wait_hook,
1193            ..
1194        } = &mut *state
1195        else {
1196            panic!("only watch an open session receive");
1197        };
1198        assert!(incoming.is_empty());
1199        *wait_hook = Some(sender);
1200        receiver
1201    }
1202}
1203
1204#[cfg(any(test, feature = "fuzz"))]
1205impl Drop for SessionInner {
1206    /// Notifies the test when the last `Arc<SessionInner>` is dropped.
1207    fn drop(&mut self) {
1208        if let Some(sender) = self.drop_hook.get_mut().unwrap().take() {
1209            let _ = sender.send(());
1210        }
1211    }
1212}
1213
1214/// Checks session ownership bounds and compiles the client construction API.
1215#[cfg(test)]
1216#[cfg_attr(coverage_nightly, coverage(off))]
1217mod tests {
1218    use crate::protocol::{self, Error, Session};
1219    use crate::transport::{Read, Stream, Verifier, Write};
1220    use std::fmt::Debug;
1221
1222    /// Compiles client construction with verifier-specific information in the result.
1223    #[allow(dead_code)]
1224    fn connect<R, W, V>(stream: Stream<R, W>, verifier: &V) -> Result<(Session, V::Info), Error>
1225    where
1226        R: Read + Send + 'static,
1227        W: Write + Send + 'static,
1228        V: Verifier,
1229    {
1230        protocol::connect(stream, verifier)
1231    }
1232
1233    /// Checks the bounds required to move the session to an application thread
1234    /// and to print it.
1235    #[test]
1236    fn test_thread_capabilities() {
1237        /// Requires an owned value to be printable and transferable to a background thread.
1238        fn movable<T: Debug + Send + 'static>() {}
1239        movable::<Session>();
1240    }
1241}