fix: assert that header length matches data length when user provides…

smp/exceptions.py

@@ -27,3 +27,7 @@ class SMPDeserializationError(SMPException):
 
 class SMPMismatchedGroupId(SMPDeserializationError):
     ...
+
+
+class SMPMalformed(SMPException):
+    ...

smp/message.py

@@ -9,7 +9,7 @@
 from pydantic import BaseModel, ConfigDict
 
 from smp import header as smpheader
-from smp.exceptions import SMPMismatchedGroupId
+from smp.exceptions import SMPMalformed, SMPMismatchedGroupId
 
 T = TypeVar("T", bound='_MessageBase')
 
@@ -86,6 +86,10 @@ def model_post_init(self, _: None) -> None:
                     ](self._COMMAND_ID),
                 ),
             )
+        elif self.header.length != len(data_bytes):
+            raise SMPMalformed(
+                f"header.length {self.header.length} != len(data_bytes) {len(data_bytes)}"
+            )
         self._bytes = cast(smpheader.Header, self.header).BYTES + data_bytes
 
 
@@ -123,6 +127,10 @@ def model_post_init(self, _: None) -> None:
                     ](self._COMMAND_ID),
                 ),
             )
+        elif self.header.length != len(data_bytes):
+            raise SMPMalformed(
+                f"header.length {self.header.length} != len(data_bytes) {len(data_bytes)}"
+            )
         self._bytes = cast(smpheader.Header, self.header).BYTES + data_bytes
 
 

tests/test_injected_header.py

@@ -0,0 +1,135 @@
+"""Test the case where the user forms the header separately."""
+
+from typing import cast
+
+import pytest
+
+from smp import header as smphdr
+from smp import image_management as smpimg
+from smp.exceptions import SMPMalformed
+
+
+def test_ImageUploadWriteRequest_injected_header() -> None:
+    h = smphdr.Header(
+        op=smphdr.OP.WRITE,
+        version=smphdr.Version.V0,
+        flags=smphdr.Flag(0),
+        length=0,
+        group_id=smphdr.GroupId.IMAGE_MANAGEMENT,
+        sequence=0,
+        command_id=smphdr.CommandId.ImageManagement.UPLOAD,
+    )
+
+    data = bytes([0x00] * 50)
+
+    r = smpimg.ImageUploadWriteRequest(
+        header=smphdr.Header(
+            op=h.op,
+            version=h.version,
+            flags=h.flags,
+            length=76,
+            group_id=h.group_id,
+            sequence=h.sequence,
+            command_id=h.command_id,
+        ),
+        off=0,
+        data=data,
+        image=1,
+        len=50,
+    )
+
+    assert cast(smphdr.Header, r.header).length == 76
+    assert len(r.BYTES) == 76 + smphdr.Header.SIZE
+
+    with pytest.raises(SMPMalformed):
+        r = smpimg.ImageUploadWriteRequest(
+            header=smphdr.Header(
+                op=h.op,
+                version=h.version,
+                flags=h.flags,
+                length=84,
+                group_id=h.group_id,
+                sequence=h.sequence,
+                command_id=h.command_id,
+            ),
+            off=0,
+            data=data,
+            image=1,
+            len=50,
+        )
+
+    with pytest.raises(SMPMalformed):
+        r = smpimg.ImageUploadWriteRequest(
+            header=smphdr.Header(
+                op=h.op,
+                version=h.version,
+                flags=h.flags,
+                length=0,
+                group_id=h.group_id,
+                sequence=h.sequence,
+                command_id=h.command_id,
+            ),
+            off=0,
+            data=data,
+            image=1,
+            len=50,
+        )
+
+
+def test_ImageUploadWriteResponse_injected_header() -> None:
+    h = smphdr.Header(
+        op=smphdr.OP.WRITE_RSP,
+        version=smphdr.Version.V0,
+        flags=smphdr.Flag(0),
+        length=0,
+        group_id=smphdr.GroupId.IMAGE_MANAGEMENT,
+        sequence=0,
+        command_id=smphdr.CommandId.ImageManagement.UPLOAD,
+    )
+
+    r = smpimg.ImageUploadProgressWriteResponse(
+        header=smphdr.Header(
+            op=h.op,
+            version=h.version,
+            flags=h.flags,
+            length=10,
+            group_id=h.group_id,
+            sequence=h.sequence,
+            command_id=h.command_id,
+        ),
+        rc=0,
+        off=0,
+    )
+
+    assert cast(smphdr.Header, r.header).length == 10
+    assert len(r.BYTES) == 10 + smphdr.Header.SIZE
+
+    with pytest.raises(SMPMalformed):
+        r = smpimg.ImageUploadProgressWriteResponse(
+            header=smphdr.Header(
+                op=h.op,
+                version=h.version,
+                flags=h.flags,
+                length=2,
+                group_id=h.group_id,
+                sequence=h.sequence,
+                command_id=h.command_id,
+            ),
+            rc=0,
+            off=0,
+        )
+
+    with pytest.raises(SMPMalformed):
+        r = smpimg.ImageUploadProgressWriteResponse(
+            header=smphdr.Header(
+                op=h.op,
+                version=h.version,
+                flags=h.flags,
+                length=0,
+                group_id=h.group_id,
+                sequence=h.sequence,
+                command_id=h.command_id,
+            ),
+            rc=0,
+            off=0,
+        )