From 9a435b2f5dc85cd9ffd8c9bd303fd989feb0dd65 Mon Sep 17 00:00:00 2001
From: Leon Morten Richter <github@leonmortenrichter.de>
Date: Fri, 20 Dec 2024 11:19:40 +0100
Subject: [PATCH] chore: handle fragmented messages with empty payload

---
 pyais/exceptions.py       |  5 +++++
 pyais/messages.py         |  7 +++----
 tests/test_decode.py      | 25 +++++++++++++++++++++++++
 tests/test_decode_raw.py  |  1 -
 tests/test_file_stream.py |  7 +++++--
 5 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/pyais/exceptions.py b/pyais/exceptions.py
index 8c491ff..e2d3e98 100644
--- a/pyais/exceptions.py
+++ b/pyais/exceptions.py
@@ -38,3 +38,8 @@ class NonPrintableCharacterException(AISBaseException):
 
 class TagBlockNotInitializedException(Exception):
     """The TagBlock is not initialized"""
+
+
+class MissingPayloadException(AISBaseException):
+    """Valid NMEA Message without payload"""
+    pass
diff --git a/pyais/messages.py b/pyais/messages.py
index 56a13e3..ac29fe2 100644
--- a/pyais/messages.py
+++ b/pyais/messages.py
@@ -12,7 +12,7 @@
 from pyais.constants import TalkerID, NavigationStatus, ManeuverIndicator, EpfdType, ShipType, NavAid, StationType, \
     TransmitMode, StationIntervals, TurnRate
 from pyais.exceptions import InvalidNMEAMessageException, TagBlockNotInitializedException, UnknownMessageException, UnknownPartNoException, \
-    InvalidDataTypeException
+    InvalidDataTypeException, MissingPayloadException
 from pyais.util import checksum, decode_into_bit_array, compute_checksum, get_itdma_comm_state, get_sotdma_comm_state, int_to_bin, str_to_bin, \
     encode_ascii_6, from_bytes, from_bytes_signed, decode_bin_as_ascii6, get_int, chk_to_int, coerce_val, \
     bits2bytes, bytes2bits, b64encode_str
@@ -500,9 +500,6 @@ def __init__(self, raw: bytes) -> None:
         except Exception as err:
             raise InvalidNMEAMessageException(raw) from err
 
-        if not len(payload):
-            raise InvalidNMEAMessageException("Invalid empty payload")
-
         if len(payload) > MAX_PAYLOAD_LEN:
             raise InvalidNMEAMessageException("AIS payload too large")
 
@@ -601,6 +598,8 @@ def decode(self) -> "ANY_MESSAGE":
         >>> nmea = NMEAMessage(b"!AIVDO,1,1,,,B>qc:003wk?8mP=18D3Q3wgTiT;T,0*13").decode()
         MessageType18(msg_type=18, ...)
         """
+        if not self.payload:
+            raise MissingPayloadException(self.raw.decode())
         try:
             return MSG_CLASS[self.ais_id].from_bitarray(self.bit_array)
         except KeyError as e:
diff --git a/tests/test_decode.py b/tests/test_decode.py
index 6bdc0d8..6e21a37 100644
--- a/tests/test_decode.py
+++ b/tests/test_decode.py
@@ -48,6 +48,7 @@
 )
 from pyais.stream import ByteStream
 from pyais.util import b64encode_str, bits2bytes, bytes2bits, decode_into_bit_array
+from pyais.exceptions import MissingPayloadException
 
 
 def ensure_type_for_msg_dict(msg_dict: typing.Dict[str, typing.Any]) -> None:
@@ -1743,3 +1744,27 @@ def test_that_decode_nmea_and_ais_works_with_proprietary_messages(self):
         self.assertEqual(decoded.msg_type, 1)
         self.assertEqual(decoded.mmsi, 538090443)
         self.assertEqual(decoded.speed, 10.9)
+
+    def test_that_decode_works_for_fragmented_messages_with_empty_payloads(self):
+        """Issue: https://github.com/M0r13n/pyais/issues/157"""
+        # WHEN decoding a fragmented message where the second message has an empty payload.
+        decoded = decode(
+            b"!AIVDM,2,1,0,A,8@2R5Ph0GhRbUqe?n>KS?wvlFR06EuOwiOl?wnSwe7wvlOwwsAwwnSGmwvwt,0*4E",
+            b"!AIVDM,2,2,0,A,,0*16",
+        )
+        # THEN the message is decoded without an error
+        # Verified against https://www.aggsoft.com/ais-decoder.htm
+        self.assertEqual(decoded.msg_type, 8)
+        self.assertEqual(decoded.repeat, 1)
+        self.assertEqual(decoded.mmsi, 2655619)
+        self.assertEqual(decoded.data, b'\x08\xaa\x97\x9bO\xd8\xe6\xe3?\xff\xb4Z \x06W\xd7\xff\xc5\xfd\x0f\xffh\xff\xb4\x7f\xfe\xd1\xff\xff\xed\x1f\xff\xda5\xf5\xff\xef\xfc')
+
+    def test_decode_with_empty_payload(self):
+        """Variation of test_that_decode_works_for_fragmented_messages_with_empty_payloads"""
+        # WHEN decoding message without payload an exception is raised
+        with self.assertRaises(MissingPayloadException) as err:
+            _ = decode(
+                b"!AIVDM,1,1,0,A,,0*16",
+            )
+
+        self.assertEqual(str(err.exception), '!AIVDM,1,1,0,A,,0*16')
diff --git a/tests/test_decode_raw.py b/tests/test_decode_raw.py
index 9337c82..9f4f66f 100644
--- a/tests/test_decode_raw.py
+++ b/tests/test_decode_raw.py
@@ -44,7 +44,6 @@ def should_raise(msg):
         should_raise(",1,1,,A,403Ovl@000Htt<tSF0l4Q@100`Pq,0*28")
         should_raise("!AIVDM,,1,,A,403Ovl@000Htt<tSF0l4Q@100`Pq,0*28")
         should_raise("!AIVDM,1,,,A,403Ovl@000Htt<tSF0l4Q@100`Pq,0*28")
-        should_raise("!AIVDM,1,1,,A,,0*28")
 
         should_raise("!AIVDM,11111111111111,1,,A,403Ovl@000Htt<tSF0l4Q@100`Pq,0*28")
         should_raise("!AIVDM,1,11111111111111111111,,A,403Ovl@000Htt<tSF0l4Q@100`Pq,0*28")
diff --git a/tests/test_file_stream.py b/tests/test_file_stream.py
index 10efdb9..835ae0f 100644
--- a/tests/test_file_stream.py
+++ b/tests/test_file_stream.py
@@ -3,7 +3,7 @@
 import unittest
 from unittest.case import skip
 
-from pyais.exceptions import UnknownMessageException
+from pyais.exceptions import UnknownMessageException, MissingPayloadException
 from pyais.messages import GatehouseSentence, NMEAMessage
 from pyais.stream import FileReaderStream, IterMessages
 
@@ -180,7 +180,10 @@ def test_marine_traffic_sample(self):
 
         with FileReaderStream(nmea_file) as stream:
             for msg in stream:
-                assert msg.decode()
+                try:
+                    assert msg.decode()
+                except MissingPayloadException:
+                    assert msg.raw.startswith(b'!AIVDM,1,1,,')
 
     def test_mixed_content(self):
         """Test that the file reader handles mixed content. That means, that is is able to handle