fix(ber,jer): handle extension groups correctly

- Adjust BER/JER encode/decode logic for extension groups
- Update extension group tests

Author: Boslx

src/ber.rs

@@ -8,6 +8,22 @@ mod rules;
 pub use identifier::Identifier;
 pub(crate) use rules::EncodingRules;
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) enum ExtensionGroupState {
+    None,
+    Pending(crate::types::Tag),
+    Active(crate::types::Tag),
+}
+
+impl ExtensionGroupState {
+    pub(crate) fn base_tag(self) -> Option<crate::types::Tag> {
+        match self {
+            Self::Pending(tag) | Self::Active(tag) => Some(tag),
+            Self::None => None,
+        }
+    }
+}
+
 /// Attempts to decode `T` from `input` using BER.
 /// # Errors
 /// Returns error specific to BER decoder if decoding is not possible.

src/ber/de.rs

@@ -3,7 +3,7 @@
 mod config;
 pub(super) mod parser;
 
-use super::identifier::Identifier;
+use super::{identifier::Identifier, ExtensionGroupState};
 use crate::{
     types::{
         self,
@@ -29,6 +29,7 @@ pub struct Decoder<'input> {
     input: &'input [u8],
     config: DecoderOptions,
     initial_len: usize,
+    extension_group: ExtensionGroupState,
 }
 
 impl<'input> Decoder<'input> {
@@ -49,9 +50,14 @@ impl<'input> Decoder<'input> {
             input,
             config,
             initial_len: input.len(),
+            extension_group: ExtensionGroupState::None,
         }
     }
 
+    fn translate_tag(&self, tag: Tag) -> Tag {
+        tag.with_context_offset(self.extension_group.base_tag())
+    }
+
     /// Return a number of the decoded bytes by this decoder
     #[must_use]
     pub fn decoded_len(&self) -> usize {
@@ -85,6 +91,9 @@ impl<'input> Decoder<'input> {
         {
             return Ok(None);
         }
+
+        let tag = self.translate_tag(tag);
+
         if tag != Tag::EOC {
             let upcoming_tag = self.peek_tag()?;
             if tag != upcoming_tag {
@@ -102,13 +111,15 @@ impl<'input> Decoder<'input> {
     }
 
     pub(crate) fn parse_value(&mut self, tag: Tag) -> Result<(Identifier, Option<&'input [u8]>)> {
+        let tag = self.translate_tag(tag);
         let (input, (identifier, contents)) =
             self::parser::parse_value(self.config, self.input, Some(tag))?;
         self.input = input;
         Ok((identifier, contents))
     }
 
     pub(crate) fn parse_primitive_value(&mut self, tag: Tag) -> Result<(Identifier, &'input [u8])> {
+        let tag = self.translate_tag(tag);
         let (input, (identifier, contents)) =
             self::parser::parse_value(self.config, self.input, Some(tag))?;
         self.input = input;
@@ -762,6 +773,28 @@ impl<'input> crate::Decoder for Decoder<'input> {
         default_initializer_fn: Option<DF>,
         decode_fn: F,
     ) -> Result<D> {
+        if tag == Tag::SEQUENCE && matches!(self.extension_group, ExtensionGroupState::Pending(_)) {
+            // Extension addition groups are encoded flattened: skip the SEQUENCE wrapper once.
+            if let ExtensionGroupState::Pending(tag) = self.extension_group {
+                self.extension_group = ExtensionGroupState::Active(tag);
+            }
+            return if D::FIELDS.is_empty() && D::EXTENDED_FIELDS.is_none()
+                || (D::FIELDS.len() == D::FIELDS.number_of_optional_and_default_fields()
+                    && self.input.is_empty())
+            {
+                if let Some(default_initializer_fn) = default_initializer_fn {
+                    Ok((default_initializer_fn)())
+                } else {
+                    Err(DecodeError::from_kind(
+                        DecodeErrorKind::UnexpectedEmptyInput,
+                        self.codec(),
+                    ))
+                }
+            } else {
+                (decode_fn)(self)
+            };
+        }
+
         self.parse_constructed_contents(tag, true, |decoder| {
             // If there are no fields, or the input is empty and we know that
             // all fields are optional or default fields, we call the default
@@ -806,7 +839,7 @@ impl<'input> crate::Decoder for Decoder<'input> {
         D: Fn(&mut Self, usize, Tag) -> Result<FIELDS, Self::Error>,
         F: FnOnce(Vec<FIELDS>) -> Result<SET, Self::Error>,
     {
-        self.parse_constructed_contents(tag, true, |decoder| {
+        let collect_fields = |decoder: &mut Self| -> Result<Vec<FIELDS>, Self::Error> {
             let mut fields = Vec::new();
 
             loop {
@@ -823,6 +856,20 @@ impl<'input> crate::Decoder for Decoder<'input> {
                 }
             }
 
+            Ok(fields)
+        };
+
+        if tag == Tag::SET && matches!(self.extension_group, ExtensionGroupState::Pending(_)) {
+            // Extension addition groups are encoded flattened: skip the SET wrapper once.
+            if let ExtensionGroupState::Pending(tag) = self.extension_group {
+                self.extension_group = ExtensionGroupState::Active(tag);
+            }
+            let fields = collect_fields(self)?;
+            return (field_fn)(fields);
+        }
+
+        self.parse_constructed_contents(tag, true, |decoder| {
+            let fields = collect_fields(decoder)?;
             (field_fn)(fields)
         })
     }
@@ -897,9 +944,26 @@ impl<'input> crate::Decoder for Decoder<'input> {
         D: Decode + crate::types::Constructed<RL, EL>,
     >(
         &mut self,
-        _tag: Tag,
+        tag: Tag,
     ) -> Result<Option<D>, Self::Error> {
-        <Option<D>>::decode(self)
+        if self.input.is_empty() {
+            return Ok(None);
+        }
+
+        let (_, identifier) = parser::parse_identifier_octet(self.input).map_err(|e| match e {
+            ParseNumberError::Nom(e) => DecodeError::map_nom_err(e, self.codec()),
+            ParseNumberError::Overflow => DecodeError::integer_overflow(32u32, self.codec()),
+        })?;
+
+        if identifier.tag == tag {
+            let previous = self.extension_group;
+            self.extension_group = ExtensionGroupState::Pending(tag);
+            let result = D::decode(self).map(Some);
+            self.extension_group = previous;
+            result
+        } else {
+            Ok(None)
+        }
     }
 }
 

src/ber/enc.rs

@@ -5,7 +5,7 @@ mod config;
 use alloc::{borrow::ToOwned, collections::VecDeque, string::ToString, vec::Vec};
 use chrono::Timelike;
 
-use super::Identifier;
+use super::{ExtensionGroupState, Identifier};
 use crate::{
     bits::octet_string_ascending,
     types::{
@@ -28,6 +28,7 @@ pub struct Encoder {
     config: EncoderOptions,
     is_set_encoding: bool,
     set_buffer: alloc::collections::BTreeMap<Tag, Vec<u8>>,
+    extension_group: ExtensionGroupState,
 }
 
 /// A convenience type around results needing to return one or many bytes.
@@ -45,6 +46,7 @@ impl Encoder {
             is_set_encoding: false,
             output: <_>::default(),
             set_buffer: <_>::default(),
+            extension_group: ExtensionGroupState::None,
         }
     }
 
@@ -63,6 +65,7 @@ impl Encoder {
             is_set_encoding: true,
             output: <_>::default(),
             set_buffer: <_>::default(),
+            extension_group: ExtensionGroupState::None,
         }
     }
 
@@ -78,9 +81,14 @@ impl Encoder {
             config,
             is_set_encoding: false,
             set_buffer: <_>::default(),
+            extension_group: ExtensionGroupState::None,
         }
     }
 
+    fn translate_tag(&self, tag: Tag) -> Tag {
+        tag.with_context_offset(self.extension_group.base_tag())
+    }
+
     /// Consumes the encoder and returns the output of the encoding.
     #[must_use]
     pub fn output(self) -> Vec<u8> {
@@ -243,6 +251,10 @@ impl Encoder {
 
     /// Encodes a given ASN.1 BER value with the `identifier`.
     fn encode_value(&mut self, identifier: Identifier, value: &[u8]) {
+        let identifier = Identifier::from_tag(
+            self.translate_tag(identifier.tag),
+            identifier.is_constructed,
+        );
         let ident_bytes = self.encode_identifier(identifier);
         self.append_byte_or_bytes(ident_bytes);
         self.encode_length(identifier, value);
@@ -711,6 +723,14 @@ impl crate::Encoder<'_> for Encoder {
         C: crate::types::Constructed<RC, EC>,
         F: FnOnce(&mut Self::AnyEncoder<'b, 0, 0>) -> Result<(), Self::Error>,
     {
+        if tag == Tag::SEQUENCE && matches!(self.extension_group, ExtensionGroupState::Pending(_)) {
+            // Extension addition groups are encoded flattened: skip the SEQUENCE wrapper once.
+            if let ExtensionGroupState::Pending(tag) = self.extension_group {
+                self.extension_group = ExtensionGroupState::Active(tag);
+            }
+            return (encoder_scope)(self);
+        }
+
         let mut encoder = Self::new(self.config);
 
         (encoder_scope)(&mut encoder)?;
@@ -730,6 +750,14 @@ impl crate::Encoder<'_> for Encoder {
         C: crate::types::Constructed<RC, EC>,
         F: FnOnce(&mut Self::AnyEncoder<'b, 0, 0>) -> Result<(), Self::Error>,
     {
+        if tag == Tag::SET && matches!(self.extension_group, ExtensionGroupState::Pending(_)) {
+            // Extension addition groups are encoded flattened: skip the SET wrapper once.
+            if let ExtensionGroupState::Pending(tag) = self.extension_group {
+                self.extension_group = ExtensionGroupState::Active(tag);
+            }
+            return (encoder_scope)(self);
+        }
+
         let mut encoder = Self::new_set(self.config);
 
         (encoder_scope)(&mut encoder)?;
@@ -757,14 +785,23 @@ impl crate::Encoder<'_> for Encoder {
     /// Encode a extension addition group value.
     fn encode_extension_addition_group<const RC: usize, const EC: usize, E>(
         &mut self,
-        _tag: Tag,
+        tag: Tag,
         value: Option<&E>,
         _: crate::types::Identifier,
     ) -> Result<Self::Ok, Self::Error>
     where
         E: Encode + crate::types::Constructed<RC, EC>,
     {
-        value.encode(self)
+        match value {
+            Some(v) => {
+                let previous = self.extension_group;
+                self.extension_group = ExtensionGroupState::Pending(tag);
+                let result = v.encode(self);
+                self.extension_group = previous;
+                result
+            }
+            None => Ok(()),
+        }
     }
 }
 

src/jer/de.rs

@@ -179,7 +179,7 @@ impl crate::Decoder for Decoder {
         F: FnOnce(&mut Self) -> Result<D, Self::Error>,
     {
         let mut last = self.stack.pop().ok_or_else(JerDecodeErrorKind::eoi)?;
-        let value_map = last
+        let _ = last
             .as_object_mut()
             .ok_or_else(|| JerDecodeErrorKind::TypeMismatch {
                 needed: "object",
@@ -193,12 +193,24 @@ impl crate::Decoder for Decoder {
             field_names.extend(extended_fields.iter().map(|f| f.name));
         }
         field_names.reverse();
+        // Push the (now partially consumed) object onto the stack so extension-addition-group
+        // decoding can pull group fields from the same flattened object.
+        self.stack.push(last);
+        let scope_index = self.stack.len() - 1;
         for name in field_names {
-            self.stack
-                .push(value_map.remove(name).unwrap_or(Value::Null));
+            let value = self
+                .stack
+                .get_mut(scope_index)
+                .and_then(|v| v.as_object_mut())
+                .and_then(|obj| obj.remove(name))
+                .unwrap_or(Value::Null);
+            self.stack.push(value);
         }
 
-        (decode_fn)(self)
+        let result = (decode_fn)(self);
+        // Pop the scope object frame.
+        let _ = self.stack.pop();
+        result
     }
 
     fn decode_sequence_of<D: crate::Decode>(
@@ -478,7 +490,51 @@ impl crate::Decoder for Decoder {
         &mut self,
         _tag: Tag,
     ) -> Result<Option<D>, Self::Error> {
-        self.decode_optional()
+        // The SEQUENCE decoder pushes a placeholder for the extension group field (which is not
+        // explicitly present in JER because extension groups are flattened).
+        //
+        // We decode a group by extracting only the group's fields from the current object and
+        // decoding the group from that scoped object.
+        let _ = self.stack.pop().ok_or_else(JerDecodeErrorKind::eoi)?;
+
+        let index = self
+            .stack
+            .iter()
+            .rposition(|v| v.is_object())
+            .ok_or_else(JerDecodeErrorKind::eoi)?;
+        let obj = self
+            .stack
+            .get_mut(index)
+            .and_then(|v| v.as_object_mut())
+            .ok_or_else(|| JerDecodeErrorKind::TypeMismatch {
+                needed: "object",
+                found: "unknown".into(),
+            })?;
+
+        let mut group_obj = serde_json::Map::with_capacity(
+            D::FIELDS.len() + D::EXTENDED_FIELDS.as_ref().map_or(0, |fields| fields.len()),
+        );
+        let mut is_present = false;
+        for field in D::FIELDS.iter() {
+            if let Some(value) = obj.remove(field.name) {
+                is_present |= !value.is_null();
+                group_obj.insert(alloc::string::String::from(field.name), value);
+            }
+        }
+        if let Some(extended_fields) = D::EXTENDED_FIELDS {
+            for field in extended_fields.iter() {
+                if let Some(value) = obj.remove(field.name) {
+                    is_present |= !value.is_null();
+                    group_obj.insert(alloc::string::String::from(field.name), value);
+                }
+            }
+        }
+        if is_present {
+            self.stack.push(Value::Object(group_obj));
+            D::decode(self).map(Some)
+        } else {
+            Ok(None)
+        }
     }
 
     fn codec(&self) -> crate::Codec {

src/jer/enc.rs

@@ -536,9 +536,22 @@ impl crate::Encoder<'_> for Encoder {
     where
         E: crate::Encode + crate::types::Constructed<RL, EL>,
     {
+        self.stack.pop();
         match value {
-            Some(v) => v.encode(self),
-            None => self.encode_none::<E>(Identifier::EMPTY),
+            Some(v) => {
+                let mut inner = Self::new();
+                v.encode(&mut inner)?;
+                if let Value::Object(obj) = inner.to_json()? {
+                    self.constructed_stack
+                        .last_mut()
+                        .ok_or_else(|| JerEncodeErrorKind::JsonEncoder {
+                            msg: "Internal stack mismatch!".into(),
+                        })?
+                        .extend(obj);
+                }
+                Ok(())
+            }
+            None => Ok(()),
         }
     }