harden proto storage — recovery, atomicity, persistence

crates/ember-core/src/engine.rs

@@ -71,6 +71,8 @@ impl Engine {
                     shard_config,
                     config.persistence.clone(),
                     Some(drop_handle.clone()),
+                    #[cfg(feature = "protobuf")]
+                    config.schema_registry.clone(),
                 )
             })
             .collect();

crates/ember-core/src/memory.rs

@@ -223,9 +223,10 @@ pub fn value_size(value: &Value) -> usize {
             let member_bytes: usize = set.iter().map(|m| m.len() + HASHSET_MEMBER_OVERHEAD).sum();
             HASHSET_BASE_OVERHEAD + member_bytes
         }
-        // type_name String (24 bytes ptr+len+cap on heap) + data Bytes (24 bytes).
+        // type_name: String struct = 24 bytes (ptr+len+cap) on 64-bit.
+        // data: Bytes struct = ~24 bytes (ptr+len+vtable/arc).
         #[cfg(feature = "protobuf")]
-        Value::Proto { type_name, data } => type_name.len() + data.len() + 24,
+        Value::Proto { type_name, data } => type_name.len() + data.len() + 48,
     }
 }
 

crates/ember-core/src/schema.rs

@@ -25,6 +25,14 @@ const MAX_DESCRIPTOR_BYTES: usize = 10 * 1024 * 1024;
 /// Deep nesting beyond this is almost certainly a bug or an abuse vector.
 const MAX_FIELD_PATH_DEPTH: usize = 16;
 
+/// Maximum allowed size for a single proto value in bytes (64 MB).
+/// Prevents a single PROTO.SET from exhausting shard memory.
+const MAX_PROTO_VALUE_BYTES: usize = 64 * 1024 * 1024;
+
+/// Maximum number of schemas that can be registered.
+/// Prevents unbounded growth of the schema registry.
+const MAX_SCHEMAS: usize = 1024;
+
 /// Errors that can occur during schema operations.
 #[derive(Debug, Error)]
 pub enum SchemaError {
@@ -48,6 +56,12 @@ pub enum SchemaError {
 
     #[error("field path too deep: {0} segments (max {1})")]
     PathTooDeep(usize, usize),
+
+    #[error("proto value too large: {0} bytes (max {1})")]
+    ValueTooLarge(usize, usize),
+
+    #[error("schema limit reached: {0} schemas (max {1})")]
+    TooManySchemas(usize, usize),
 }
 
 /// A registered schema: the raw descriptor bytes and the parsed pool.
@@ -106,6 +120,10 @@ impl SchemaRegistry {
             return Err(SchemaError::AlreadyExists(name));
         }
 
+        if self.schemas.len() >= MAX_SCHEMAS {
+            return Err(SchemaError::TooManySchemas(self.schemas.len(), MAX_SCHEMAS));
+        }
+
         if descriptor_bytes.len() > MAX_DESCRIPTOR_BYTES {
             return Err(SchemaError::DescriptorTooLarge(
                 descriptor_bytes.len(),
@@ -145,8 +163,16 @@ impl SchemaRegistry {
 
     /// Validates that `data` is a valid encoding of `message_type`.
     ///
+    /// Checks the value size limit before decoding.
     /// Searches all registered schemas for the type name.
     pub fn validate(&self, message_type: &str, data: &[u8]) -> Result<(), SchemaError> {
+        if data.len() > MAX_PROTO_VALUE_BYTES {
+            return Err(SchemaError::ValueTooLarge(
+                data.len(),
+                MAX_PROTO_VALUE_BYTES,
+            ));
+        }
+
         let descriptor = self.find_message(message_type)?;
 
         DynamicMessage::decode(descriptor, data)
@@ -1364,4 +1390,52 @@ mod tests {
         let frame = registry.get_field("test.BigNum", &buf, "val").unwrap();
         assert_eq!(frame, Frame::Integer(42));
     }
+
+    // --- value size limit / schema count limit tests ---
+
+    #[test]
+    fn value_too_large_rejected() {
+        let mut registry = SchemaRegistry::new();
+        let desc = make_descriptor("test", "User", "name");
+        registry.register("users".into(), desc).unwrap();
+
+        let oversized = vec![0u8; MAX_PROTO_VALUE_BYTES + 1];
+        let err = registry.validate("test.User", &oversized).unwrap_err();
+        assert!(matches!(err, SchemaError::ValueTooLarge(_, _)));
+    }
+
+    #[test]
+    fn value_at_limit_allowed() {
+        let mut registry = SchemaRegistry::new();
+        let desc = make_descriptor("test", "User", "name");
+        registry.register("users".into(), desc).unwrap();
+
+        // a value exactly at the limit should pass the size check
+        // (it will fail validation since it's not valid protobuf, but
+        // it should NOT fail with ValueTooLarge)
+        let at_limit = vec![0u8; MAX_PROTO_VALUE_BYTES];
+        let err = registry.validate("test.User", &at_limit).unwrap_err();
+        assert!(
+            !matches!(err, SchemaError::ValueTooLarge(_, _)),
+            "expected validation error, not size limit"
+        );
+    }
+
+    #[test]
+    fn schema_count_limit() {
+        let mut registry = SchemaRegistry::new();
+
+        // register up to the limit
+        for i in 0..MAX_SCHEMAS {
+            let desc = make_descriptor(&format!("pkg{i}"), &format!("Msg{i}"), "val");
+            registry
+                .register(format!("schema-{i}"), desc)
+                .unwrap_or_else(|e| panic!("failed to register schema {i}: {e}"));
+        }
+
+        // the next one should fail
+        let desc = make_descriptor("overflow", "Overflow", "val");
+        let err = registry.register("overflow".into(), desc).unwrap_err();
+        assert!(matches!(err, SchemaError::TooManySchemas(_, _)));
+    }
 }