asn1: Add support for SET

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

Author: facutuesca

src/cryptography/hazmat/asn1/__init__.py

@@ -19,6 +19,7 @@
     decode_der,
     encode_der,
     sequence,
+    set,
 )
 
 __all__ = [
@@ -38,4 +39,5 @@
     "decode_der",
     "encode_der",
     "sequence",
+    "set",
 ]

src/cryptography/hazmat/asn1/asn1.py

@@ -326,6 +326,16 @@ def _register_asn1_sequence(cls: type[U]) -> None:
     setattr(cls, "__asn1_root__", root)
 
 
+def _register_asn1_set(cls: type[U]) -> None:
+    raw_fields = get_type_hints(cls, include_extras=True)
+    root = declarative_asn1.AnnotatedType(
+        declarative_asn1.Type.Set(cls, _annotate_fields(raw_fields)),
+        declarative_asn1.Annotation(),
+    )
+
+    setattr(cls, "__asn1_root__", root)
+
+
 # Due to https://github.com/python/mypy/issues/19731, we can't define an alias
 # for `dataclass_transform` that conditionally points to `typing` or
 # `typing_extensions` depending on the Python version (like we do for
@@ -357,6 +367,29 @@ def sequence(cls: type[U]) -> type[U]:
         _register_asn1_sequence(dataclass_cls)
         return dataclass_cls
 
+    @typing_extensions.dataclass_transform(kw_only_default=True)
+    def set(cls: type[U]) -> type[U]:
+        # We use `dataclasses.dataclass` to add an __init__ method
+        # to the class with keyword-only parameters.
+        if sys.version_info >= (3, 10):
+            dataclass_cls = dataclasses.dataclass(
+                repr=False,
+                eq=False,
+                # `match_args` was added in Python 3.10 and defaults
+                # to True
+                match_args=False,
+                # `kw_only` was added in Python 3.10 and defaults to
+                # False
+                kw_only=True,
+            )(cls)
+        else:
+            dataclass_cls = dataclasses.dataclass(
+                repr=False,
+                eq=False,
+            )(cls)
+        _register_asn1_set(dataclass_cls)
+        return dataclass_cls
+
 else:
 
     @typing.dataclass_transform(kw_only_default=True)
@@ -372,6 +405,19 @@ def sequence(cls: type[U]) -> type[U]:
         _register_asn1_sequence(dataclass_cls)
         return dataclass_cls
 
+    @typing.dataclass_transform(kw_only_default=True)
+    def set(cls: type[U]) -> type[U]:
+        # Only add an __init__ method, with keyword-only
+        # parameters.
+        dataclass_cls = dataclasses.dataclass(
+            repr=False,
+            eq=False,
+            match_args=False,
+            kw_only=True,
+        )(cls)
+        _register_asn1_set(dataclass_cls)
+        return dataclass_cls
+
 
 # TODO: replace with `Default[U]` once the min Python version is >= 3.12
 @dataclasses.dataclass(frozen=True)

src/cryptography/hazmat/bindings/_rust/declarative_asn1.pyi

@@ -14,6 +14,7 @@ def non_root_python_to_rust(cls: type) -> Type: ...
 class Type:
     Sequence: typing.ClassVar[type]
     SequenceOf: typing.ClassVar[type]
+    Set: typing.ClassVar[type]
     SetOf: typing.ClassVar[type]
     Option: typing.ClassVar[type]
     Choice: typing.ClassVar[type]

src/rust/src/declarative_asn1/decode.rs

@@ -286,6 +286,21 @@ pub(crate) fn decode_annotated_type<'a>(
                 Ok(list.into_any())
             })?
         }
+        Type::Set(cls, fields) => {
+            let set_parse_result = read_value::<asn1::Set<'_>>(parser, encoding)?;
+
+            set_parse_result.parse(|d| -> ParseResult<pyo3::Bound<'a, pyo3::PyAny>> {
+                let kwargs = pyo3::types::PyDict::new(py);
+                let fields = fields.bind(py);
+                for (name, ann_type) in fields.into_iter() {
+                    let ann_type = ann_type.cast::<AnnotatedType>()?;
+                    let value = decode_annotated_type(py, d, ann_type.get())?;
+                    kwargs.set_item(name, value)?;
+                }
+                let val = cls.call(py, (), Some(&kwargs))?.into_bound(py);
+                Ok(val)
+            })?
+        }
         Type::SetOf(cls) => {
             let setof_parse_result = read_value::<asn1::Set<'_>>(parser, encoding)?;
 

src/rust/src/declarative_asn1/encode.rs

@@ -78,6 +78,22 @@ impl asn1::Asn1Writable for AnnotatedTypeObject<'_> {
 
                 write_value(writer, &asn1::SequenceOfWriter::new(values), encoding)
             }
