Refactor

libp2p/security/noise/io.py

@@ -3,10 +3,12 @@ from typing import cast
 
 from noise.connection import NoiseConnection as NoiseState
 
-from libp2p.io.abc import ReadWriter
+from libp2p.io.abc import ReadWriteCloser, MsgReadWriter, EncryptedMsgReadWriter
+from libp2p.io.msgio import BaseMsgReadWriter, encode_msg_with_length
 from libp2p.io.utils import read_exactly
 from libp2p.network.connection.raw_connection_interface import IRawConnection
 
+
 SIZE_NOISE_MESSAGE_LEN = 2
 MAX_NOISE_MESSAGE_LEN = 2 ** (8 * SIZE_NOISE_MESSAGE_LEN) - 1
 SIZE_NOISE_MESSAGE_BODY_LEN = 2
@@ -20,53 +22,12 @@ BYTE_ORDER = "big"
 #                   <-2 bytes-><-       max=65533 bytes        ->
 
 
-def encode_data(data: bytes, size_len: int) -> bytes:
-    len_data = len(data)
-    try:
-        len_bytes = len_data.to_bytes(size_len, BYTE_ORDER)
-    except OverflowError as e:
-        raise ValueError from e
-    return len_bytes + data
+class NoisePacketReadWriter(BaseMsgReadWriter):
+    size_len_bytes = SIZE_NOISE_MESSAGE_LEN
 
 
-class MsgReader(ABC):
-    @abstractmethod
-    async def read_msg(self) -> bytes:
-        ...
-
-
-class MsgWriter(ABC):
-    @abstractmethod
-    async def write_msg(self, msg: bytes) -> None:
-        ...
-
-
-class MsgReadWriter(MsgReader, MsgWriter):
-    pass
-
-
-# TODO: Add comments
-class NoisePacketReadWriter(MsgReadWriter):
-    """Encode and decode the low level noise messages."""
-
-    read_writer: ReadWriter
-
-    def __init__(self, read_writer: ReadWriter) -> None:
-        self.read_writer = read_writer
-
-    async def read_msg(self) -> bytes:
-        len_bytes = await read_exactly(self.read_writer, SIZE_NOISE_MESSAGE_LEN)
-        len_int = int.from_bytes(len_bytes, BYTE_ORDER)
-        return await read_exactly(self.read_writer, len_int)
-
-    async def write_msg(self, msg: bytes) -> None:
-        encoded_data = encode_data(msg, SIZE_NOISE_MESSAGE_LEN)
-        await self.read_writer.write(encoded_data)
-
-
-# TODO: Add comments
 def encode_msg_body(msg_body: bytes) -> bytes:
-    encoded_msg_body = encode_data(msg_body, SIZE_NOISE_MESSAGE_BODY_LEN)
+    encoded_msg_body = encode_msg_with_length(msg_body, SIZE_NOISE_MESSAGE_BODY_LEN)
     if len(encoded_msg_body) > MAX_NOISE_MESSAGE_BODY_LEN:
         raise ValueError(
             f"msg_body is too long: {len(msg_body)}, "
@@ -88,39 +49,36 @@ def decode_msg_body(noise_msg: bytes) -> bytes:
     ]
 
 
-class NoiseHandshakeReadWriter(MsgReadWriter):
+class BaseNoiseMsgReadWriter(EncryptedMsgReadWriter):
     read_writer: MsgReadWriter
     noise_state: NoiseState
 
     def __init__(self, conn: IRawConnection, noise_state: NoiseState) -> None:
-        self.read_writer = NoisePacketReadWriter(cast(ReadWriter, conn))
+        self.read_writer = NoisePacketReadWriter(cast(ReadWriteCloser, conn))
         self.noise_state = noise_state
 
     async def write_msg(self, data: bytes) -> None:
         noise_msg = encode_msg_body(data)
-        data_encrypted = self.noise_state.write_message(noise_msg)
+        data_encrypted = self.encrypt(noise_msg)
         await self.read_writer.write_msg(data_encrypted)
 
     async def read_msg(self) -> bytes:
         noise_msg_encrypted = await self.read_writer.read_msg()
-        noise_msg = self.noise_state.read_message(noise_msg_encrypted)
+        noise_msg = self.decrypt(noise_msg_encrypted)
         return decode_msg_body(noise_msg)
 
 
-class NoiseTransportReadWriter(MsgReadWriter):
-    read_writer: MsgReadWriter
-    noise_state: NoiseState
+class NoiseHandshakeReadWriter(BaseNoiseMsgReadWriter):
+    def encrypt(self, data: bytes) -> bytes:
+        return self.noise_state.write_message(data)
 
-    def __init__(self, conn: IRawConnection, noise_state: NoiseState) -> None:
-        self.read_writer = NoisePacketReadWriter(cast(ReadWriter, conn))
-        self.noise_state = noise_state
+    def decrypt(self, data: bytes) -> bytes:
+        return self.noise_state.read_message(data)
 
-    async def write_msg(self, data: bytes) -> None:
-        noise_msg = encode_msg_body(data)
-        data_encrypted = self.noise_state.encrypt(noise_msg)
-        await self.read_writer.write_msg(data_encrypted)
 
-    async def read_msg(self) -> bytes:
-        noise_msg_encrypted = await self.read_writer.read_msg()
-        noise_msg = self.noise_state.decrypt(noise_msg_encrypted)
-        return decode_msg_body(noise_msg)
+class NoiseTransportReadWriter(BaseNoiseMsgReadWriter):
+    def encrypt(self, data: bytes) -> bytes:
+        return self.noise_state.encrypt(data)
+
+    def decrypt(self, data: bytes) -> bytes:
+        return self.noise_state.decrypt(data)

libp2p/security/noise/patterns.py

@@ -16,7 +16,7 @@ from .exceptions import (
     NoiseStateError,
     PeerIDMismatchesPubkey,
 )
-from .io import NoiseHandshakeReadWriter
+from .io import encode_msg_body, decode_msg_body, NoiseHandshakeReadWriter
 from .messages import (
     NoiseHandshakePayload,
     make_handshake_payload_sig,
@@ -56,6 +56,16 @@ class BasePattern(IPattern):
         )
         return NoiseHandshakePayload(self.libp2p_privkey.get_public_key(), signature)
 
+    async def write_msg(self, conn: IRawConnection, data: bytes) -> None:
+        noise_msg = encode_msg_body(data)
+        data_encrypted = self.noise_state.write_message(noise_msg)
+        await self.read_writer.write_msg(data_encrypted)
+
+    async def read_msg(self) -> bytes:
+        noise_msg_encrypted = await self.read_writer.read_msg()
+        noise_msg = self.noise_state.read_message(noise_msg_encrypted)
+        return decode_msg_body(noise_msg)
+
 
 class PatternXX(BasePattern):
     def __init__(