diff --git a/server/recceiver/protocol/messages.py b/server/recceiver/protocol/messages.py index 58566e02..82f2e394 100644 --- a/server/recceiver/protocol/messages.py +++ b/server/recceiver/protocol/messages.py @@ -17,6 +17,14 @@ class ProtocolError(ValueError): """Raised when a RecSync message violates the wire protocol.""" +class KeyValueDecodeError(ValueError): + """Raised when recsync message content cannot be decoded as a valid utf-8 string""" + + def __init__(self, key: bytes, value: bytes): + self.key = key + self.value = value + + def _require_length(data, expected_length, message_name, exact=True): if exact and len(data) == expected_length: return @@ -154,8 +162,11 @@ def decode(cls, body): text = body[cls.payload.size :] _validate_record_lengths(kind, record_type_length, record_name_length) _require_length(text, record_type_length + record_name_length, "add record text") - record_type = text[:record_type_length].decode() - record_name = text[record_type_length:].decode() + try: + record_type = text[:record_type_length].decode() + record_name = text[record_type_length:].decode() + except UnicodeDecodeError: + raise KeyValueDecodeError(text[:record_type_length], text[record_type_length:]) return cls(record_id, kind, record_type, record_name) @@ -230,8 +241,12 @@ def decode(cls, body): if key_length == 0: raise ProtocolError("add info key must not be empty") _require_length(text, key_length + value_length, "add info text") - key = text[:key_length].decode() - value = text[key_length:].decode() + key = value = None + try: + key = text[:key_length].decode() + value = text[key_length:].decode() + except UnicodeDecodeError: + raise KeyValueDecodeError(text[:key_length], text[key_length:]) return cls(record_id, key, value) diff --git a/server/recceiver/recast.py b/server/recceiver/recast.py index e2790537..abc3cb39 100644 --- a/server/recceiver/recast.py +++ b/server/recceiver/recast.py @@ -140,6 +140,15 @@ def recvInfo(self, body): except messages.ProtocolError: log.error("Ignoring info update") return self.getInitialState() + except messages.KeyValueDecodeError as e: + log.exception( + "Error decoding key/value pair from {}:{}: '{}': '{}'", + self.sess.ep.host, + self.sess.ep.port, + e.key, + e.value, + ) + return self.getInitialState() if info.record_id: self.sess.rec_info(info.record_id, info.key, info.value) else: @@ -153,6 +162,15 @@ def recvAddRec(self, body): except messages.ProtocolError: log.error("Ignoring record update") return self.getInitialState() + except messages.KeyValueDecodeError as e: + log.exception( + "Error decoding record type/name pair from {}:{}: '{}': '{}'", + self.sess.ep.host, + self.sess.ep.port, + e.key, + e.value, + ) + return self.getInitialState() if record.is_alias: self.sess.add_alias(record.record_id, record.record_name) else: diff --git a/server/tests/unit/test_protocol.py b/server/tests/unit/test_protocol.py index d15699a3..ba43fb38 100644 --- a/server/tests/unit/test_protocol.py +++ b/server/tests/unit/test_protocol.py @@ -89,6 +89,18 @@ def test_decode_add_info_for_ioc_property(): assert info == messages.AddInfo(record_id=0, key="iocName", value="IOC-1") +def test_decode_raises_KeyValudeDecodeError_for_invalid_utf8_key(): + body = struct.pack(">IBxH", 0, 1, 3) + b"\xffval" + with pytest.raises(messages.KeyValueDecodeError): + messages.AddInfo.decode(body) + + +def test_decode_raises_KeyValudeDecodeError_for_invalid_utf8_value(): + body = struct.pack(">IBxH", 0, 3, 1) + b"key\xff" + with pytest.raises(messages.KeyValueDecodeError): + messages.AddInfo.decode(body) + + @pytest.mark.parametrize( "body", [