Merge pull request #77 from ngjunsiang/master

fix: Create bitfield structs with the correct size

Author: leif81

opendis/record.py

@@ -4,40 +4,85 @@
 """
 
 from collections.abc import Sequence
-from ctypes import _SimpleCData, BigEndianStructure, c_uint
+from ctypes import (
+    _SimpleCData,
+    BigEndianStructure,
+    c_uint8,
+    c_uint16,
+    c_uint32,
+    sizeof,
+)
+from typing import Literal
+
+from .stream import DataInputStream, DataOutputStream
 from .types import (
     bf_enum,
     bf_int,
     bf_uint,
 )
 
-from .DataInputStream import DataInputStream
-from .DataOutputStream import DataOutputStream
-
-
-def _bitfield(
-        name: str,
-        bytesize: int,
-        fields: Sequence[
-            tuple[str, type[_SimpleCData]] | tuple[str, type[_SimpleCData], int]
-        ],
-    ):
+# Type definitions for bitfield field descriptors
+CTypeFieldDescription = tuple[str, type[_SimpleCData], int]
+DisFieldDescription = tuple[str, "DisFieldType", int]
+
+# Field type constants simplify the construction of bitfields
+# which would otherwise require manually specifying ctypes types.
+# The currently implemented bitfields only use integers, but DIS7
+# mentions CHAR types which may be needed in future.
+DisFieldType = Literal["INTEGER"]
+INTEGER = "INTEGER"
+
+
+def field(name: str,
+          ftype: DisFieldType,
+          bits: int) -> CTypeFieldDescription:
+    """Helper function to create the field description tuple used by ctypes."""
+    match (ftype, bits):
+        case (INTEGER, b) if 0 < b <= 8:
+            return (name, c_uint8, bits)
+        case (INTEGER, b) if 8 < b <= 16:
+            return (name, c_uint16, bits)
+        case (INTEGER, b) if 16 < b <= 32:
+            return (name, c_uint32, bits)
+        case _:
+            raise ValueError(f"Unrecognized (ftype, bits): {ftype}, {bits}")
+
+
+def _bitfield(name: str,
+              fields: Sequence[DisFieldDescription]):
     """Factory function for bitfield structs, which are subclasses of
     ctypes.Structure.
     These are used in records that require them to unpack non-octet-sized fields.
 
     Args:
         name: Name of the bitfield struct.
         bytesize: Size of the bitfield in bytes.
-        fields: Sequence of tuples defining the fields of the bitfield.
-                See https://docs.python.org/3/library/ctypes.html#ctypes.Structure._fields_
+        fields: Sequence of tuples defining fields of the bitfield, in the form
+            (field_name, "INTEGER", field_size_in_bits).
     """
-    if bytesize <= 0:
-        raise ValueError("Cannot create bitfield with less than one byte")
-
+    # Argument validation
+    struct_fields = []
+    bitsize = 0
+    for name, ftype, bits in fields:
+        if ftype not in (INTEGER,):
+            raise ValueError(f"Unsupported field type: {ftype}")
+        if not isinstance(bits, int):
+            raise ValueError(f"Field size must be int: {bits!r}")
+        if bits <= 0 or bits > 32:
+            raise ValueError(f"Field size must be between 1 and 32: got {bits}")
+        bitsize += bits
+        struct_fields.append(field(name, ftype, bits))
+
+    if bitsize == 0:
+        raise ValueError(f"Bitfield size cannot be zero")
+    elif bitsize % 8 != 0:
+        raise ValueError(f"Bitfield size must be multiple of 8, got {bitsize}")
+    bytesize = bitsize // 8
+
+    # Create the struct class
     class Bitfield(BigEndianStructure):
-        _fields_ = fields
-    
+        _fields_ = struct_fields
+
         @staticmethod
         def marshalledSize() -> int:
             return bytesize
@@ -48,6 +93,12 @@ def serialize(self, outputStream: DataOutputStream) -> None:
         @classmethod
         def parse(cls, inputStream: DataInputStream) -> "Bitfield":
             return cls.from_buffer_copy(inputStream.read_bytes(bytesize))
+
+    # Sanity check: ensure the struct size matches expected size
+    assert sizeof(Bitfield) == bytesize, \
+        f"Bitfield size mismatch: expected {bytesize}, got {sizeof(Bitfield)}"
+
+    # Assign the class name
     Bitfield.__name__ = name
     return Bitfield
 
@@ -61,11 +112,11 @@ class NetId:
     YY = Frequency Table
     """
 
-    _struct = _bitfield(name="NetId", bytesize=2, fields=[
-        ("netNumber", c_uint, 10),
-        ("frequencyTable", c_uint, 2),
-        ("mode", c_uint, 2),
-        ("padding", c_uint, 2)
+    _struct = _bitfield(name="NetId", fields=[
+        ("netNumber", INTEGER, 10),
+        ("frequencyTable", INTEGER, 2),
+        ("mode", INTEGER, 2),
+        ("padding", INTEGER, 2)
     ])
 
     def __init__(self,
@@ -113,11 +164,11 @@ class SpreadSpectrum:
     In Python, the presence or absence of each technique is indicated by a bool.
     """
 
-    _struct = _bitfield("SpreadSpectrum", 2, [
-        ("frequencyHopping", c_uint, 1),
-        ("pseudoNoise", c_uint, 1),
-        ("timeHopping", c_uint, 1),
-        ("padding", c_uint, 13)
+    _struct = _bitfield(name="SpreadSpectrum", fields=[
+        ("frequencyHopping", INTEGER, 1),
+        ("pseudoNoise", INTEGER, 1),
+        ("timeHopping", INTEGER, 1),
+        ("padding", INTEGER, 13)
     ])
 
     def __init__(self,