verit_core/

registry.rs

//! Schema registry + distribution bundle for multi-service deployments.
//!
//! In a polyglot backend, a producer may send **hash-only** messages (just the
//! 128-bit schema id on the wire, no inline schema — the compact, fast form). A
//! consumer that receives one needs the *writer's* schema to resolve it against
//! its own reader schema. The registry is where writer schemas live, keyed by
//! their id.
//!
//! The design leans on the one fact that makes a Veritate registry simpler than
//! a conventional one: **the id is the content hash of the canonical schema**.
//! So the store is *content-addressed* — registration is idempotent, ids never
//! collide with different content, and there is no "which version is v3?"
//! question. A [`Schema`] carries its own authoritative id, so nothing can be
//! registered under a claimed-but-wrong id; [`from_bundle`](SchemaRegistry::from_bundle)
//! recomputes every id on import by decoding the canonical bytes.
//!
//! [`SchemaRegistry::to_bundle`] / [`from_bundle`](SchemaRegistry::from_bundle)
//! are the **distribution protocol**: a portable `VRSB` blob carrying a set of
//! schemas that one service publishes and others load — a schema-set snapshot
//! for a deployment.

use std::collections::HashMap;

use crate::error::{Error, Result};
use crate::resolve::Resolver;
use crate::schema::Schema;

/// Distribution-bundle magic: "VRSB" (Veritate Schema Bundle), version 1.
pub const BUNDLE_MAGIC: &[u8; 4] = b"VRSB";
/// Bundle format version this build reads and writes.
pub const BUNDLE_VERSION: u8 = 1;
const BUNDLE_HEADER_LEN: usize = 12;

/// A content-addressed store of schemas, keyed by their 128-bit id.
#[derive(Clone, Debug, Default)]
pub struct SchemaRegistry {
    by_id: HashMap<u128, Schema>,
}

impl SchemaRegistry {
    pub fn new() -> SchemaRegistry {
        SchemaRegistry {
            by_id: HashMap::new(),
        }
    }

    /// Register a schema, returning its id. Idempotent: registering the same
    /// content twice is a no-op (same id, same bytes — content-addressed).
    pub fn register(&mut self, schema: Schema) -> u128 {
        let id = schema.id();
        self.by_id.entry(id).or_insert(schema);
        id
    }

    /// Register from canonical bytes (e.g. a message's inline schema, or a
    /// bundle entry): decodes and validates them first, so a malformed or
    /// non-canonical blob is rejected rather than stored.
    pub fn register_canonical(&mut self, bytes: &[u8]) -> Result<u128> {
        Ok(self.register(Schema::from_canonical(bytes)?))
    }

    /// The schema with this id, if registered.
    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()
    }

    /// The ids of every registered schema.
    pub fn ids(&self) -> impl Iterator<Item = u128> + '_ {
        self.by_id.keys().copied()
    }

    /// Build a [`Resolver`] to read a message written with `writer_id` into
    /// `reader` — the multi-service payoff. Look up the writer schema here (it
    /// must be registered), then resolve it against the consumer's own reader
    /// schema, so a hash-only message from a peer becomes fully readable with
    /// schema evolution intact.
    ///
    /// `writer_id` is typically `msg.schema_id()`. Errors with
    /// [`Error::SchemaIdMismatch`] if the writer schema is not registered.
    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)
    }

    /// Serialize every registered schema into a portable distribution bundle.
    /// Entries are ordered by id, so the bytes are deterministic (diff-friendly)
    /// and independent of insertion order.
    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]); // reserved
        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
    }

    /// Load a registry from a distribution bundle. Every schema is decoded and
    /// validated (and thereby its id recomputed from its bytes), so a tampered
    /// bundle is rejected rather than trusted.
    pub fn from_bundle(bytes: &[u8]) -> Result<SchemaRegistry> {
        let mut reg = SchemaRegistry::new();
        reg.merge_bundle(bytes)?;
        Ok(reg)
    }

    /// Merge a distribution bundle into this registry, returning how many
    /// schemas were newly added (already-present ids are skipped — idempotent).
    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), // widened from u16
                ],
            )
            .build("LogEvent")
            .unwrap()
    }

    #[test]
    fn register_is_content_addressed_and_idempotent() {
        let mut reg = SchemaRegistry::new();
        let id_a = reg.register(v0());
        // Re-registering the same content is a no-op and yields the same id.
        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() {
        // A producer writes v0, hash-only (no inline schema on the wire).
        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"
        );

        // The consumer knows only v1, but has the producer's schema in its 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"
        );

        // An unknown writer id is a typed error, not a panic.
        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();
        // Deterministic regardless of insertion order.
        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()); // duplicate
        b.register(v1()); // new
        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();
        // Corrupt the first entry's length prefix to point past the buffer.
        bad_len[8..12].copy_from_slice(&1u32.to_le_bytes()); // count stays 1
        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"
        );
    }
}