Merge pull request #342 from njgheorghita/whoareyou-bugfix

Replace encoded message cipher with empty bytes for WHOAREYOU packet

ddht/base_message.py

@@ -15,6 +15,11 @@ def to_bytes(self) -> bytes:
         return b"".join((int_to_big_endian(self.message_type), rlp.encode(self)))
 
 
+class EmptyMessage(BaseMessage):
+    def to_bytes(self) -> bytes:
+        return b""
+
+
 TMessage = TypeVar("TMessage")
 TBaseMessage = TypeVar("TBaseMessage", bound=BaseMessage)
 TResponseMessage = TypeVar("TResponseMessage", bound=BaseMessage)

ddht/v5_1/packets.py

@@ -10,7 +10,7 @@
 from eth_utils.toolz import take
 import rlp
 
-from ddht.base_message import BaseMessage
+from ddht.base_message import BaseMessage, EmptyMessage
 from ddht.constants import UINT8_TO_BYTES
 from ddht.encryption import aesctr_decrypt_stream, aesctr_encrypt, aesgcm_encrypt
 from ddht.exceptions import DecodingError
@@ -244,12 +244,18 @@ def prepare(
             auth_data_size,
         )
         authenticated_data = b"".join((iv, header.to_wire_bytes(), auth_data_bytes,))
-        message_cipher_text = aesgcm_encrypt(
-            key=initiator_key,
-            nonce=aes_gcm_nonce,
-            plain_text=message.to_bytes(),
-            authenticated_data=authenticated_data,
-        )
+        # encrypted empty bytestring results in a bytestring,
+        # which we don't want if message is supposed to be empty
+        if type(message) is EmptyMessage:
+            message_cipher_text = b""
+        else:
+            message_cipher_text = aesgcm_encrypt(
+                key=initiator_key,
+                nonce=aes_gcm_nonce,
+                plain_text=message.to_bytes(),
+                authenticated_data=authenticated_data,
+            )
+
         return cls(
             iv=iv,
             header=header,

ddht/v5_1/session.py

@@ -13,7 +13,12 @@
 
 from ddht._utils import humanize_node_id
 from ddht.abc import HandshakeSchemeAPI, HandshakeSchemeRegistryAPI
-from ddht.base_message import AnyInboundMessage, AnyOutboundMessage, BaseMessage
+from ddht.base_message import (
+    AnyInboundMessage,
+    AnyOutboundMessage,
+    BaseMessage,
+    EmptyMessage,
+)
 from ddht.endpoint import Endpoint
 from ddht.exceptions import DecryptionError, HandshakeFailure
 from ddht.message_registry import MessageTypeRegistry
@@ -205,11 +210,6 @@ def to_bytes(self) -> bytes:
         return self.random_data
 
 
-class EmptyMessage(BaseMessage):
-    def to_bytes(self) -> bytes:
-        return b""
-
-
 #
 # Initiator
 #