netget/server/capnp_rpc/
schema.rs

1//! A Cap'n Proto schema loaded at startup from `capnp compile -o-` output (a
2//! CodeGeneratorRequest), and the JSON mapping it drives: a struct's fields by name, enums by
3//! enumerant name, unions as the one member that is set, groups as nested objects, Data as
4//! `{"$hex": …}`. The layout offsets of schema.capnp itself are written out below, read from
5//! `capnp compile -ocapnp schema.capnp`.
6use super::layout::{
7    Builder, ElementSize, ListReader, Message, StructBuilder, StructReader, Target,
8};
9use anyhow::{bail, ensure, Context, Result};
10use serde_json::{json, Map, Value};
11use std::collections::HashMap;
12
13/// Largest compiled schema accepted.
14pub const MAX_SCHEMA_BYTES: usize = 4 * 1024 * 1024;
15/// Most elements in one list the codec writes from JSON.
16pub const MAX_LIST: usize = 65_536;
17
18#[derive(Clone, Debug)]
19pub enum Type {
20    Void,
21    Bool,
22    Int(u8),
23    UInt(u8),
24    Float32,
25    Float64,
26    Text,
27    Data,
28    List(Box<Type>),
29    Enum(u64),
30    Struct(u64),
31    /// Interfaces and AnyPointer: not mapped (read as null, written as null).
32    Opaque,
33}
34
35#[derive(Clone, Debug)]
36pub enum FieldKind {
37    Slot { offset: u32, ty: Type, default: u64 },
38    Group(u64),
39}
40
41#[derive(Clone, Debug)]
42pub struct Field {
43    pub name: String,
44    /// Set for a member of the struct's union.
45    pub discriminant: Option<u16>,
46    pub kind: FieldKind,
47}
48
49#[derive(Clone, Debug)]
50pub struct StructInfo {
51    pub name: String,
52    pub data_words: u16,
53    pub ptr_count: u16,
54    pub discriminant_offset: u32,
55    pub fields: Vec<Field>,
56}
57
58#[derive(Clone, Debug)]
59pub struct Method {
60    pub name: String,
61    pub id: u16,
62    pub params: u64,
63    pub results: u64,
64}
65
66#[derive(Clone, Debug)]
67pub struct InterfaceInfo {
68    pub name: String,
69    pub id: u64,
70    pub methods: Vec<Method>,
71    pub superclasses: Vec<u64>,
72}
73
74#[derive(Debug, Default)]
75pub struct Schema {
76    pub structs: HashMap<u64, StructInfo>,
77    pub enums: HashMap<u64, Vec<String>>,
78    pub interfaces: HashMap<u64, InterfaceInfo>,
79}
80
81fn short_name(display: &str, prefix: u32) -> String {
82    display
83        .get(prefix as usize..)
84        .unwrap_or(display)
85        .to_string()
86}
87
88fn read_type(t: StructReader<'_>) -> Result<Type> {
89    Ok(match t.u16(0) {
90        0 => Type::Void,
91        1 => Type::Bool,
92        2 => Type::Int(8),
93        3 => Type::Int(16),
94        4 => Type::Int(32),
95        5 => Type::Int(64),
96        6 => Type::UInt(8),
97        7 => Type::UInt(16),
98        8 => Type::UInt(32),
99        9 => Type::UInt(64),
100        10 => Type::Float32,
101        11 => Type::Float64,
102        12 => Type::Text,
103        13 => Type::Data,
104        14 => Type::List(Box::new(read_type(
105            t.struct_field(0)?.context("list without element type")?,
106        )?)),
107        15 => Type::Enum(t.u64(1)),
108        16 => Type::Struct(t.u64(1)),
109        _ => Type::Opaque,
110    })
111}
112
113/// The raw bits of a primitive default, so a field reads as `stored ^ default`.
114fn read_default(v: Option<StructReader<'_>>, ty: &Type) -> u64 {
115    let Some(v) = v else { return 0 };
116    match ty {
117        Type::Bool => v.bits(16, 1),
118        Type::Int(8) | Type::UInt(8) => v.bits(16, 8),
119        Type::Int(16) | Type::UInt(16) | Type::Enum(_) => v.bits(16, 16),
120        Type::Int(32) | Type::UInt(32) | Type::Float32 => v.bits(32, 32),
121        Type::Int(64) | Type::UInt(64) | Type::Float64 => v.u64(1),
122        _ => 0,
123    }
124}
125
126impl Schema {
127    /// Parse a CodeGeneratorRequest.
128    pub fn load(bytes: &[u8]) -> Result<Self> {
129        ensure!(bytes.len() <= MAX_SCHEMA_BYTES, "compiled schema too large");
130        let segments = frame_segments(bytes)?;
131        let msg = Message::new(segments);
132        let request = msg.root_struct()?;
133        let nodes = request.list_field(0)?.context("no nodes in the schema")?;
134        let mut schema = Schema::default();
135        for i in 0..nodes.len {
136            let node = nodes.struct_at(i)?;
137            let id = node.u64(0);
138            let name = short_name(&node.text(0)?.unwrap_or_default(), node.u32(2));
139            match node.u16(6) {
140                1 => {
141                    let mut fields = Vec::new();
142                    if let Some(list) = node.list_field(3)? {
143                        for f in 0..list.len {
144                            let f = list.struct_at(f)?;
145                            let discriminant = match f.u16(1) ^ 0xffff {
146                                0xffff => None,
147                                d => Some(d),
148                            };
149                            let kind = match f.u16(4) {
150                                0 => {
151                                    let ty = read_type(
152                                        f.struct_field(2)?.context("slot without type")?,
153                                    )?;
154                                    let default = read_default(f.struct_field(3)?, &ty);
155                                    FieldKind::Slot {
156                                        offset: f.u32(1),
157                                        ty,
158                                        default,
159                                    }
160                                }
161                                _ => FieldKind::Group(f.u64(2)),
162                            };
163                            fields.push(Field {
164                                name: f.text(0)?.unwrap_or_default(),
165                                discriminant,
166                                kind,
167                            });
168                        }
169                    }
170                    schema.structs.insert(
171                        id,
172                        StructInfo {
173                            name,
174                            data_words: node.u16(7),
175                            ptr_count: node.u16(12),
176                            discriminant_offset: node.u32(8),
177                            fields,
178                        },
179                    );
180                }
181                2 => {
182                    let mut names = Vec::new();
183                    if let Some(list) = node.list_field(3)? {
184                        for e in 0..list.len {
185                            names.push(list.struct_at(e)?.text(0)?.unwrap_or_default());
186                        }
187                    }
188                    schema.enums.insert(id, names);
189                }
190                3 => {
191                    let mut methods = Vec::new();
192                    if let Some(list) = node.list_field(3)? {
193                        for m in 0..list.len {
194                            let m_reader = list.struct_at(m)?;
195                            methods.push(Method {
196                                name: m_reader.text(0)?.unwrap_or_default(),
197                                id: m as u16,
198                                params: m_reader.u64(1),
199                                results: m_reader.u64(2),
200                            });
201                        }
202                    }
203                    let mut superclasses = Vec::new();
204                    if let Some(list) = node.list_field(4)? {
205                        for s in 0..list.len {
206                            superclasses.push(list.struct_at(s)?.u64(0));
207                        }
208                    }
209                    schema.interfaces.insert(
210                        id,
211                        InterfaceInfo {
212                            name,
213                            id,
214                            methods,
215                            superclasses,
216                        },
217                    );
218                }
219                _ => {}
220            }
221        }
222        Ok(schema)
223    }
224
225    /// The interface named `name` (its short name, or its full display name).
226    pub fn interface(&self, name: &str) -> Result<&InterfaceInfo> {
227        let mut found = self
228            .interfaces
229            .values()
230            .filter(|i| i.name == name || i.name.rsplit(['.', ':']).next() == Some(name));
231        let first = found
232            .next()
233            .with_context(|| format!("the schema has no interface named {name}"))?;
234        ensure!(found.next().is_none(), "interface name {name} is ambiguous");
235        Ok(first)
236    }
237
238    /// Every method `iface` answers, its own and its superclasses', as (interface id, method).
239    pub fn methods_of(&self, iface: u64) -> Vec<(u64, &InterfaceInfo, &Method)> {
240        let mut out = Vec::new();
241        let mut pending = vec![iface];
242        let mut seen = Vec::new();
243        while let Some(id) = pending.pop() {
244            if seen.contains(&id) {
245                continue;
246            }
247            seen.push(id);
248            if let Some(i) = self.interfaces.get(&id) {
249                out.extend(i.methods.iter().map(|m| (id, i, m)));
250                pending.extend(i.superclasses.iter().copied());
251            }
252        }
253        out
254    }
255
256    fn info(&self, id: u64) -> Result<&StructInfo> {
257        self.structs
258            .get(&id)
259            .with_context(|| format!("struct {id:#x} is not in the schema"))
260    }
261
262    pub fn struct_size(&self, id: u64) -> Result<(u16, u16)> {
263        let s = self.info(id)?;
264        Ok((s.data_words, s.ptr_count))
265    }
266
267    /// A struct as JSON.
268    pub fn to_json(&self, id: u64, s: StructReader<'_>) -> Result<Value> {
269        let info = self.info(id)?;
270        let active = s.u16(u64::from(info.discriminant_offset));
271        let mut out = Map::new();
272        for f in &info.fields {
273            if f.discriminant.is_some_and(|d| d != active) {
274                continue;
275            }
276            let v = match &f.kind {
277                FieldKind::Group(gid) => self.to_json(*gid, s)?,
278                FieldKind::Slot {
279                    offset,
280                    ty,
281                    default,
282                } => self.slot_to_json(s, *offset, ty, *default)?,
283            };
284            out.insert(f.name.clone(), v);
285        }
286        Ok(Value::Object(out))
287    }
288
289    fn slot_to_json(
290        &self,
291        s: StructReader<'_>,
292        offset: u32,
293        ty: &Type,
294        default: u64,
295    ) -> Result<Value> {
296        let o = u64::from(offset);
297        Ok(match ty {
298            Type::Void => Value::Null,
299            Type::Bool => json!(s.bits(o, 1) ^ default != 0),
300            Type::Int(w) | Type::UInt(w) => {
301                let raw = s.bits(o * u64::from(*w), u64::from(*w)) ^ default;
302                integer(ty, raw)
303            }
304            Type::Float32 => json!(f32::from_bits((s.bits(o * 32, 32) ^ default) as u32)),
305            Type::Float64 => json!(f64::from_bits(s.bits(o * 64, 64) ^ default)),
306            Type::Enum(eid) => self.enum_name(*eid, (s.bits(o * 16, 16) ^ default) as u16),
307            Type::Text => s.text(offset as u16)?.map_or(Value::Null, Value::String),
308            Type::Data => s
309                .data(offset as u16)?
310                .map_or(Value::Null, |b| json!({"$hex": hex::encode(b)})),
311            Type::Struct(sid) => match s.struct_field(offset as u16)? {
312                Some(c) => self.to_json(*sid, c)?,
313                None => Value::Null,
314            },
315            Type::List(el) => match s.list_field(offset as u16)? {
316                Some(l) => self.list_to_json(el, l)?,
317                None => Value::Null,
318            },
319            Type::Opaque => Value::Null,
320        })
321    }
322
323    fn enum_name(&self, id: u64, v: u16) -> Value {
324        match self.enums.get(&id).and_then(|n| n.get(v as usize)) {
325            Some(n) => json!(n),
326            None => json!(v),
327        }
328    }
329
330    fn list_to_json(&self, el: &Type, l: ListReader<'_>) -> Result<Value> {
331        let mut out = Vec::with_capacity(l.len.min(4096) as usize);
332        for i in 0..l.len {
333            out.push(match el {
334                Type::Struct(sid) => self.to_json(*sid, l.struct_at(i)?)?,
335                Type::Text | Type::Data | Type::List(_) | Type::Opaque => match l.pointer_at(i)? {
336                    Target::Null => Value::Null,
337                    Target::List(inner) => match el {
338                        Type::Text => json!(inner.text()?),
339                        Type::Data => json!({"$hex": hex::encode(inner.bytes()?)}),
340                        Type::List(e) => self.list_to_json(e, inner)?,
341                        _ => Value::Null,
342                    },
343                    _ => Value::Null,
344                },
345                Type::Void => Value::Null,
346                Type::Bool => json!(l.primitive(i)? != 0),
347                Type::Int(_) | Type::UInt(_) => integer(el, l.primitive(i)?),
348                Type::Float32 => json!(f32::from_bits(l.primitive(i)? as u32)),
349                Type::Float64 => json!(f64::from_bits(l.primitive(i)?)),
350                Type::Enum(eid) => self.enum_name(*eid, l.primitive(i)? as u16),
351            });
352        }
353        Ok(Value::Array(out))
354    }
355
356    /// Fill struct `s` (already allocated with this struct's size) from a JSON object. Unknown
357    /// keys are refused, so a misspelt field is an error rather than silently dropped.
358    pub fn from_json(&self, b: &mut Builder, id: u64, s: StructBuilder, v: &Value) -> Result<()> {
359        let info = self.info(id)?;
360        let obj = match v {
361            Value::Null => return self.set_defaults(b, info, s),
362            Value::Object(o) => o,
363            _ => bail!("{} must be a JSON object", info.name),
364        };
365        for key in obj.keys() {
366            ensure!(
367                info.fields.iter().any(|f| &f.name == key),
368                "{} has no field {key}",
369                info.name
370            );
371        }
372        let members: Vec<&Field> = info
373            .fields
374            .iter()
375            .filter(|f| f.discriminant.is_some() && obj.contains_key(&f.name))
376            .collect();
377        ensure!(
378            members.len() <= 1,
379            "{}: set at most one union member",
380            info.name
381        );
382        self.set_defaults(b, info, s)?;
383        if let Some(d) = members.first().and_then(|f| f.discriminant) {
384            b.set_u16(s, u64::from(info.discriminant_offset), d);
385        }
386        for f in &info.fields {
387            let Some(value) = obj.get(&f.name) else {
388                continue;
389            };
390            let at = |e: anyhow::Error| e.context(format!("{}.{}", info.name, f.name));
391            match &f.kind {
392                FieldKind::Group(gid) => self.from_json(b, *gid, s, value).map_err(at)?,
393                FieldKind::Slot {
394                    offset,
395                    ty,
396                    default,
397                } => self
398                    .set_slot(b, s, *offset, ty, *default, value)
399                    .map_err(at)?,
400            }
401        }
402        Ok(())
403    }
404
405    /// Write every primitive default, so an absent field reads as its declared default.
406    fn set_defaults(&self, b: &mut Builder, info: &StructInfo, s: StructBuilder) -> Result<()> {
407        // Stored bits are value ^ default, so a zero data section already reads as every
408        // default. Nothing to write.
409        let _ = (b, info, s);
410        Ok(())
411    }
412
413    fn set_slot(
414        &self,
415        b: &mut Builder,
416        s: StructBuilder,
417        offset: u32,
418        ty: &Type,
419        default: u64,
420        v: &Value,
421    ) -> Result<()> {
422        let o = u64::from(offset);
423        if v.is_null() {
424            return Ok(());
425        }
426        match ty {
427            Type::Void | Type::Opaque => {}
428            Type::Bool => {
429                let x = v.as_bool().context("expected a boolean")?;
430                b.set_bits(s, o, 1, u64::from(x) ^ default);
431            }
432            Type::Int(w) | Type::UInt(w) => b.set_bits(
433                s,
434                o * u64::from(*w),
435                u64::from(*w),
436                int_bits(ty, v)? ^ default,
437            ),
438            Type::Float32 => {
439                let x = v.as_f64().context("expected a number")? as f32;
440                b.set_bits(s, o * 32, 32, u64::from(x.to_bits()) ^ default);
441            }
442            Type::Float64 => {
443                let x = v.as_f64().context("expected a number")?;
444                b.set_bits(s, o * 64, 64, x.to_bits() ^ default);
445            }
446            Type::Enum(eid) => {
447                let x = self.enum_value(*eid, v)?;
448                b.set_bits(s, o * 16, 16, u64::from(x) ^ default);
449            }
450            Type::Text => b.set_text(s, offset as u16, v.as_str().context("expected a string")?)?,
451            Type::Data => b.set_bytes(s, offset as u16, &data_bytes(v)?)?,
452            Type::Struct(sid) => {
453                let (dw, pc) = self.struct_size(*sid)?;
454                let child = b.init_struct(s, offset as u16, dw, pc)?;
455                self.from_json(b, *sid, child, v)?;
456            }
457            Type::List(el) => {
458                let slot = Builder::pointer_slot_of(s, offset as u16)?;
459                self.list_from_json(b, slot, el, v)?;
460            }
461        }
462        Ok(())
463    }
464
465    fn enum_value(&self, id: u64, v: &Value) -> Result<u16> {
466        if let Some(n) = v.as_u64() {
467            return u16::try_from(n).context("enum value out of range");
468        }
469        let name = v.as_str().context("expected an enumerant name")?;
470        self.enums
471            .get(&id)
472            .and_then(|names| names.iter().position(|n| n == name))
473            .map(|p| p as u16)
474            .with_context(|| format!("no enumerant {name}"))
475    }
476
477    fn list_from_json(&self, b: &mut Builder, slot: usize, el: &Type, v: &Value) -> Result<()> {
478        let items = v.as_array().context("expected an array")?;
479        ensure!(items.len() <= MAX_LIST, "list longer than {MAX_LIST}");
480        let len = items.len() as u32;
481        match el {
482            Type::Struct(sid) => {
483                let (dw, pc) = self.struct_size(*sid)?;
484                let elements = b.slot_struct_list(slot, len, dw, pc)?;
485                for (e, item) in elements.into_iter().zip(items) {
486                    self.from_json(b, *sid, e, item)?;
487                }
488            }
489            Type::Text | Type::Data | Type::List(_) | Type::Opaque => {
490                let slots = b.slot_pointer_list(slot, len)?;
491                for (p, item) in slots.into_iter().zip(items) {
492                    if item.is_null() {
493                        continue;
494                    }
495                    match el {
496                        Type::Text => {
497                            let mut bytes = item
498                                .as_str()
499                                .context("expected a string")?
500                                .as_bytes()
501                                .to_vec();
502                            bytes.push(0);
503                            b.slot_bytes(p, &bytes)?;
504                        }
505                        Type::Data => b.slot_bytes(p, &data_bytes(item)?)?,
506                        Type::List(e) => self.list_from_json(b, p, e, item)?,
507                        _ => {}
508                    }
509                }
510            }
511            _ => {
512                let size = match el {
513                    Type::Void => ElementSize::Void,
514                    Type::Bool => ElementSize::Bit,
515                    Type::Int(8) | Type::UInt(8) => ElementSize::Byte,
516                    Type::Int(16) | Type::UInt(16) | Type::Enum(_) => ElementSize::TwoBytes,
517                    Type::Int(32) | Type::UInt(32) | Type::Float32 => ElementSize::FourBytes,
518                    _ => ElementSize::EightBytes,
519                };
520                let values = items
521                    .iter()
522                    .map(|item| {
523                        Ok(match el {
524                            Type::Void => 0,
525                            Type::Bool => u64::from(item.as_bool().context("expected a boolean")?),
526                            Type::Int(_) | Type::UInt(_) => int_bits(el, item)?,
527                            Type::Float32 => u64::from(
528                                (item.as_f64().context("expected a number")? as f32).to_bits(),
529                            ),
530                            Type::Float64 => item.as_f64().context("expected a number")?.to_bits(),
531                            Type::Enum(eid) => u64::from(self.enum_value(*eid, item)?),
532                            _ => 0,
533                        })
534                    })
535                    .collect::<Result<Vec<u64>>>()?;
536                b.slot_primitive_list(slot, size, &values)?;
537            }
538        }
539        Ok(())
540    }
541
542    /// A short description of a method's parameters or results, for the model's prompt. A
543    /// struct already being described (a recursive type) is named rather than expanded.
544    pub fn describe(&self, id: u64) -> String {
545        self.describe_in(id, &mut Vec::new())
546    }
547
548    fn describe_in(&self, id: u64, open: &mut Vec<u64>) -> String {
549        let Ok(info) = self.info(id) else {
550            return String::new();
551        };
552        open.push(id);
553        let out = info
554            .fields
555            .iter()
556            .map(|f| match &f.kind {
557                FieldKind::Slot { ty, .. } => format!("{}: {}", f.name, self.type_name(ty, open)),
558                FieldKind::Group(g) => format!("{}: {{{}}}", f.name, self.describe_in(*g, open)),
559            })
560            .collect::<Vec<_>>()
561            .join(", ");
562        open.pop();
563        out
564    }
565
566    fn type_name(&self, ty: &Type, open: &mut Vec<u64>) -> String {
567        match ty {
568            Type::Void => "Void".into(),
569            Type::Bool => "Bool".into(),
570            Type::Int(w) => format!("Int{w}"),
571            Type::UInt(w) => format!("UInt{w}"),
572            Type::Float32 => "Float32".into(),
573            Type::Float64 => "Float64".into(),
574            Type::Text => "Text".into(),
575            Type::Data => "Data {\"$hex\"}".into(),
576            Type::List(e) => format!("List({})", self.type_name(e, open)),
577            Type::Enum(id) => match self.enums.get(id) {
578                Some(n) => format!("enum {}", n.join("|")),
579                None => "enum".into(),
580            },
581            Type::Struct(id) if open.contains(id) || open.len() >= 8 => self
582                .structs
583                .get(id)
584                .map_or_else(|| "struct".into(), |s| s.name.clone()),
585            Type::Struct(id) => format!("{{{}}}", self.describe_in(*id, open)),
586            Type::Opaque => "unsupported".into(),
587        }
588    }
589}
590
591fn integer(ty: &Type, raw: u64) -> Value {
592    match ty {
593        Type::Int(8) => json!(raw as u8 as i8),
594        Type::Int(16) => json!(raw as u16 as i16),
595        Type::Int(32) => json!(raw as u32 as i32),
596        Type::Int(_) => json!(raw as i64),
597        _ => json!(raw),
598    }
599}
600
601fn int_bits(ty: &Type, v: &Value) -> Result<u64> {
602    let (signed, width) = match ty {
603        Type::Int(w) => (true, *w),
604        Type::UInt(w) => (false, *w),
605        _ => bail!("not an integer type"),
606    };
607    if signed {
608        let x = v.as_i64().context("expected an integer")?;
609        let (min, max) = if width == 64 {
610            (i64::MIN, i64::MAX)
611        } else {
612            (-(1i64 << (width - 1)), (1i64 << (width - 1)) - 1)
613        };
614        ensure!((min..=max).contains(&x), "Int{width} out of range");
615        Ok(if width == 64 {
616            x as u64
617        } else {
618            (x as u64) & ((1u64 << width) - 1)
619        })
620    } else {
621        let x = v.as_u64().context("expected a non-negative integer")?;
622        ensure!(
623            width == 64 || x < (1u64 << width),
624            "UInt{width} out of range"
625        );
626        Ok(x)
627    }
628}
629
630fn data_bytes(v: &Value) -> Result<Vec<u8>> {
631    let h = v["$hex"]
632        .as_str()
633        .context("Data must be {\"$hex\": \"…\"}")?;
634    hex::decode(h).context("invalid hex")
635}
636
637/// A framed message held in memory.
638pub fn frame_segments(bytes: &[u8]) -> Result<Vec<Vec<u64>>> {
639    ensure!(bytes.len() >= 8, "truncated message");
640    let count = u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize + 1;
641    ensure!(count <= super::layout::MAX_SEGMENTS, "too many segments");
642    let header = (4 + 4 * count).div_ceil(8) * 8;
643    ensure!(bytes.len() >= header, "truncated segment table");
644    let mut at = header;
645    let mut out = Vec::with_capacity(count);
646    for i in 0..count {
647        let size = u32::from_le_bytes(bytes[4 + 4 * i..8 + 4 * i].try_into().unwrap()) as usize;
648        let end = at
649            .checked_add(size.checked_mul(8).context("segment size overflow")?)
650            .context("segment size overflow")?;
651        ensure!(end <= bytes.len(), "truncated segment");
652        out.push(
653            bytes[at..end]
654                .chunks_exact(8)
655                .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
656                .collect(),
657        );
658        at = end;
659    }
660    Ok(out)
661}
netget/server/capnp_rpc/schema.rs

netget/server/capnp_rpc/
schema.rs
