protocol/session/
agent.rs

1//! The device agent's end of the connection.
2
3use std::collections::{HashMap, VecDeque};
4use std::time::{Duration, Instant};
5
6use super::{DisconnectReason, HandshakeError, Negotiated, PEER_TIMEOUT, PING_INTERVAL, Phase};
7use crate::capability::Capabilities;
8use crate::error::Error;
9use crate::event::{Event, EventPayload, Stream};
10use crate::frame::{DEFAULT_MAX_FRAME, FrameDecoder, MIN_MAX_FRAME, Oversized, write_frame};
11use crate::liveness::{Alive, Liveness};
12use crate::message::{
13    Call, DecodeError, Hello, Message, Method, Ping, Pong, Reply, Request, Response, Welcome, kind,
14};
15
16/// Settings for an [`AgentSession`].
17#[derive(Clone, Debug, PartialEq, Eq)]
18pub struct AgentConfig {
19    /// Protocol versions the agent speaks.
20    pub versions: Vec<u16>,
21    pub capabilities: Capabilities,
22    /// The largest frame payload the agent accepts.
23    pub max_frame: u32,
24    /// Agent name and version for the hello.
25    pub agent: String,
26    /// The per-launch token the host started the agent with.
27    pub token: Option<String>,
28    pub ping_interval: Duration,
29    pub peer_timeout: Duration,
30}
31
32impl Default for AgentConfig {
33    fn default() -> AgentConfig {
34        AgentConfig {
35            versions: crate::SUPPORTED_VERSIONS.to_vec(),
36            capabilities: Capabilities::NONE,
37            max_frame: DEFAULT_MAX_FRAME,
38            agent: String::new(),
39            token: None,
40            ping_interval: PING_INTERVAL,
41            peer_timeout: PEER_TIMEOUT,
42        }
43    }
44}
45
46/// What happened on an [`AgentSession`].
47#[derive(Clone, Debug, PartialEq, Eq)]
48pub enum AgentEvent {
49    /// The host accepted the hello.
50    Ready(Negotiated),
51    /// A request to work on; answer it with [`AgentSession::respond`].
52    Request(Request),
53    /// The session closed.
54    Closed(DisconnectReason),
55}
56
57/// What became of an answer passed to [`AgentSession::respond`].
58#[derive(Clone, Copy, Debug, PartialEq, Eq)]
59pub enum Responded {
60    /// Sent as given.
61    Sent,
62    /// Over the negotiated frame size: a `frame_too_large` error went instead.
63    TooLarge,
64    /// Not sent: the request was already answered, by its deadline for
65    /// instance, or never arrived.
66    NotInFlight,
67}
68
69#[derive(Debug)]
70enum State {
71    AwaitingWelcome,
72    Ready(Negotiated),
73    Closed(DisconnectReason),
74}
75
76/// A request handed to the agent and not answered yet.
77#[derive(Debug)]
78struct InFlight {
79    method: Method,
80    deadline_ms: Option<u32>,
81    deadline: Option<Instant>,
82}
83
84/// The device agent's end of the connection.
85#[derive(Debug)]
86pub struct AgentSession {
87    config: AgentConfig,
88    state: State,
89    decoder: FrameDecoder,
90    outbox: Vec<u8>,
91    events: VecDeque<AgentEvent>,
92    in_flight: HashMap<u32, InFlight>,
93    liveness: Liveness,
94    next_nonce: u64,
95    /// The last sequence number used on each stream.
96    seqs: HashMap<Stream, u64>,
97}
98
99impl AgentSession {
100    /// Starts the session; the hello is ready to send at once, and the
101    /// welcome must arrive within the peer timeout.
102    pub fn new(config: AgentConfig, now: Instant) -> AgentSession {
103        let hello = Message::Hello(Hello {
104            versions: config.versions.clone(),
105            capabilities: config.capabilities,
106            max_frame: config.max_frame,
107            agent: config.agent.clone(),
108            token: config.token.clone(),
109        });
110        let mut session = AgentSession {
111            // Wide enough for a refusal even when the agent asks for less.
112            decoder: FrameDecoder::new(config.max_frame.max(MIN_MAX_FRAME)),
113            liveness: Liveness::new(config.ping_interval, config.peer_timeout, now),
114            config,
115            state: State::AwaitingWelcome,
116            outbox: Vec::new(),
117            events: VecDeque::new(),
118            in_flight: HashMap::new(),
119            next_nonce: 0,
120            seqs: HashMap::new(),
121        };
122        session.send_control(&hello);
123        session
124    }
125
126    /// Feeds bytes read from the connection.
127    pub fn receive(&mut self, bytes: &[u8], now: Instant) {
128        if self.is_closed() {
129            return;
130        }
131        self.decoder.push(bytes);
132        // Bytes that only feed the discard of an oversized frame are no sign of life.
133        let mut alive = !self.decoder.is_discarding();
134        while !self.is_closed() {
135            let frame = match self.decoder.next_frame() {
136                None => break,
137                Some(Ok(payload)) => Ok(Message::decode(payload)),
138                Some(Err(oversized)) => Err(oversized),
139            };
140            alive = true;
141            match frame {
142                Ok(Ok(msg)) => self.handle(msg, now),
143                Ok(Err(DecodeError::UnknownKind(_))) => {}
144                Ok(Err(e)) => self.close_with(DisconnectReason::Protocol(e.to_string())),
145                Err(oversized) => self.handle_oversized(oversized),
146            }
147        }
148        if alive {
149            self.liveness.heard(now);
150        }
151    }
152
153    /// The connection reached its end.
154    pub fn receive_eof(&mut self, now: Instant) {
155        let _ = now;
156        if !self.is_closed() {
157            self.close_with(DisconnectReason::Eof);
158        }
159    }
160
161    /// Answers requests whose deadline passed and keeps the connection alive.
162    pub fn handle_timeout(&mut self, now: Instant) {
163        match self.phase() {
164            Phase::Closed => return,
165            Phase::Handshake => {
166                if let Alive::Dead(silence) = self.liveness.check(now, false) {
167                    self.close_with(DisconnectReason::Silent(silence));
168                }
169                return;
170            }
171            Phase::Ready => {}
172        }
173        let mut expired: Vec<u32> = self
174            .in_flight
175            .iter()
176            .filter(|(_, f)| f.deadline.is_some_and(|at| at <= now))
177            .map(|(id, _)| *id)
178            .collect();
179        expired.sort_unstable();
180        for id in expired {
181            let flight = self.in_flight.remove(&id).expect("listed above");
182            let error =
183                Error::deadline_exceeded(flight.method, id, flight.deadline_ms.unwrap_or(0));
184            self.send_response(id, flight.method, Err(error));
185        }
186        match self.liveness.check(now, true) {
187            Alive::Yes => {}
188            Alive::Ping => {
189                let nonce = self.next_nonce;
190                self.next_nonce = self.next_nonce.wrapping_add(1);
191                self.send_control(&Message::Ping(Ping { nonce }));
192            }
193            Alive::Dead(silence) => self.close_with(DisconnectReason::Silent(silence)),
194        }
195    }
196
197    /// When `handle_timeout` must next run; `None` once closed.
198    pub fn poll_timeout(&self) -> Option<Instant> {
199        if self.is_closed() {
200            return None;
201        }
202        let liveness = self.liveness.next_check(self.phase() == Phase::Ready);
203        let deadline = self.in_flight.values().filter_map(|f| f.deadline).min();
204        Some(deadline.map_or(liveness, |d| d.min(liveness)))
205    }
206
207    /// The instant request `id` must be answered by, if it has a deadline.
208    pub fn deadline(&self, id: u32) -> Option<Instant> {
209        self.in_flight.get(&id).and_then(|f| f.deadline)
210    }
211
212    /// Answers request `id`.
213    pub fn respond(&mut self, id: u32, result: Result<Reply, Error>) -> Responded {
214        match self.in_flight.remove(&id) {
215            Some(flight) => self.send_response(id, flight.method, result),
216            None => Responded::NotInFlight,
217        }
218    }
219
220    /// Sends an event on its payload's stream with the stream's next sequence
221    /// number and returns the number. `None` when the session is not ready,
222    /// or when the event exceeds the negotiated frame size; its number is
223    /// used up then, so the host sees the gap and asks for a resync.
224    pub fn emit(&mut self, payload: EventPayload, time_us: u64) -> Option<u64> {
225        let State::Ready(negotiated) = &self.state else {
226            return None;
227        };
228        let limit = negotiated.max_frame;
229        let seq = self.next_seq(payload.stream());
230        let event = Message::Event(Event {
231            seq,
232            time_us,
233            payload,
234        });
235        write_frame(&event, limit, &mut self.outbox)
236            .ok()
237            .map(|()| seq)
238    }
239
240    /// Uses up the next sequence number of `stream` without sending an event,
241    /// as when a full event queue drops one. The host sees the gap: a
242    /// lossless stream then asks for a resync, a lossy one counts a drop.
243    pub fn skip_event(&mut self, stream: Stream) -> Option<u64> {
244        if self.phase() != Phase::Ready {
245            return None;
246        }
247        Some(self.next_seq(stream))
248    }
249
250    /// The last sequence number sent or skipped on `stream`; 0 before the first.
251    pub fn last_seq(&self, stream: Stream) -> u64 {
252        self.seqs.get(&stream).copied().unwrap_or(0)
253    }
254
255    fn next_seq(&mut self, stream: Stream) -> u64 {
256        let seq = self.seqs.entry(stream).or_insert(0);
257        *seq += 1;
258        *seq
259    }
260
261    /// Bytes to write to the connection.
262    pub fn poll_transmit(&mut self) -> Option<Vec<u8>> {
263        if self.outbox.is_empty() {
264            None
265        } else {
266            Some(std::mem::take(&mut self.outbox))
267        }
268    }
269
270    pub fn poll_event(&mut self) -> Option<AgentEvent> {
271        self.events.pop_front()
272    }
273
274    pub fn is_closed(&self) -> bool {
275        matches!(self.state, State::Closed(_))
276    }
277
278    /// Why the session closed, once it has.
279    pub fn close_reason(&self) -> Option<&DisconnectReason> {
280        match &self.state {
281            State::Closed(reason) => Some(reason),
282            _ => None,
283        }
284    }
285
286    fn phase(&self) -> Phase {
287        match self.state {
288            State::AwaitingWelcome => Phase::Handshake,
289            State::Ready(_) => Phase::Ready,
290            State::Closed(_) => Phase::Closed,
291        }
292    }
293
294    fn handle(&mut self, msg: Message, now: Instant) {
295        match (self.phase(), msg) {
296            (Phase::Handshake, Message::Welcome(welcome)) => self.accept_welcome(welcome),
297            (Phase::Ready, Message::Request(request)) => self.accept_request(request, now),
298            (Phase::Ready, Message::Ping(ping)) => {
299                self.send_control(&Message::Pong(Pong { nonce: ping.nonce }));
300            }
301            (Phase::Ready, Message::Pong(_)) | (Phase::Closed, _) => {}
302            (phase, other) => {
303                let expected = if phase == Phase::Handshake {
304                    "a welcome"
305                } else {
306                    "a request, a ping or a pong"
307                };
308                self.close_with(DisconnectReason::Protocol(format!(
309                    "expected {expected} from the host, received a {}",
310                    kind::name(other.kind())
311                )));
312            }
313        }
314    }
315
316    fn accept_welcome(&mut self, welcome: Welcome) {
317        if let Some(error) = welcome.error {
318            self.close_with(DisconnectReason::Handshake(HandshakeError::Refused(error)));
319            return;
320        }
321        let Some(version) = welcome.version else {
322            self.close_with(DisconnectReason::Protocol(
323                "a welcome with neither a version nor an error".to_owned(),
324            ));
325            return;
326        };
327        if !self.config.versions.contains(&version) {
328            self.close_with(DisconnectReason::Handshake(
329                HandshakeError::VersionNotOffered {
330                    version,
331                    offered: self.config.versions.clone(),
332                },
333            ));
334            return;
335        }
336        let max_frame = welcome.max_frame.min(self.config.max_frame);
337        if max_frame < MIN_MAX_FRAME {
338            self.close_with(DisconnectReason::Handshake(
339                HandshakeError::FrameLimitTooSmall { limit: max_frame },
340            ));
341            return;
342        }
343        let negotiated = Negotiated {
344            version,
345            max_frame,
346            capabilities: self.config.capabilities,
347            agent: self.config.agent.clone(),
348            host_versions: welcome.versions,
349            agent_versions: self.config.versions.clone(),
350        };
351        self.decoder.set_limit(max_frame);
352        self.state = State::Ready(negotiated.clone());
353        self.events.push_back(AgentEvent::Ready(negotiated));
354    }
355
356    fn accept_request(&mut self, request: Request, now: Instant) {
357        if self.in_flight.contains_key(&request.id) {
358            self.close_with(DisconnectReason::Protocol(format!(
359                "request id {} reused while in flight",
360                request.id
361            )));
362            return;
363        }
364        let method = request.call.method();
365        if let Call::Unknown(_) = request.call {
366            self.send_response(request.id, method, Err(Error::unknown_method(method)));
367            return;
368        }
369        let needed = method.required_capability();
370        if !self.config.capabilities.contains(needed) {
371            let missing = Capabilities(needed.0 & !self.config.capabilities.0);
372            self.send_response(
373                request.id,
374                method,
375                Err(Error::capability_missing(method, missing)),
376            );
377            return;
378        }
379        let deadline = request
380            .deadline_ms
381            .map(|ms| now + Duration::from_millis(u64::from(ms)));
382        self.in_flight.insert(
383            request.id,
384            InFlight {
385                method,
386                deadline_ms: request.deadline_ms,
387                deadline,
388            },
389        );
390        self.events.push_back(AgentEvent::Request(request));
391    }
392
393    /// Answers a request the host sent over the limit, whose payload was dropped.
394    fn handle_oversized(&mut self, oversized: Oversized) {
395        if oversized.kind == Some(kind::REQUEST)
396            && let (Some(id), Some(method)) = (oversized.id, oversized.method)
397        {
398            let error =
399                Error::frame_too_large(method, id, u64::from(oversized.size), oversized.limit);
400            self.send_response(id, method, Err(error));
401        }
402    }
403
404    /// Writes a response, or a `frame_too_large` error in its place when it
405    /// exceeds the negotiated limit.
406    fn send_response(
407        &mut self,
408        id: u32,
409        method: Method,
410        result: Result<Reply, Error>,
411    ) -> Responded {
412        let State::Ready(negotiated) = &self.state else {
413            return Responded::NotInFlight;
414        };
415        let limit = negotiated.max_frame;
416        let response = Message::Response(Response { id, method, result });
417        match write_frame(&response, limit, &mut self.outbox) {
418            Ok(()) => Responded::Sent,
419            Err(too_large) => {
420                let refusal = Message::Response(Response {
421                    id,
422                    method,
423                    result: Err(Error::frame_too_large(
424                        method,
425                        id,
426                        too_large.size as u64,
427                        limit,
428                    )),
429                });
430                write_frame(&refusal, limit, &mut self.outbox)
431                    .expect("an error response fits in the smallest frame limit");
432                Responded::TooLarge
433            }
434        }
435    }
436
437    /// Queues a hello, ping or pong: messages far below any negotiable limit.
438    fn send_control(&mut self, msg: &Message) {
439        write_frame(msg, u32::MAX, &mut self.outbox).expect("a control message fits in a frame");
440    }
441
442    fn close_with(&mut self, reason: DisconnectReason) {
443        self.in_flight.clear();
444        self.state = State::Closed(reason.clone());
445        self.events.push_back(AgentEvent::Closed(reason));
446    }
447}
protocol/session/agent.rs

protocol/session/
agent.rs
