asn1: refactor expected tag logic

Signed-off-by: Facundo Tuesca <facundo.tuesca@trailofbits.com>

Author: facutuesca

src/rust/src/declarative_asn1/decode.rs

@@ -7,7 +7,7 @@ use pyo3::types::{PyAnyMethods, PyListMethods};
 
 use crate::asn1::big_byte_slice_to_py_int;
 use crate::declarative_asn1::types::{
-    check_size_constraint, type_to_tag, AnnotatedType, Annotation, BitString, Encoding,
+    check_size_constraint, is_tag_valid_for_type, AnnotatedType, Annotation, BitString, Encoding,
     GeneralizedTime, IA5String, PrintableString, Type, UtcTime,
 };
 use crate::error::CryptographyError;
@@ -172,10 +172,9 @@ pub(crate) fn decode_annotated_type<'a>(
     // Handle DEFAULT annotation if field is not present (by
     // returning the default value)
     if let Some(default) = &ann_type.annotation.get().default {
-        let expected_tag = type_to_tag(inner, encoding);
-        let next_tag = parser.peek_tag();
-        if next_tag != Some(expected_tag) {
-            return Ok(default.clone_ref(py).into_bound(py));
+        match parser.peek_tag() {
+            Some(next_tag) if is_tag_valid_for_type(next_tag, inner, encoding) => (),
+            _ => return Ok(default.clone_ref(py).into_bound(py)),
         }
     }
 
@@ -210,9 +209,8 @@ pub(crate) fn decode_annotated_type<'a>(
             })?
         }
         Type::Option(cls) => {
-            let inner_tag = type_to_tag(cls.get().inner.get(), encoding);
             match parser.peek_tag() {
-                Some(t) if t == inner_tag => {
+                Some(t) if is_tag_valid_for_type(t, cls.get().inner.get(), encoding) => {
                     // For optional types, annotations will always be associated to the `Optional` type
                     // i.e: `Annotated[Optional[T], annotation]`, as opposed to the inner `T` type.
                     // Therefore, when decoding the inner type `T` we must pass the annotation of the `Optional`

src/rust/src/declarative_asn1/types.rs

@@ -419,29 +419,52 @@ pub(crate) fn python_class_to_annotated<'p>(
     }
 }
 
-pub(crate) fn type_to_tag(t: &Type, encoding: &Option<pyo3::Py<Encoding>>) -> asn1::Tag {
-    let inner_tag = match t {
-        Type::Sequence(_, _) => asn1::Sequence::TAG,
-        Type::SequenceOf(_) => asn1::Sequence::TAG,
-        Type::Option(t) => type_to_tag(t.get().inner.get(), encoding),
-        Type::PyBool() => bool::TAG,
-        Type::PyInt() => asn1::BigInt::TAG,
-        Type::PyBytes() => <&[u8] as SimpleAsn1Readable>::TAG,
-        Type::PyStr() => asn1::Utf8String::TAG,
-        Type::PrintableString() => asn1::PrintableString::TAG,
-        Type::IA5String() => asn1::IA5String::TAG,
-        Type::ObjectIdentifier() => asn1::ObjectIdentifier::TAG,
-        Type::UtcTime() => asn1::UtcTime::TAG,
-        Type::GeneralizedTime() => asn1::GeneralizedTime::TAG,
-        Type::BitString() => asn1::BitString::TAG,
-    };
-
-    match encoding {
+// Checks if encoding `tag_without_encoding` using `encoding` results
+// in `tag`
+fn check_tag_with_encoding(
+    tag_without_encoding: asn1::Tag,
+    encoding: &Option<pyo3::Py<Encoding>>,
+    tag: asn1::Tag,
+) -> bool {
+    let tag_with_encoding = match encoding {
         Some(e) => match e.get() {
-            Encoding::Implicit(n) => asn1::implicit_tag(*n, inner_tag),
+            Encoding::Implicit(n) => asn1::implicit_tag(*n, tag_without_encoding),
             Encoding::Explicit(n) => asn1::explicit_tag(*n),
         },
-        None => inner_tag,
+        None => tag_without_encoding,
+    };
+    tag_with_encoding == tag
+}
+
+// Given `tag` and `encoding`, returns whether that tag with that encoding
+// matches what one would expect to see when decoding `type_`
+pub(crate) fn is_tag_valid_for_type(
+    tag: asn1::Tag,
+    type_: &Type,
+    encoding: &Option<pyo3::Py<Encoding>>,
+) -> bool {
+    match type_ {
+        Type::Sequence(_, _) => check_tag_with_encoding(asn1::Sequence::TAG, encoding, tag),
+        Type::SequenceOf(_) => check_tag_with_encoding(asn1::Sequence::TAG, encoding, tag),
+        Type::Option(t) => is_tag_valid_for_type(tag, t.get().inner.get(), encoding),
+        Type::PyBool() => check_tag_with_encoding(bool::TAG, encoding, tag),
+        Type::PyInt() => check_tag_with_encoding(asn1::BigInt::TAG, encoding, tag),
+        Type::PyBytes() => {
+            check_tag_with_encoding(<&[u8] as SimpleAsn1Readable>::TAG, encoding, tag)
+        }
+        Type::PyStr() => check_tag_with_encoding(asn1::Utf8String::TAG, encoding, tag),
+        Type::PrintableString() => {
+            check_tag_with_encoding(asn1::PrintableString::TAG, encoding, tag)
+        }
+        Type::IA5String() => check_tag_with_encoding(asn1::IA5String::TAG, encoding, tag),
+        Type::ObjectIdentifier() => {
+            check_tag_with_encoding(asn1::ObjectIdentifier::TAG, encoding, tag)
+        }
+        Type::UtcTime() => check_tag_with_encoding(asn1::UtcTime::TAG, encoding, tag),
+        Type::GeneralizedTime() => {
+            check_tag_with_encoding(asn1::GeneralizedTime::TAG, encoding, tag)
+        }
+        Type::BitString() => check_tag_with_encoding(asn1::BitString::TAG, encoding, tag),
     }
 }
 
@@ -468,14 +491,15 @@ pub(crate) fn check_size_constraint(
 #[cfg(test)]
 mod tests {
 
+    use asn1::SimpleAsn1Readable;
     use pyo3::IntoPyObject;
 
-    use super::{type_to_tag, AnnotatedType, Annotation, Type};
+    use super::{is_tag_valid_for_type, AnnotatedType, Annotation, Type};
 
     #[test]
-    // Needed for coverage of `type_to_tag(Type::Option(..))`, since
-    // `type_to_tag` is never called with an optional value.
-    fn test_option_type_to_tag() {
+    // Needed for coverage of `is_tag_valid_for_type(Type::Option(..))`, since
+    // `is_tag_valid_for_type` is never called with an optional value.
+    fn test_option_is_tag_valid_for_type() {
         pyo3::Python::initialize();
 
         pyo3::Python::attach(|py| {
@@ -509,8 +533,11 @@ mod tests {
                 },
             )
             .unwrap();
-            let expected_tag = type_to_tag(&Type::Option(optional_type), &None);
-            assert_eq!(expected_tag, type_to_tag(&Type::PyInt(), &None))
+            assert!(is_tag_valid_for_type(
+                asn1::BigInt::TAG,
+                &Type::Option(optional_type),
+                &None
+            ));
         })
     }
 }

tests/hazmat/asn1/test_serialization.py

@@ -467,6 +467,10 @@ class Example:
             h: typing.Union[asn1.BitString, None]
             i: typing.Union[asn1.IA5String, None]
             j: typing.Union[x509.ObjectIdentifier, None]
+            k: Annotated[typing.Union[str, None], asn1.Implicit(0)]
+            only_field_present: Annotated[
+                typing.Union[str, None], asn1.Implicit(1)
+            ]
 
         assert_roundtrips(
             [
@@ -482,8 +486,10 @@ class Example:
                         h=None,
                         i=None,
                         j=None,
+                        k=None,
+                        only_field_present="a",
                     ),
-                    b"\x30\x00",
+                    b"\x30\x03\x81\x01a",
                 )
             ]
         )