+            Type::Set(_cls, fields) => write_value(
+                writer,
+                &asn1::SetWriter::new(&|w| {
+                    for (name, ann_type) in fields.bind(py).into_iter() {
+                        let name = name.cast::<pyo3::types::PyString>()?;
+                        let ann_type = ann_type.cast::<AnnotatedType>()?;
+                        let object = AnnotatedTypeObject {
+                            annotated_type: ann_type.get(),
+                            value: self.value.getattr(name)?,
+                        };
+                        w.write_element(&object)?;
+                    }
+                    Ok(())
+                }),
+                encoding,
+            ),
             Type::SetOf(cls) => {
                 let setof = value.cast::<super::types::SetOf>()?;
                 let values: Vec<AnnotatedTypeObject<'_>> = setof

src/rust/src/declarative_asn1/types.rs

@@ -23,6 +23,10 @@ pub enum Type {
     Sequence(pyo3::Py<pyo3::types::PyType>, pyo3::Py<pyo3::types::PyDict>),
     /// SEQUENCE OF (`list[`T`]`)
     SequenceOf(pyo3::Py<AnnotatedType>),
+    /// SET(`class`, `dict`)
+    /// The first element is the Python class that represents the set,
+    /// the second element is a dict of the (already converted) fields of the class.
+    Set(pyo3::Py<pyo3::types::PyType>, pyo3::Py<pyo3::types::PyDict>),
     /// SET OF (`list[`T`]`)
     SetOf(pyo3::Py<AnnotatedType>),
     /// OPTIONAL (`T | None`)
@@ -640,6 +644,7 @@ pub(crate) fn is_tag_valid_for_type(
 ) -> bool {
     match type_ {
         Type::Sequence(_, _) => check_tag_with_encoding(asn1::Sequence::TAG, encoding, tag),
+        Type::Set(_, _) => check_tag_with_encoding(asn1::Set::TAG, encoding, tag),
         Type::SequenceOf(_) => check_tag_with_encoding(asn1::Sequence::TAG, encoding, tag),
         Type::SetOf(_) => check_tag_with_encoding(asn1::SetOf::<()>::TAG, encoding, tag),
         Type::Option(t) => is_tag_valid_for_type(py, tag, t.get().inner.get(), encoding),

tests/hazmat/asn1/test_serialization.py

@@ -1037,6 +1037,77 @@ class Example:
             )
 
 
+class TestSet:
+    def test_ok_set_single_field(self) -> None:
+        @asn1.set
+        @_comparable_dataclass
+        class Example:
+            foo: int
+
+        assert_roundtrips([(Example(foo=9), b"\x31\x03\x02\x01\x09")])
+
+    def test_ok_set_multiple_fields(self) -> None:
+        @asn1.set
+        @_comparable_dataclass
+        class Example:
+            foo: int
+            bar: int
+
+        assert_roundtrips(
+            [(Example(foo=6, bar=9), b"\x31\x06\x02\x01\x06\x02\x01\x09")]
+        )
+
+    def test_fail_set_multiple_fields_wrong_order(self) -> None:
+        @asn1.set
+        @_comparable_dataclass
+        class Example:
+            foo: int
+            bar: int
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "invalid SET ordering while performing ASN.1 serialization"
+            ),
+        ):
+            assert_roundtrips(
+                [(Example(foo=9, bar=6), b"\x31\x06\x02\x01\x06\x02\x01\x09")]
+            )
+
+    def test_ok_nested_set(self) -> None:
+        @asn1.set
+        @_comparable_dataclass
+        class Child:
+            foo: int
+
+        @asn1.set
+        @_comparable_dataclass
+        class Parent:
+            foo: Child
+
+        assert_roundtrips(
+            [(Parent(foo=Child(foo=9)), b"\x31\x05\x31\x03\x02\x01\x09")]
+        )
+
+    def test_ok_set_multiple_types(self) -> None:
+        @asn1.set
+        @_comparable_dataclass
+        class Example:
+            a: bool
+            b: int
+            c: bytes
+            d: str
+
+        assert_roundtrips(
+            [
+                (
+                    Example(a=True, b=9, c=b"c", d="d"),
+                    b"\x31\x0c\x01\x01\xff\x02\x01\x09\x04\x01c\x0c\x01d",
+                )
+            ]
+        )
+
+
 class TestSize:
     def test_ok_sequenceof_size_restriction(self) -> None:
         @asn1.sequence