diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 27295c12..2fa38290 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -6,6 +6,8 @@ ### New Features and Improvements +- Added `MessageRegistry::from_descriptors` (behind the `zeroparser` feature) for building a registry from a parse root plus sibling/imported top-level messages, each keyed by its package-qualified FQN. This lets the zeroparser decode messages whose root fields reference imported well-known types (e.g. `google.protobuf.Timestamp`, `Duration`, and the wrapper types) that `from_descriptor` alone does not register. `others` takes borrowed `&[(&str, &DescriptorProto)]` pairs. + ### Bug Fixes ### Documentation diff --git a/rust/sdk/src/zeroparser/registry.rs b/rust/sdk/src/zeroparser/registry.rs index 19e9deb7..39f2b5ed 100644 --- a/rust/sdk/src/zeroparser/registry.rs +++ b/rust/sdk/src/zeroparser/registry.rs @@ -185,15 +185,58 @@ impl MessageRegistry { } } + /// Build a registry from the parse root plus additional top-level messages. + /// + /// Use this when root fields reference sibling/imported messages that are + /// not nested under `root`. + #[inline(always)] + pub fn from_descriptors(root: &DescriptorProto, others: &[(&str, &DescriptorProto)]) -> Self { + let mut messages = HashMap::new(); + + // Register the root (and its nested types) at the root prefix, exactly + // as `from_descriptor` does, so root-relative lookups are unchanged. + Self::collect_messages(root, &mut messages); + + // Register every other top-level message under its package-qualified FQN. + for (package, desc) in others { + let prefix = if package.is_empty() { + ROOT_PREFIX.to_string() + } else { + format!(".{package}") + }; + Self::collect_messages_with_prefix(desc, prefix, &mut messages); + } + + let root_descriptor = DescriptorWithFieldCache::from_descriptor(root); + + MessageRegistry { + messages, + root_descriptor, + } + } + /// Iteratively collect all message descriptors using an explicit stack. /// This avoids stack overflow for deeply nested message types. #[inline(always)] fn collect_messages( root: &DescriptorProto, acc: &mut HashMap, + ) { + Self::collect_messages_with_prefix(root, ROOT_PREFIX.to_string(), acc); + } + + /// Collect `root` and all of its nested types into `acc`, threading + /// `initial_prefix` as the enclosing namespace. An empty prefix yields + /// `.Name` (root behaviour); a package prefix like `.google.protobuf` + /// yields `.google.protobuf.Name`. + #[inline(always)] + fn collect_messages_with_prefix( + root: &DescriptorProto, + initial_prefix: String, + acc: &mut HashMap, ) { // Use an explicit stack: (descriptor, prefix) - let mut stack: Vec<(&DescriptorProto, String)> = vec![(root, ROOT_PREFIX.to_string())]; + let mut stack: Vec<(&DescriptorProto, String)> = vec![(root, initial_prefix)]; while let Some((desc, current_prefix)) = stack.pop() { let name = desc.name.as_deref().unwrap_or(""); @@ -355,4 +398,87 @@ pub mod tests { assert_eq!(registry.get_field_type_name(2), Some(".Level1.Level2")); assert_eq!(registry.get_field_type_name(99), None); } + + #[test] + fn message_registry_from_descriptors_registers_siblings() { + // Root message references an imported well-known type by its FQN. + let root = make_descriptor( + "Bet", + vec![ + make_field(1, "id", field_descriptor_proto::Type::String, false, None), + make_field( + 2, + "note", + field_descriptor_proto::Type::Message, + false, + Some(".google.protobuf.StringValue"), + ), + ], + ); + + // A sibling/imported top-level message from another file + package. + let string_value = make_descriptor( + "StringValue", + vec![make_field( + 1, + "value", + field_descriptor_proto::Type::String, + false, + None, + )], + ); + + // `from_descriptor` alone never registers the sibling: lookup misses. + let root_only = MessageRegistry::from_descriptor(&root); + assert!(root_only.get(".google.protobuf.StringValue").is_none()); + + // `from_descriptors` registers it under its true package-qualified FQN. + let registry = + MessageRegistry::from_descriptors(&root, &[("google.protobuf", &string_value)]); + + assert!(registry.get(".Bet").is_some()); + let wrapper = registry + .get(".google.protobuf.StringValue") + .expect("sibling well-known type must be registered"); + assert_eq!(&*wrapper.get_field(1).unwrap().name, "value"); + + // Root remains the parse root, and the field still resolves to the FQN. + assert_eq!( + registry.get_field_type_name(2), + Some(".google.protobuf.StringValue") + ); + } + + #[test] + fn message_registry_from_descriptors_empty_package_and_nesting() { + // Empty-package sibling registers at the root prefix (`.Name`), and its + // nested types are collected too. + let mut sibling = make_descriptor( + "Sibling", + vec![make_field( + 1, + "inner", + field_descriptor_proto::Type::Message, + false, + Some(".Sibling.Inner"), + )], + ); + sibling.nested_type.push(make_descriptor( + "Inner", + vec![make_field( + 1, + "x", + field_descriptor_proto::Type::Int32, + false, + None, + )], + )); + + let root = make_descriptor("Root", vec![]); + let registry = MessageRegistry::from_descriptors(&root, &[("", &sibling)]); + + assert!(registry.get(".Root").is_some()); + assert!(registry.get(".Sibling").is_some()); + assert!(registry.get(".Sibling.Inner").is_some()); + } } diff --git a/rust/sdk/src/zeroparser/tests/common/mod.rs b/rust/sdk/src/zeroparser/tests/common/mod.rs index 97b8c954..e83ea830 100644 --- a/rust/sdk/src/zeroparser/tests/common/mod.rs +++ b/rust/sdk/src/zeroparser/tests/common/mod.rs @@ -60,14 +60,25 @@ pub fn get_message_descriptor( } pub fn create_registry_for_version(version: ProtoVersion, message_name: &str) -> MessageRegistry { - let file_desc_set = load_descriptor_set(version); + let file_desc_set = + FileDescriptorSet::decode(E2E_DESCRIPTOR_SET).expect("decode descriptor file"); let (msg_desc, file) = find_message_and_file(&file_desc_set, version.package(), message_name); let mut descriptor = msg_desc.clone(); descriptor.name = Some(format!( "{}.{message_name}", file.package.as_deref().unwrap_or("") )); - MessageRegistry::from_descriptor(&descriptor) + + let others: Vec<(&str, &DescriptorProto)> = file_desc_set + .file + .iter() + .flat_map(|f| { + let package = f.package.as_deref().unwrap_or(""); + f.message_type.iter().map(move |msg| (package, msg)) + }) + .collect(); + + MessageRegistry::from_descriptors(&descriptor, &others) } pub fn encode_message_for_version( diff --git a/rust/sdk/src/zeroparser/tests/e2e.rs b/rust/sdk/src/zeroparser/tests/e2e.rs index 14ea2713..65f7a63d 100644 --- a/rust/sdk/src/zeroparser/tests/e2e.rs +++ b/rust/sdk/src/zeroparser/tests/e2e.rs @@ -2925,13 +2925,10 @@ fn test_invalid_utf8_various_errors(#[case] version: ProtoVersion) { #[rstest] #[case(ProtoVersion::Proto2)] #[case(ProtoVersion::Proto3)] -fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) { +fn test_google_protobuf_types_decode_success(#[case] version: ProtoVersion) { use prost_types::{Duration, Timestamp}; - let registry = create_registry_for_version(version, "GoogleProtobufTypesMessage"); - - // Test google.protobuf types (Timestamp/Duration) - let (buf_timestamp, _) = match version { + let (buf, registry) = match version { ProtoVersion::Proto2 => { let msg = proto2::GoogleProtobufTypesMessage { timestamp: Some(Timestamp { @@ -2942,6 +2939,15 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) { seconds: 3600, nanos: 500000000, }), + string_value: Some("wrapped_string".to_string()), + int32_value: Some(-42), + int64_value: Some(-4242), + uint32_value: Some(42), + uint64_value: Some(4242), + float_value: Some(3.5), + double_value: Some(7.25), + bool_value: Some(true), + bytes_value: Some(b"wrapped_bytes".to_vec()), timestamps: vec![ Timestamp { seconds: 1000000000, @@ -2952,7 +2958,6 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) { nanos: 999999999, }, ], - ..Default::default() }; encode_message_for_version(version, &msg, "GoogleProtobufTypesMessage") } @@ -2966,6 +2971,15 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) { seconds: 3600, nanos: 500000000, }), + string_value: Some("wrapped_string".to_string()), + int32_value: Some(-42), + int64_value: Some(-4242), + uint32_value: Some(42), + uint64_value: Some(4242), + float_value: Some(3.5), + double_value: Some(7.25), + bool_value: Some(true), + bytes_value: Some(b"wrapped_bytes".to_vec()), timestamps: vec![ Timestamp { seconds: 1000000000, @@ -2976,37 +2990,81 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) { nanos: 999999999, }, ], - ..Default::default() }; encode_message_for_version(version, &msg, "GoogleProtobufTypesMessage") } }; - let result = ParsedMessage::parse(&buf_timestamp, ®istry); - assert!( - matches!(result, Err(ParseError::UnknownTypeName { .. })), - "Expected UnknownTypeName error for google.protobuf.Timestamp, got {:?}", - result - ); + let parsed = ParsedMessage::parse(&buf, ®istry).unwrap(); - // Test wrapper types (manually encoded) - let mut buf_wrappers = Vec::new(); + let timestamp = parsed.get_message(1).unwrap(); + assert_eq!( + timestamp.get_scalar(1), + Some(&FieldValueRef::Int64(1234567890)) + ); + assert_eq!( + timestamp.get_scalar(2), + Some(&FieldValueRef::Int32(123456789)) + ); - // StringValue (field 3) - message with field 1 = string - let mut string_value_buf = Vec::new(); - prost::encoding::string::encode(1, &"wrapped_string".to_string(), &mut string_value_buf); - prost::encoding::message::encode(3, &string_value_buf, &mut buf_wrappers); + let duration = parsed.get_message(2).unwrap(); + assert_eq!(duration.get_scalar(1), Some(&FieldValueRef::Int64(3600))); + assert_eq!( + duration.get_scalar(2), + Some(&FieldValueRef::Int32(500000000)) + ); - // Int32Value (field 4) - message with field 1 = int32 - let mut int32_value_buf = Vec::new(); - prost::encoding::int32::encode(1, &-42, &mut int32_value_buf); - prost::encoding::message::encode(4, &int32_value_buf, &mut buf_wrappers); + assert_eq!( + parsed.get_message(3).unwrap().get_scalar(1), + Some(&FieldValueRef::String("wrapped_string")) + ); + assert_eq!( + parsed.get_message(4).unwrap().get_scalar(1), + Some(&FieldValueRef::Int32(-42)) + ); + assert_eq!( + parsed.get_message(5).unwrap().get_scalar(1), + Some(&FieldValueRef::Int64(-4242)) + ); + assert_eq!( + parsed.get_message(6).unwrap().get_scalar(1), + Some(&FieldValueRef::UInt32(42)) + ); + assert_eq!( + parsed.get_message(7).unwrap().get_scalar(1), + Some(&FieldValueRef::UInt64(4242)) + ); + assert_eq!( + parsed.get_message(8).unwrap().get_scalar(1), + Some(&FieldValueRef::Float(3.5)) + ); + assert_eq!( + parsed.get_message(9).unwrap().get_scalar(1), + Some(&FieldValueRef::Double(7.25)) + ); + assert_eq!( + parsed.get_message(10).unwrap().get_scalar(1), + Some(&FieldValueRef::Bool(true)) + ); + assert_eq!( + parsed.get_message(11).unwrap().get_scalar(1), + Some(&FieldValueRef::Bytes(b"wrapped_bytes")) + ); - let result = ParsedMessage::parse(&buf_wrappers, ®istry); - assert!( - matches!(result, Err(ParseError::UnknownTypeName { .. })), - "Expected UnknownTypeName error for google.protobuf wrapper types, got {:?}", - result + let timestamps = parsed.get_repeated_messages(12); + assert_eq!(timestamps.len(), 2); + assert_eq!( + timestamps[0].get_scalar(1), + Some(&FieldValueRef::Int64(1000000000)) + ); + assert_eq!(timestamps[0].get_scalar(2), None); + assert_eq!( + timestamps[1].get_scalar(1), + Some(&FieldValueRef::Int64(2000000000)) + ); + assert_eq!( + timestamps[1].get_scalar(2), + Some(&FieldValueRef::Int32(999999999)) ); }