protocol/session/
host.rs

1//! The daemon's end of a device connection.
2
3use std::collections::{HashMap, VecDeque};
4use std::time::{Duration, Instant};
5
6use super::{
7    DEADLINE_GRACE, DisconnectReason, HandshakeError, Negotiated, PEER_TIMEOUT, PING_INTERVAL,
8    Phase, RESYNC_BACKOFF, RESYNC_BACKOFF_MAX, select_version,
9};
10use crate::capability::Capabilities;
11use crate::error::{Error, code};
12use crate::event::{Event, Stream};
13use crate::frame::{DEFAULT_MAX_FRAME, FrameDecoder, MIN_MAX_FRAME, Oversized, write_frame};
14use crate::liveness::{Alive, Liveness};
15use crate::message::{
16    Call, DecodeError, Hello, Message, Method, Ping, Pong, Reply, Request, Response, ResyncParams,
17    Welcome, kind,
18};
19use crate::sequence::{SeqCheck, StreamTracker};
20
21/// Settings for a [`HostSession`].
22#[derive(Clone, Debug, PartialEq, Eq)]
23pub struct HostConfig {
24    /// Protocol versions the daemon speaks.
25    pub versions: Vec<u16>,
26    /// The largest frame payload the daemon accepts.
27    pub max_frame: u32,
28    /// The per-launch token the agent's hello must carry, if any.
29    pub token: Option<String>,
30    pub ping_interval: Duration,
31    pub peer_timeout: Duration,
32    pub deadline_grace: Duration,
33}
34
35impl Default for HostConfig {
36    fn default() -> HostConfig {
37        HostConfig {
38            versions: crate::SUPPORTED_VERSIONS.to_vec(),
39            max_frame: DEFAULT_MAX_FRAME,
40            token: None,
41            ping_interval: PING_INTERVAL,
42            peer_timeout: PEER_TIMEOUT,
43            deadline_grace: DEADLINE_GRACE,
44        }
45    }
46}
47
48/// What happened on a [`HostSession`].
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub enum HostEvent {
51    /// The handshake completed.
52    Ready(Negotiated),
53    /// The answer to a request: the agent's, or an error the session reports
54    /// on its behalf (a passed deadline, an oversized frame, a lost connection).
55    Response {
56        id: u32,
57        method: Method,
58        result: Result<Reply, Error>,
59    },
60    /// An event, in order.
61    Event(Event),
62    /// A lossless stream skipped or repeated a sequence number; its events
63    /// are dropped until the resync the session requested completes.
64    OutOfSync {
65        stream: Stream,
66        expected: u64,
67        received: u64,
68    },
69    /// The stream is in sync again: `seq` is the last number the agent sent
70    /// on it before its state was reset.
71    Resynced { stream: Stream, seq: u64 },
72    /// The agent could not resynchronize the stream, which stays out of
73    /// sync: a resync asked for on demand failed, or an automatic one failed
74    /// in a way no retry fixes. An automatic resync that fails in a way a
75    /// retry can fix, such as a busy agent, is sent again after a backoff
76    /// instead, until the stream is in sync.
77    ResyncFailed { stream: Stream, error: Error },
78    /// A frame over the limit that answered no pending request.
79    FrameTooLarge(Oversized),
80    /// The session closed; later requests fail at once.
81    Closed(DisconnectReason),
82}
83
84#[derive(Debug)]
85enum State {
86    AwaitingHello,
87    Ready(Negotiated),
88    Closed(DisconnectReason),
89}
90
91/// A request sent and not answered yet.
92#[derive(Debug)]
93struct Pending {
94    method: Method,
95    deadline_ms: Option<u32>,
96    /// When the host stops waiting for the agent's answer: the deadline plus
97    /// the grace period.
98    expires: Option<Instant>,
99    /// Set for resync requests; their outcome is reported as a resync, not
100    /// as a response.
101    resync: Option<ResyncOf>,
102}
103
104/// A resync request's stream, and whether the session sent it on its own
105/// after a sequence gap, in which case it is retried.
106#[derive(Clone, Copy, Debug)]
107struct ResyncOf {
108    stream: Stream,
109    automatic: bool,
110}
111
112/// An automatic resync that failed and is to be sent again.
113#[derive(Clone, Copy, Debug)]
114struct Retry {
115    /// Tries that failed in a row.
116    failures: u32,
117    /// When it is due; `None` until the session next learns the time.
118    at: Option<Instant>,
119}
120
121/// The daemon's end of a device connection.
122#[derive(Debug)]
123pub struct HostSession {
124    config: HostConfig,
125    state: State,
126    decoder: FrameDecoder,
127    outbox: Vec<u8>,
128    events: VecDeque<HostEvent>,
129    pending: HashMap<u32, Pending>,
130    next_id: u32,
131    liveness: Liveness,
132    next_nonce: u64,
133    streams: StreamTracker,
134    /// Failed automatic resyncs to send again, by stream.
135    retries: HashMap<Stream, Retry>,
136    /// Failures in a row of the automatic resync in flight, by stream, so a
137    /// retry that fails again backs off longer.
138    in_flight: HashMap<Stream, u32>,
139}
140
141impl HostSession {
142    /// Starts waiting for the agent's hello; it must arrive within the peer
143    /// timeout.
144    pub fn new(config: HostConfig, now: Instant) -> HostSession {
145        HostSession {
146            decoder: FrameDecoder::new(config.max_frame.max(MIN_MAX_FRAME)),
147            liveness: Liveness::new(config.ping_interval, config.peer_timeout, now),
148            config,
149            state: State::AwaitingHello,
150            outbox: Vec::new(),
151            events: VecDeque::new(),
152            pending: HashMap::new(),
153            next_id: 1,
154            next_nonce: 0,
155            streams: StreamTracker::new(),
156            retries: HashMap::new(),
157            in_flight: HashMap::new(),
158        }
159    }
160
161    /// Feeds bytes read from the connection.
162    pub fn receive(&mut self, bytes: &[u8], now: Instant) {
163        if self.is_closed() {
164            return;
165        }
166        self.decoder.push(bytes);
167        // Bytes that only feed the discard of an oversized frame are no sign
168        // of life: a corrupt length must not keep a dead peer alive.
169        let mut alive = !self.decoder.is_discarding();
170        while !self.is_closed() {
171            let frame = match self.decoder.next_frame() {
172                None => break,
173                Some(Ok(payload)) => Ok(Message::decode(payload)),
174                Some(Err(oversized)) => Err(oversized),
175            };
176            alive = true;
177            match frame {
178                Ok(Ok(msg)) => self.handle(msg, now),
179                // A kind a newer agent may send; nothing to do with it.
180                Ok(Err(DecodeError::UnknownKind(_))) => {}
181                Ok(Err(e)) => self.close_with(DisconnectReason::Protocol(e.to_string())),
182                Err(oversized) => self.handle_oversized(oversized),
183            }
184        }
185        if alive {
186            self.liveness.heard(now);
187        }
188        self.schedule_retries(now);
189    }
190
191    /// The connection reached its end.
192    pub fn receive_eof(&mut self, now: Instant) {
193        let _ = now;
194        if !self.is_closed() {
195            self.close_with(DisconnectReason::Eof);
196        }
197    }
198
199    /// Acts on deadlines and liveness that are due at `now`.
200    pub fn handle_timeout(&mut self, now: Instant) {
201        if self.is_closed() {
202            return;
203        }
204        let mut expired: Vec<u32> = self
205            .pending
206            .iter()
207            .filter(|(_, p)| p.expires.is_some_and(|at| at <= now))
208            .map(|(id, _)| *id)
209            .collect();
210        expired.sort_unstable();
211        for id in expired {
212            let pending = self.pending.remove(&id).expect("listed above");
213            let error =
214                Error::deadline_exceeded(pending.method, id, pending.deadline_ms.unwrap_or(0));
215            self.resolve(id, pending, Err(error));
216        }
217        self.schedule_retries(now);
218        self.send_due_retries(now);
219        match self.liveness.check(now, self.phase() == Phase::Ready) {
220            Alive::Yes => {}
221            Alive::Ping => {
222                let nonce = self.next_nonce;
223                self.next_nonce = self.next_nonce.wrapping_add(1);
224                self.send_control(&Message::Ping(Ping { nonce }));
225            }
226            Alive::Dead(silence) => self.close_with(DisconnectReason::Silent(silence)),
227        }
228    }
229
230    /// When `handle_timeout` must next run; `None` once closed.
231    pub fn poll_timeout(&self) -> Option<Instant> {
232        if self.is_closed() {
233            return None;
234        }
235        let liveness = self.liveness.next_check(self.phase() == Phase::Ready);
236        let deadline = self.pending.values().filter_map(|p| p.expires).min();
237        let retry = self.retries.values().filter_map(|r| r.at).min();
238        Some(
239            [deadline, retry]
240                .into_iter()
241                .flatten()
242                .fold(liveness, Instant::min),
243        )
244    }
245
246    /// Sends a request, returning its id. Fails at once, sending nothing, when
247    /// the session is not ready or closed, when the agent did not advertise
248    /// the capability the method needs, or when the request exceeds the
249    /// negotiated frame size.
250    pub fn send_request(
251        &mut self,
252        call: Call,
253        deadline: Option<Duration>,
254        now: Instant,
255    ) -> Result<u32, Error> {
256        self.send(call, deadline, now, None)
257    }
258
259    /// Asks the agent to resynchronize `stream`; the outcome arrives as
260    /// [`HostEvent::Resynced`] or [`HostEvent::ResyncFailed`].
261    pub fn resync(&mut self, stream: Stream, now: Instant) -> Result<u32, Error> {
262        self.send(
263            Call::Resync(ResyncParams { stream }),
264            None,
265            now,
266            Some(ResyncOf {
267                stream,
268                automatic: false,
269            }),
270        )
271    }
272
273    /// Sends the resync the session asks for on its own after a gap on
274    /// `stream`, the `failures`-th try; a failure to send it is reported at
275    /// once, since no retry sends it either.
276    fn resync_automatically(&mut self, stream: Stream, failures: u32, now: Instant) {
277        let sent = self.send(
278            Call::Resync(ResyncParams { stream }),
279            None,
280            now,
281            Some(ResyncOf {
282                stream,
283                automatic: true,
284            }),
285        );
286        match sent {
287            Ok(_) => {
288                self.in_flight.insert(stream, failures);
289            }
290            Err(error) => self
291                .events
292                .push_back(HostEvent::ResyncFailed { stream, error }),
293        }
294    }
295
296    /// Gives each failed automatic resync without a time the time it is
297    /// due: the backoff after its latest failure.
298    fn schedule_retries(&mut self, now: Instant) {
299        for retry in self.retries.values_mut() {
300            if retry.at.is_none() {
301                let doublings = retry.failures.saturating_sub(1).min(16);
302                let wait = RESYNC_BACKOFF
303                    .saturating_mul(1 << doublings)
304                    .min(RESYNC_BACKOFF_MAX);
305                retry.at = Some(now + wait);
306            }
307        }
308    }
309
310    /// Sends the automatic resyncs due at `now` again.
311    fn send_due_retries(&mut self, now: Instant) {
312        if self.phase() != Phase::Ready {
313            return;
314        }
315        let mut due: Vec<(Stream, u32)> = self
316            .retries
317            .iter()
318            .filter(|(_, retry)| retry.at.is_some_and(|at| at <= now))
319            .map(|(stream, retry)| (*stream, retry.failures))
320            .collect();
321        due.sort_unstable_by_key(|(stream, _)| stream.0);
322        for (stream, failures) in due {
323            self.retries.remove(&stream);
324            self.resync_automatically(stream, failures, now);
325        }
326    }
327
328    /// Sequence state of the event streams.
329    pub fn streams(&self) -> &StreamTracker {
330        &self.streams
331    }
332
333    fn send(
334        &mut self,
335        call: Call,
336        deadline: Option<Duration>,
337        now: Instant,
338        resync: Option<ResyncOf>,
339    ) -> Result<u32, Error> {
340        let (capabilities, limit) = match &self.state {
341            State::Ready(n) => (n.capabilities, n.max_frame),
342            State::AwaitingHello => {
343                return Err(Error::new(
344                    code::NOT_READY,
345                    "the agent has not finished its handshake",
346                ));
347            }
348            State::Closed(reason) => return Err(Error::disconnected(reason)),
349        };
350        let method = call.method();
351        let needed = method.required_capability();
352        if !capabilities.contains(needed) {
353            let missing = Capabilities(needed.0 & !capabilities.0);
354            return Err(Error::capability_missing(method, missing));
355        }
356        let id = self.allocate_id();
357        let deadline_ms = deadline.map(|d| u32::try_from(d.as_millis()).unwrap_or(u32::MAX));
358        let request = Message::Request(Request {
359            id,
360            deadline_ms,
361            call,
362        });
363        write_frame(&request, limit, &mut self.outbox)
364            .map_err(|e| Error::frame_too_large(method, id, e.size as u64, e.limit))?;
365        let expires = deadline.map(|d| now + d + self.config.deadline_grace);
366        self.pending.insert(
367            id,
368            Pending {
369                method,
370                deadline_ms,
371                expires,
372                resync,
373            },
374        );
375        Ok(id)
376    }
377
378    /// Bytes to write to the connection.
379    pub fn poll_transmit(&mut self) -> Option<Vec<u8>> {
380        if self.outbox.is_empty() {
381            None
382        } else {
383            Some(std::mem::take(&mut self.outbox))
384        }
385    }
386
387    pub fn poll_event(&mut self) -> Option<HostEvent> {
388        self.events.pop_front()
389    }
390
391    /// What the handshake settled, once it has.
392    pub fn negotiated(&self) -> Option<&Negotiated> {
393        match &self.state {
394            State::Ready(n) => Some(n),
395            _ => None,
396        }
397    }
398
399    pub fn is_closed(&self) -> bool {
400        matches!(self.state, State::Closed(_))
401    }
402
403    /// Why the session closed, once it has.
404    pub fn close_reason(&self) -> Option<&DisconnectReason> {
405        match &self.state {
406            State::Closed(reason) => Some(reason),
407            _ => None,
408        }
409    }
410
411    /// Closes the session from this side, failing every pending request.
412    pub fn close(&mut self) {
413        if !self.is_closed() {
414            self.close_with(DisconnectReason::Closed);
415        }
416    }
417
418    fn phase(&self) -> Phase {
419        match self.state {
420            State::AwaitingHello => Phase::Handshake,
421            State::Ready(_) => Phase::Ready,
422            State::Closed(_) => Phase::Closed,
423        }
424    }
425
426    fn handle(&mut self, msg: Message, now: Instant) {
427        match (self.phase(), msg) {
428            (Phase::Handshake, Message::Hello(hello)) => self.accept_hello(hello),
429            (Phase::Ready, Message::Response(response)) => self.accept_response(response),
430            (Phase::Ready, Message::Event(event)) => self.accept_event(event, now),
431            (Phase::Ready, Message::Ping(ping)) => {
432                self.send_control(&Message::Pong(Pong { nonce: ping.nonce }));
433            }
434            (Phase::Ready, Message::Pong(_)) | (Phase::Closed, _) => {}
435            (phase, other) => {
436                let expected = if phase == Phase::Handshake {
437                    "a hello"
438                } else {
439                    "a response, an event, a ping or a pong"
440                };
441                self.close_with(DisconnectReason::Protocol(format!(
442                    "expected {expected} from the agent, received a {}",
443                    kind::name(other.kind())
444                )));
445            }
446        }
447    }
448
449    fn accept_hello(&mut self, hello: Hello) {
450        match self.negotiate(&hello) {
451            Ok(negotiated) => {
452                self.send_control(&Message::Welcome(Welcome {
453                    versions: self.config.versions.clone(),
454                    version: Some(negotiated.version),
455                    max_frame: negotiated.max_frame,
456                    error: None,
457                }));
458                self.decoder.set_limit(negotiated.max_frame);
459                self.state = State::Ready(negotiated.clone());
460                self.events.push_back(HostEvent::Ready(negotiated));
461            }
462            Err(failure) => {
463                // The agent learns why before the connection closes.
464                self.send_control(&Message::Welcome(Welcome {
465                    versions: self.config.versions.clone(),
466                    version: None,
467                    max_frame: self.config.max_frame.min(hello.max_frame),
468                    error: Some(failure.to_error()),
469                }));
470                self.close_with(DisconnectReason::Handshake(failure));
471            }
472        }
473    }
474
475    fn negotiate(&self, hello: &Hello) -> Result<Negotiated, HandshakeError> {
476        if let Some(token) = &self.config.token
477            && hello.token.as_ref() != Some(token)
478        {
479            return Err(HandshakeError::TokenMismatch);
480        }
481        let version = select_version(&self.config.versions, &hello.versions).ok_or_else(|| {
482            HandshakeError::NoCommonVersion {
483                host: self.config.versions.clone(),
484                agent: hello.versions.clone(),
485            }
486        })?;
487        let max_frame = self.config.max_frame.min(hello.max_frame);
488        if max_frame < MIN_MAX_FRAME {
489            return Err(HandshakeError::FrameLimitTooSmall { limit: max_frame });
490        }
491        Ok(Negotiated {
492            version,
493            max_frame,
494            capabilities: hello.capabilities,
495            agent: hello.agent.clone(),
496            host_versions: self.config.versions.clone(),
497            agent_versions: hello.versions.clone(),
498        })
499    }
500
501    fn accept_response(&mut self, response: Response) {
502        // Absent: already answered by its deadline, or never sent. Dropped.
503        let Some(pending) = self.pending.remove(&response.id) else {
504            return;
505        };
506        let result = if response.method == pending.method {
507            response.result
508        } else {
509            Err(Error::new(
510                code::PROTOCOL_ERROR,
511                format!(
512                    "the agent answered {} request {} as {}",
513                    pending.method, response.id, response.method
514                ),
515            ))
516        };
517        self.resolve(response.id, pending, result);
518    }
519
520    /// Reports the outcome of a request: as a response, or for resync
521    /// requests as the stream's new state. An automatic resync that failed
522    /// in a way a retry can fix is queued to be sent again instead, while
523    /// the session lives.
524    fn resolve(&mut self, id: u32, pending: Pending, result: Result<Reply, Error>) {
525        let Some(ResyncOf { stream, automatic }) = pending.resync else {
526            self.events.push_back(HostEvent::Response {
527                id,
528                method: pending.method,
529                result,
530            });
531            return;
532        };
533        let failures = if automatic {
534            self.in_flight.remove(&stream).unwrap_or(0)
535        } else {
536            0
537        };
538        match result {
539            Ok(Reply::Resync(ack)) if ack.stream == stream => {
540                self.streams.resynced(stream, ack.seq);
541                self.events.push_back(HostEvent::Resynced {
542                    stream,
543                    seq: ack.seq,
544                });
545            }
546            Ok(other) => self.events.push_back(HostEvent::ResyncFailed {
547                stream,
548                error: Error::new(
549                    code::PROTOCOL_ERROR,
550                    format!("the resync of the {stream} stream was answered with {other:?}"),
551                ),
552            }),
553            Err(error) if automatic && error.retryable && !self.is_closed() => {
554                self.retries.insert(
555                    stream,
556                    Retry {
557                        failures: failures + 1,
558                        at: None,
559                    },
560                );
561            }
562            Err(error) => self
563                .events
564                .push_back(HostEvent::ResyncFailed { stream, error }),
565        }
566    }
567
568    fn accept_event(&mut self, event: Event, now: Instant) {
569        let stream = event.stream();
570        match self.streams.check(stream, event.seq) {
571            SeqCheck::Deliver { .. } => self.events.push_back(HostEvent::Event(event)),
572            SeqCheck::Gap { expected, received } => {
573                self.events.push_back(HostEvent::OutOfSync {
574                    stream,
575                    expected,
576                    received,
577                });
578                self.retries.remove(&stream);
579                self.resync_automatically(stream, 0, now);
580            }
581            SeqCheck::Discard => {}
582        }
583    }
584
585    fn handle_oversized(&mut self, oversized: Oversized) {
586        if oversized.kind == Some(kind::RESPONSE)
587            && let Some(id) = oversized.id
588            && let Some(pending) = self.pending.remove(&id)
589        {
590            let error = Error::frame_too_large(
591                pending.method,
592                id,
593                u64::from(oversized.size),
594                oversized.limit,
595            );
596            self.resolve(id, pending, Err(error));
597            return;
598        }
599        self.events.push_back(HostEvent::FrameTooLarge(oversized));
600    }
601
602    /// Queues a hello, welcome, ping or pong: messages far below any
603    /// negotiable frame limit.
604    fn send_control(&mut self, msg: &Message) {
605        write_frame(msg, u32::MAX, &mut self.outbox).expect("a control message fits in a frame");
606    }
607
608    fn allocate_id(&mut self) -> u32 {
609        loop {
610            let id = self.next_id;
611            self.next_id = self.next_id.wrapping_add(1);
612            if id != 0 && !self.pending.contains_key(&id) {
613                return id;
614            }
615        }
616    }
617
618    fn close_with(&mut self, reason: DisconnectReason) {
619        self.retries.clear();
620        let mut ids: Vec<u32> = self.pending.keys().copied().collect();
621        ids.sort_unstable();
622        for id in ids {
623            let pending = self.pending.remove(&id).expect("listed above");
624            self.resolve(id, pending, Err(Error::disconnected(&reason)));
625        }
626        self.state = State::Closed(reason.clone());
627        self.events.push_back(HostEvent::Closed(reason));
628    }
629}
protocol/session/host.rs

protocol/session/
host.rs
