use std::collections::HashMap;
use crate::error::{Error, Result};
use crate::resolve::Resolver;
use crate::schema::Schema;
pub const BUNDLE_MAGIC: &[u8; 4] = b"VRSB";
pub const BUNDLE_VERSION: u8 = 1;
const BUNDLE_HEADER_LEN: usize = 12;
#[derive(Clone, Debug, Default)]
pub struct SchemaRegistry {
by_id: HashMap<u128, Schema>,
}
impl SchemaRegistry {
pub fn new() -> SchemaRegistry {
SchemaRegistry {
by_id: HashMap::new(),
}
}
pub fn register(&mut self, schema: Schema) -> u128 {
let id = schema.id();
self.by_id.entry(id).or_insert(schema);
id
}
pub fn register_canonical(&mut self, bytes: &[u8]) -> Result<u128> {
Ok(self.register(Schema::from_canonical(bytes)?))
}
pub fn get(&self, id: u128) -> Option<&Schema> {
self.by_id.get(&id)
}
pub fn contains(&self, id: u128) -> bool {
self.by_id.contains_key(&id)
}
pub fn len(&self) -> usize {
self.by_id.len()
}
pub fn is_empty(&self) -> bool {
self.by_id.is_empty()
}
pub fn ids(&self) -> impl Iterator<Item = u128> + '_ {
self.by_id.keys().copied()
}
pub fn resolver_for(&self, writer_id: u128, reader: &Schema) -> Result<Resolver> {
let writer = self.get(writer_id).ok_or(Error::SchemaIdMismatch {
message: writer_id,
expected: reader.id(),
})?;
Resolver::new(writer, reader)
}
pub fn to_bundle(&self) -> Vec<u8> {
let mut ids: Vec<u128> = self.by_id.keys().copied().collect();
ids.sort_unstable();
let mut buf = Vec::new();
buf.extend_from_slice(BUNDLE_MAGIC);
buf.push(BUNDLE_VERSION);
buf.extend_from_slice(&[0u8; 3]); buf.extend_from_slice(&(ids.len() as u32).to_le_bytes());
for id in ids {
let canonical = self.by_id[&id].canonical_bytes();
buf.extend_from_slice(&(canonical.len() as u32).to_le_bytes());
buf.extend_from_slice(canonical);
}
buf
}
pub fn from_bundle(bytes: &[u8]) -> Result<SchemaRegistry> {
let mut reg = SchemaRegistry::new();
reg.merge_bundle(bytes)?;
Ok(reg)
}
pub fn merge_bundle(&mut self, bytes: &[u8]) -> Result<usize> {
if bytes.len() < BUNDLE_HEADER_LEN {
return Err(Error::BadSchema("schema bundle truncated".into()));
}
if &bytes[0..4] != BUNDLE_MAGIC {
return Err(Error::BadSchema("bad schema-bundle magic".into()));
}
if bytes[4] != BUNDLE_VERSION {
return Err(Error::BadSchema("unsupported schema-bundle version".into()));
}
let count = u32::from_le_bytes(bytes[8..12].try_into().unwrap()) as usize;
let mut pos = BUNDLE_HEADER_LEN;
let mut added = 0;
for _ in 0..count {
let end = pos
.checked_add(4)
.filter(|&e| e <= bytes.len())
.ok_or_else(|| Error::BadSchema("schema bundle truncated".into()))?;
let len = u32::from_le_bytes(bytes[pos..end].try_into().unwrap()) as usize;
pos = end;
let entry_end = pos
.checked_add(len)
.filter(|&e| e <= bytes.len())
.ok_or_else(|| Error::BadSchema("schema bundle entry out of bounds".into()))?;
let before = self.len();
self.register_canonical(&bytes[pos..entry_end])?;
if self.len() > before {
added += 1;
}
pos = entry_end;
}
if pos != bytes.len() {
return Err(Error::BadSchema(
"trailing bytes after schema bundle".into(),
));
}
Ok(added)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::{Dt, SchemaBuilder};
use crate::{encode, Message, SchemaMode, Value};
fn v0() -> Schema {
SchemaBuilder::new()
.add_struct(
"LogEvent",
vec![(1, "service", Dt::Str), (4, "latency", Dt::U16)],
)
.build("LogEvent")
.unwrap()
}
fn v1() -> Schema {
SchemaBuilder::new()
.add_struct(
"LogEvent",
vec![
(1, "service", Dt::Str),
(3, "message", Dt::Str),
(4, "latency", Dt::U32), ],
)
.build("LogEvent")
.unwrap()
}
#[test]
fn register_is_content_addressed_and_idempotent() {
let mut reg = SchemaRegistry::new();
let id_a = reg.register(v0());
let id_b = reg.register(v0());
assert_eq!(id_a, id_b);
assert_eq!(reg.len(), 1);
assert!(reg.contains(id_a));
assert_eq!(reg.get(id_a).unwrap().id(), id_a);
}
#[test]
fn registry_resolves_a_hash_only_peer_message() {
let writer = v0();
let bytes = encode(
&writer,
&Value::Struct(vec![(1, Value::str("checkout")), (4, Value::U16(900))]),
SchemaMode::HashOnly,
)
.unwrap();
let msg = Message::parse(&bytes).unwrap();
assert!(
!msg.has_inline_schema(),
"hash-only: nothing to read without the registry"
);
let mut reg = SchemaRegistry::new();
reg.register(writer);
let reader = v1();
let resolver = reg.resolver_for(msg.schema_id(), &reader).unwrap();
let root = msg.root(&resolver).unwrap();
assert_eq!(root.get_str(1).unwrap(), Some("checkout"));
assert_eq!(root.get_u32(4).unwrap(), Some(900), "u16 widened to u32");
assert_eq!(
root.get_str(3).unwrap(),
None,
"field added in v1, absent in v0 data"
);
assert!(reg.resolver_for(0xdead_beef, &reader).is_err());
}
#[test]
fn bundle_round_trips_and_is_deterministic() {
let mut reg = SchemaRegistry::new();
reg.register(v0());
reg.register(v1());
let bundle = reg.to_bundle();
let mut reg2 = SchemaRegistry::new();
reg2.register(v1());
reg2.register(v0());
assert_eq!(bundle, reg2.to_bundle());
let loaded = SchemaRegistry::from_bundle(&bundle).unwrap();
assert_eq!(loaded.len(), 2);
for id in reg.ids() {
assert_eq!(loaded.get(id).unwrap().id(), id);
}
}
#[test]
fn merge_reports_new_additions_and_skips_duplicates() {
let mut a = SchemaRegistry::new();
a.register(v0());
let mut b = SchemaRegistry::new();
b.register(v0()); b.register(v1()); let added = a.merge_bundle(&b.to_bundle()).unwrap();
assert_eq!(added, 1, "only v1 is new");
assert_eq!(a.len(), 2);
}
#[test]
fn rejects_corrupt_bundle() {
let mut reg = SchemaRegistry::new();
reg.register(v0());
let good = reg.to_bundle();
let mut bad_magic = good.clone();
bad_magic[0] = b'X';
assert!(SchemaRegistry::from_bundle(&bad_magic).is_err());
let mut bad_len = good.clone();
bad_len[8..12].copy_from_slice(&1u32.to_le_bytes()); bad_len[12..16].copy_from_slice(&u32::MAX.to_le_bytes());
assert!(SchemaRegistry::from_bundle(&bad_len).is_err());
assert!(
SchemaRegistry::from_bundle(&good[..6]).is_err(),
"truncated header"
);
}
}