microsandbox_agent_client/
protocol.rs1use 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
16pub const LEGACY_PROTOCOL_VERSION: u8 = 1;
22const LEGACY_RELAY_ID_RANGE_STEP: u32 = u32::MAX / 16;
23
24pub struct AgentProtocol;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum AgentWireFormat {
34 Current,
36 LegacyV1,
38}
39
40#[derive(Debug, Clone)]
42pub struct AgentReady {
43 pub wire_format: AgentWireFormat,
45 pub negotiated_version: u8,
47 pub agent: Ready,
49 ready_body: Vec<u8>,
50}
51
52impl AgentWireFormat {
57 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 pub fn ready_bytes(&self) -> &[u8] {
72 &self.ready_body
73 }
74
75 pub fn supports(&self, message_type: MessageType) -> bool {
77 message_type.is_available_at(self.negotiated_version)
78 }
79
80 pub fn agent_version(&self) -> &str {
82 &self.agent.agent_version
83 }
84
85 pub fn is_legacy_protocol(&self) -> bool {
87 self.wire_format == AgentWireFormat::LegacyV1
88 }
89}
90
91impl AgentProtocol {
92 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
105impl Protocol for AgentProtocol {
110 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 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 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 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 None => 0,
189 };
190 Ok(SendMetadata {
191 generation: ready.wire_format.version(),
192 flags,
193 })
194 }
195}
196
197async 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}