Catch KeyError for TYPE_TYPE missing in message

Added additional resiliency to the network protocol, such that if an
entirely garbled message is sent to the server (such as making a test
connection with nc and sending random characters), the protocol parser
reports that the message is invalid rather than producing KeyError due to
not being able to locate a message type.

Fixes: #4
Change-Id: I4839af0dc30fa1d6682129e60c13f9685fd845cf

Author: zzzeek

sqlalchemy_collectd/protocol.py

@@ -261,7 +261,12 @@ def __init__(self, *types):
 
     def receive(self, buf):
         result = self._unpack_packet(buf)
-        type_name = result[TYPE_TYPE]
+        try:
+            type_name = result[TYPE_TYPE]
+        except KeyError:
+            log.warn("Message did not have TYPE_TYPE block, skipping")
+            return None
+
         type_ = None
         try:
             type_ = self._types[type_name]

sqlalchemy_collectd/tests/test_protocol.py

@@ -81,3 +81,24 @@ def test_decode_packet(self):
             },
             result,
         )
+
+    def test_decode_unknown_type(self):
+        type_ = protocol.Type(
+            "my_type",
+            ("some_val", protocol.VALUE_GAUGE),
+            ("some_other_val", protocol.VALUE_DERIVE),
+        )
+
+        message_receiver = protocol.MessageReceiver(type_)
+        with mock.patch.object(protocol, "log") as log:
+            result = message_receiver.receive(b"asdfjq34kt2n34kjnas")
+        self.assertEqual(result, None)
+        self.assertEqual(
+            log.mock_calls,
+            [
+                mock.call.warn("Message %s not known, skipping", mock.ANY),
+                mock.call.warn(
+                    "Message did not have TYPE_TYPE block, skipping"
+                ),
+            ],
+        )

unreleased_changes/4.rst

@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, protocl
+    :tickets: 4
+
+    Added additional resiliency to the network protocol, such that if an
+    entirely garbled message is sent to the server (such as making a test
+    connection with nc and sending random characters), the protocol parser
+    reports that the message is invalid rather than producing KeyError due to
+    not being able to locate a message type.