fix: correctly set the post-handshake codec limit

Signed-off-by: ljedrz <[email protected]>
vitalmervin · Apr 20, 2023 · 13cd025 · 13cd025
1 parent 096c0a3
commit 13cd025
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 13 deletions.
diff --git a/node/messages/src/helpers/codec.rs b/node/messages/src/helpers/codec.rs
@@ -34,19 +34,17 @@ pub struct MessageCodec<N: Network> {
 }
 
 impl<N: Network> MessageCodec<N> {
-    /// Increases the maximum permitted message size post-handshake.
-    pub fn update_max_message_len(&mut self) {
-        self.codec.set_max_frame_length(MAXIMUM_MESSAGE_SIZE);
+    pub fn handshake() -> Self {
+        let mut codec = Self::default();
+        codec.codec.set_max_frame_length(MAXIMUM_HANDSHAKE_MESSAGE_SIZE);
+        codec
     }
 }
 
 impl<N: Network> Default for MessageCodec<N> {
     fn default() -> Self {
         Self {
-            codec: LengthDelimitedCodec::builder()
-                .max_frame_length(MAXIMUM_HANDSHAKE_MESSAGE_SIZE)
-                .little_endian()
-                .new_codec(),
+            codec: LengthDelimitedCodec::builder().max_frame_length(MAXIMUM_MESSAGE_SIZE).little_endian().new_codec(),
             _phantom: Default::default(),
         }
     }

diff --git a/node/router/src/handshake.rs b/node/router/src/handshake.rs
@@ -106,7 +106,7 @@ impl<N: Network> Router<N> {
         };
 
         // Perform the handshake; we pass on a mutable reference to peer_ip in case the process is broken at any point in time.
-        let mut handshake_result = if peer_side == ConnectionSide::Responder {
+        let handshake_result = if peer_side == ConnectionSide::Responder {
             self.handshake_inner_initiator(peer_addr, &mut peer_ip, stream, genesis_header).await
         } else {
             self.handshake_inner_responder(peer_addr, &mut peer_ip, stream, genesis_header).await
@@ -117,10 +117,9 @@ impl<N: Network> Router<N> {
             self.connecting_peers.lock().remove(&ip);
         }
 
-        // If the handshake succeeded, announce it and increase the message size limit.
-        if let Ok((ref peer_ip, ref mut framed)) = handshake_result {
+        // If the handshake succeeded, announce it.
+        if let Ok((ref peer_ip, _)) = handshake_result {
             info!("Connected to '{peer_ip}'");
-            framed.codec_mut().update_max_message_len();
         }
 
         handshake_result
@@ -135,7 +134,7 @@ impl<N: Network> Router<N> {
         genesis_header: Header<N>,
     ) -> io::Result<(SocketAddr, Framed<&mut TcpStream, MessageCodec<N>>)> {
         // Construct the stream.
-        let mut framed = Framed::new(stream, MessageCodec::<N>::default());
+        let mut framed = Framed::new(stream, MessageCodec::<N>::handshake());
 
         // This value is immediately guaranteed to be present, so it can be unwrapped.
         let peer_ip = peer_ip.unwrap();
@@ -199,7 +198,7 @@ impl<N: Network> Router<N> {
         genesis_header: Header<N>,
     ) -> io::Result<(SocketAddr, Framed<&mut TcpStream, MessageCodec<N>>)> {
         // Construct the stream.
-        let mut framed = Framed::new(stream, MessageCodec::<N>::default());
+        let mut framed = Framed::new(stream, MessageCodec::<N>::handshake());
 
         /* Step 1: Receive the challenge request. */