Skip to main content

microsandbox_agent_client/
protocol.rs

1//! Agent relay setup and metadata over the shared framed router.
2
3use std::sync::Arc;
4
5use microsandbox_protocol::{
6    codec::{self, MAX_FRAME_SIZE, RawFrame},
7    core::Ready,
8    message::{FRAME_HEADER_SIZE, MessageType, PROTOCOL_VERSION},
9};
10use microsandbox_protocol_client::{
11    BoxFuture, BoxTransport, CborEnvelopeCodec, ClientError, ClientResult, ConnectOptions,
12    ErrorKind, Established, IdRange, Protocol, SendMetadata,
13};
14use tokio::io::AsyncReadExt;
15
16//--------------------------------------------------------------------------------------------------
17// Constants
18//--------------------------------------------------------------------------------------------------
19
20/// Protocol generation used by the supported pre-0.5 agent wire path.
21pub const LEGACY_PROTOCOL_VERSION: u8 = 1;
22const LEGACY_RELAY_ID_RANGE_STEP: u32 = u32::MAX / 16;
23
24//--------------------------------------------------------------------------------------------------
25// Types
26//--------------------------------------------------------------------------------------------------
27
28/// Agent relay protocol specialization; it owns no router or SDK services.
29pub struct AgentProtocol;
30
31/// Wire form selected from the relay prologue, separate from capability gates.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum AgentWireFormat {
34    /// Current relay prologue and envelope generation.
35    Current,
36    /// Supported pre-0.5 relay prologue and generation-one envelope.
37    LegacyV1,
38}
39
40/// Immutable metadata captured once during relay setup.
41#[derive(Debug, Clone)]
42pub struct AgentReady {
43    /// Relay wire form selected during the prologue.
44    pub wire_format: AgentWireFormat,
45    /// Smaller of host and peer generations, used for known-operation gates.
46    pub negotiated_version: u8,
47    /// Decoded original `core.ready` payload.
48    pub agent: Ready,
49    ready_body: Vec<u8>,
50}
51
52//--------------------------------------------------------------------------------------------------
53// Methods
54//--------------------------------------------------------------------------------------------------
55
56impl AgentWireFormat {
57    /// Envelope generation emitted on this wire form.
58    ///
59    /// Current peers historically receive the host generation even when the
60    /// feature gate negotiated lower; preserve that existing byte contract.
61    pub fn version(self) -> u8 {
62        match self {
63            Self::Current => PROTOCOL_VERSION,
64            Self::LegacyV1 => LEGACY_PROTOCOL_VERSION,
65        }
66    }
67}
68
69impl AgentReady {
70    /// Exact original ready envelope, including unknown fields.
71    pub fn ready_bytes(&self) -> &[u8] {
72        &self.ready_body
73    }
74
75    /// Whether the connected generation supports this known message.
76    pub fn supports(&self, message_type: MessageType) -> bool {
77        message_type.is_available_at(self.negotiated_version)
78    }
79
80    /// Self-reported package version, empty on older agents lacking the field.
81    pub fn agent_version(&self) -> &str {
82        &self.agent.agent_version
83    }
84
85    /// Whether the connection selected the supported pre-0.5 wire form.
86    pub fn is_legacy_protocol(&self) -> bool {
87        self.wire_format == AgentWireFormat::LegacyV1
88    }
89}
90
91impl AgentProtocol {
92    /// Validate a known operation against a separately retained generation.
93    pub fn ensure_version_compat_for(
94        message_type: MessageType,
95        negotiated: u8,
96    ) -> ClientResult<()> {
97        if message_type.is_available_at(negotiated) {
98            Ok(())
99        } else {
100            Err(ClientError::new(ErrorKind::UnsupportedOperation))
101        }
102    }
103}
104
105//--------------------------------------------------------------------------------------------------
106// Trait Implementations
107//--------------------------------------------------------------------------------------------------
108
109impl Protocol for AgentProtocol {
110    // The generation-eight relay permanently retires completed correlations.
111    const REUSE_IDS: bool = false;
112    type Ready = AgentReady;
113
114    fn establish(
115        mut stream: BoxTransport,
116        options: ConnectOptions,
117    ) -> BoxFuture<'static, ClientResult<Established<AgentReady>>> {
118        Box::pin(async move {
119            // The engine bounds this entire future with one setup deadline.
120            // Eight bytes disambiguate current [min,max] and legacy [offset,len].
121            let mut prologue = [0u8; 8];
122            stream.read_exact(&mut prologue).await?;
123            let first = u32::from_be_bytes(prologue[..4].try_into().unwrap());
124            let second = u32::from_be_bytes(prologue[4..].try_into().unwrap());
125            let legacy = (FRAME_HEADER_SIZE as u32..=MAX_FRAME_SIZE).contains(&second)
126                && (first == 0 || first >= second);
127            let (ids, frame, wire_format) = if legacy {
128                let frame = read_after_prefix(&mut stream, second).await?;
129                (
130                    IdRange {
131                        start: first.saturating_add(1),
132                        end_exclusive: first.saturating_add(LEGACY_RELAY_ID_RANGE_STEP).into(),
133                    },
134                    frame,
135                    AgentWireFormat::LegacyV1,
136                )
137            } else {
138                // Current ranges can start at zero, but ID zero is setup-only.
139                let ids = IdRange {
140                    start: first.max(1),
141                    end_exclusive: second.into(),
142                };
143                ids.validate()?;
144                let frame = codec::read_raw_frame(&mut stream)
145                    .await
146                    .map_err(|_| ClientError::new(ErrorKind::InvalidData))?;
147                (ids, frame, AgentWireFormat::Current)
148            };
149            ids.validate()?;
150            // Keep the historical ready decoder and feature-generation rule.
151            let ready_message = codec::raw_frame_to_message(frame.clone())
152                .map_err(|_| ClientError::new(ErrorKind::InvalidData))?;
153            if ready_message.t != MessageType::Ready {
154                return Err(ClientError::new(ErrorKind::InvalidData));
155            }
156            let agent = ready_message
157                .payload::<Ready>()
158                .map_err(|_| ClientError::new(ErrorKind::InvalidData))?;
159            let ready = AgentReady {
160                wire_format,
161                negotiated_version: wire_format.version().min(ready_message.v),
162                agent,
163                ready_body: frame.body,
164            };
165            if ready.is_legacy_protocol() {
166                tracing::warn!(
167                    "agent client: legacy pre-0.5 exec protocol; filesystem operations require a newer agent"
168                );
169            }
170            Ok(Established {
171                transport: stream,
172                codec: Arc::new(CborEnvelopeCodec),
173                ids,
174                ready,
175                limits: options.limits,
176            })
177        })
178    }
179
180    fn prepare(ready: &AgentReady, wire_name: &str) -> ClientResult<SendMetadata> {
181        let flags = match MessageType::from_wire_str(wire_name) {
182            Some(message_type) => {
183                Self::ensure_version_compat_for(message_type, ready.negotiated_version)?;
184                message_type.flags()
185            }
186            // Dynamic names need no schema registration. Raw sends also support
187            // caller-selected flags and generations.
188            None => 0,
189        };
190        Ok(SendMetadata {
191            generation: ready.wire_format.version(),
192            flags,
193        })
194    }
195}
196
197//--------------------------------------------------------------------------------------------------
198// Functions
199//--------------------------------------------------------------------------------------------------
200
201async fn read_after_prefix(stream: &mut BoxTransport, length: u32) -> ClientResult<RawFrame> {
202    if !(FRAME_HEADER_SIZE as u32..=MAX_FRAME_SIZE).contains(&length) {
203        return Err(ClientError::new(ErrorKind::InvalidData));
204    }
205    let mut header = [0u8; FRAME_HEADER_SIZE];
206    stream.read_exact(&mut header).await?;
207    let mut body = vec![0; length as usize - FRAME_HEADER_SIZE];
208    stream.read_exact(&mut body).await?;
209    Ok(RawFrame {
210        id: u32::from_be_bytes(header[..4].try_into().unwrap()),
211        flags: header[4],
212        body,
213    })
214}