antigravity_codes/
handshake.rs

1//! The length-prefixed stdio handshake that precedes the WebSocket session.
2//!
3//! Everything else in this protocol is protobuf-*JSON*, but the very first
4//! exchange is binary protobuf on the harness's stdin/stdout:
5//!
6//! ```text
7//! client -> harness   u32le length, then a serialised InputConfig
8//! harness -> client   u32le length, then a serialised OutputConfig
9//! ```
10//!
11//! `OutputConfig` carries the loopback port the harness just bound and the
12//! per-process API key the WebSocket upgrade must present. Two messages, seven
13//! fields between them, so this module hand-rolls the wire format rather than
14//! taking a protobuf runtime dependency for it — see
15//! `scripts/codegen_antigravity.py` for why a runtime is otherwise unnecessary.
16//!
17//! The field numbers below are pinned by the descriptor snapshot in
18//! `tests/schemas/localharness.descriptor.bin`, and this module's tests lock them
19//! against
20//! bytes captured from a live 0.1.10 harness.
21
22use crate::protocol::{ClientInfo, InputConfig, OutputConfig};
23
24/// A protobuf wire type. Only these two occur in the handshake messages.
25const WIRE_VARINT: u8 = 0;
26const WIRE_LEN: u8 = 2;
27
28fn tag(out: &mut Vec<u8>, field: u32, wire: u8) {
29    put_varint(out, u64::from(field) << 3 | u64::from(wire));
30}
31
32fn put_varint(out: &mut Vec<u8>, mut v: u64) {
33    loop {
34        let byte = (v & 0x7f) as u8;
35        v >>= 7;
36        if v == 0 {
37            out.push(byte);
38            return;
39        }
40        out.push(byte | 0x80);
41    }
42}
43
44fn put_string(out: &mut Vec<u8>, field: u32, v: &str) {
45    tag(out, field, WIRE_LEN);
46    put_varint(out, v.len() as u64);
47    out.extend_from_slice(v.as_bytes());
48}
49
50fn put_message(out: &mut Vec<u8>, field: u32, body: &[u8]) {
51    tag(out, field, WIRE_LEN);
52    put_varint(out, body.len() as u64);
53    out.extend_from_slice(body);
54}
55
56fn put_u32(out: &mut Vec<u8>, field: u32, v: u32) {
57    tag(out, field, WIRE_VARINT);
58    put_varint(out, u64::from(v));
59}
60
61fn encode_client_info(ci: &ClientInfo) -> Vec<u8> {
62    let mut out = Vec::new();
63    if let Some(v) = &ci.language {
64        put_string(&mut out, 1, v);
65    }
66    if let Some(v) = &ci.version {
67        put_string(&mut out, 2, v);
68    }
69    if let Some(v) = &ci.language_version {
70        put_string(&mut out, 3, v);
71    }
72    if let Some(v) = &ci.os {
73        put_string(&mut out, 4, v);
74    }
75    if let Some(v) = &ci.os_version {
76        put_string(&mut out, 5, v);
77    }
78    out
79}
80
81/// Serialises an [`InputConfig`] to binary protobuf, fields in tag order.
82pub fn encode_input_config(cfg: &InputConfig) -> Vec<u8> {
83    let mut out = Vec::new();
84    if let Some(v) = &cfg.storage_directory {
85        put_string(&mut out, 1, v);
86    }
87    if let Some(v) = cfg.port {
88        put_u32(&mut out, 2, v);
89    }
90    if let Some(v) = &cfg.bind_address {
91        put_string(&mut out, 3, v);
92    }
93    if let Some(ci) = &cfg.client_info {
94        put_message(&mut out, 4, &encode_client_info(ci));
95    }
96    // `map<string, string> env = 5` — each entry is a synthetic message with
97    // `key = 1` and `value = 2`. Sorted so the encoding is deterministic.
98    let mut env: Vec<_> = cfg.env.iter().collect();
99    env.sort_by(|a, b| a.0.cmp(b.0));
100    for (k, v) in env {
101        let mut entry = Vec::new();
102        put_string(&mut entry, 1, k);
103        put_string(&mut entry, 2, v);
104        put_message(&mut out, 5, &entry);
105    }
106    out
107}
108
109/// Prefixes a serialised message with its length, as the harness expects.
110pub fn frame(body: &[u8]) -> Vec<u8> {
111    let mut out = Vec::with_capacity(body.len() + 4);
112    out.extend_from_slice(&(body.len() as u32).to_le_bytes());
113    out.extend_from_slice(body);
114    out
115}
116
117/// A malformed handshake reply.
118#[derive(Debug, thiserror::Error)]
119pub enum DecodeError {
120    /// The buffer ended in the middle of a field.
121    #[error("truncated protobuf message at byte {0}")]
122    Truncated(usize),
123    /// A varint ran past the 10 bytes a u64 can occupy.
124    #[error("malformed varint at byte {0}")]
125    BadVarint(usize),
126    /// A field used a wire type the handshake never emits.
127    #[error("unsupported wire type {wire} for field {field}")]
128    UnsupportedWireType {
129        /// The protobuf field number that carried it.
130        field: u32,
131        /// The wire type encountered.
132        wire: u8,
133    },
134    /// A string field was not valid UTF-8.
135    #[error("field {0} is not valid UTF-8")]
136    NotUtf8(u32),
137}
138
139struct Reader<'a> {
140    buf: &'a [u8],
141    pos: usize,
142}
143
144impl<'a> Reader<'a> {
145    fn varint(&mut self) -> Result<u64, DecodeError> {
146        let mut value = 0u64;
147        let mut shift = 0;
148        loop {
149            let byte = *self
150                .buf
151                .get(self.pos)
152                .ok_or(DecodeError::Truncated(self.pos))?;
153            self.pos += 1;
154            value |= u64::from(byte & 0x7f) << shift;
155            if byte & 0x80 == 0 {
156                return Ok(value);
157            }
158            shift += 7;
159            if shift >= 64 {
160                return Err(DecodeError::BadVarint(self.pos));
161            }
162        }
163    }
164
165    fn bytes(&mut self) -> Result<&'a [u8], DecodeError> {
166        let len = self.varint()? as usize;
167        let end = self
168            .pos
169            .checked_add(len)
170            .ok_or(DecodeError::Truncated(self.pos))?;
171        let slice = self
172            .buf
173            .get(self.pos..end)
174            .ok_or(DecodeError::Truncated(self.pos))?;
175        self.pos = end;
176        Ok(slice)
177    }
178
179    /// Steps over a field this decoder does not care about.
180    fn skip(&mut self, field: u32, wire: u8) -> Result<(), DecodeError> {
181        match wire {
182            WIRE_VARINT => {
183                self.varint()?;
184            }
185            WIRE_LEN => {
186                self.bytes()?;
187            }
188            1 => self.pos += 8,
189            5 => self.pos += 4,
190            _ => return Err(DecodeError::UnsupportedWireType { field, wire }),
191        }
192        Ok(())
193    }
194}
195
196/// Parses the `OutputConfig` the harness writes to stdout.
197///
198/// Unknown fields are skipped, so a newer harness that adds to this message
199/// still hands back a usable port and key.
200pub fn decode_output_config(buf: &[u8]) -> Result<OutputConfig, DecodeError> {
201    let mut r = Reader { buf, pos: 0 };
202    let mut cfg = OutputConfig::default();
203    while r.pos < buf.len() {
204        let key = r.varint()?;
205        let field = (key >> 3) as u32;
206        let wire = (key & 7) as u8;
207        match (field, wire) {
208            (1, WIRE_VARINT) => cfg.port = Some(r.varint()? as i32),
209            (2, WIRE_LEN) => {
210                let raw = r.bytes()?;
211                cfg.api_key = Some(
212                    std::str::from_utf8(raw)
213                        .map_err(|_| DecodeError::NotUtf8(2))?
214                        .to_string(),
215                );
216            }
217            _ => r.skip(field, wire)?,
218        }
219    }
220    Ok(cfg)
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    /// Captured from a live 0.1.10 harness handshake.
228    const GOLDEN_INPUT: &str = "0a0c2f746d702f616773746f7265221e0a0472757374120\
2296302e312e31301a04312e383522056c696e75782a0136";
230    const GOLDEN_OUTPUT: &str = "08e5bc021220373130643035656337373862323462363338\
2313663636632396530316536613031";
232
233    fn unhex(s: &str) -> Vec<u8> {
234        (0..s.len())
235            .step_by(2)
236            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
237            .collect()
238    }
239
240    #[test]
241    fn encodes_byte_for_byte_like_the_reference_client() {
242        let cfg = InputConfig {
243            storage_directory: Some("/tmp/agstore".into()),
244            client_info: Some(ClientInfo {
245                language: Some("rust".into()),
246                version: Some("0.1.10".into()),
247                language_version: Some("1.85".into()),
248                os: Some("linux".into()),
249                os_version: Some("6".into()),
250            }),
251            ..Default::default()
252        };
253        assert_eq!(encode_input_config(&cfg), unhex(GOLDEN_INPUT));
254    }
255
256    #[test]
257    fn decodes_a_captured_output_config() {
258        let cfg = decode_output_config(&unhex(GOLDEN_OUTPUT)).unwrap();
259        assert_eq!(cfg.port, Some(40549));
260        assert_eq!(
261            cfg.api_key.as_deref(),
262            Some("710d05ec778b24b6386ccf29e01e6a01")
263        );
264    }
265
266    #[test]
267    fn skips_fields_a_newer_harness_might_add() {
268        let mut buf = unhex(GOLDEN_OUTPUT);
269        put_string(&mut buf, 9, "something-new");
270        put_u32(&mut buf, 10, 42);
271        let cfg = decode_output_config(&buf).unwrap();
272        assert_eq!(cfg.port, Some(40549));
273    }
274
275    #[test]
276    fn rejects_a_truncated_message() {
277        let buf = unhex(GOLDEN_OUTPUT);
278        assert!(decode_output_config(&buf[..buf.len() - 4]).is_err());
279    }
280
281    #[test]
282    fn frames_with_a_little_endian_length() {
283        assert_eq!(frame(&[1, 2, 3]), vec![3, 0, 0, 0, 1, 2, 3]);
284    }
285
286    #[test]
287    fn encodes_env_entries_deterministically() {
288        let cfg = InputConfig {
289            env: [
290                ("B".to_string(), "2".to_string()),
291                ("A".to_string(), "1".to_string()),
292            ]
293            .into_iter()
294            .collect(),
295            ..Default::default()
296        };
297        let once = encode_input_config(&cfg);
298        assert_eq!(once, encode_input_config(&cfg));
299        // Entry for "A" sorts first: tag 5 (0x2a), len 6, key "A", value "1".
300        assert_eq!(&once[..2], &[0x2a, 0x06]);
301    }
302}
antigravity_codes/handshake.rs

antigravity_codes/
handshake.rs
