From 4f9c5cbd6cab4d4804f503d176737d5dd3947315 Mon Sep 17 00:00:00 2001
From: Yuki Kishimoto <yukikishimoto@protonmail.com>
Date: Wed, 21 Aug 2024 06:51:27 +0000
Subject: [PATCH] Flush stream on final handshake stage (#431)

Flush the stream before finalizing the writing stage in the handshake machine. This prevents potential data loss by ensuring all buffered contents are transmitted.
---
 src/handshake/machine.rs | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/handshake/machine.rs b/src/handshake/machine.rs
index 2e3f2cbb..fc3cb551 100644
--- a/src/handshake/machine.rs
+++ b/src/handshake/machine.rs
@@ -81,7 +81,10 @@ impl<Stream: Read + Write> HandshakeMachine<Stream> {
                             ..self
                         })
                     } else {
-                        RoundResult::StageFinished(StageResult::DoneWriting(self.stream))
+                        RoundResult::Incomplete(HandshakeMachine {
+                            state: HandshakeState::Flushing,
+                            ..self
+                        })
                     })
                 } else {
                     Ok(RoundResult::WouldBlock(HandshakeMachine {
@@ -90,6 +93,13 @@ impl<Stream: Read + Write> HandshakeMachine<Stream> {
                     }))
                 }
             }
+            HandshakeState::Flushing => Ok(match self.stream.flush().no_block()? {
+                Some(()) => RoundResult::StageFinished(StageResult::DoneWriting(self.stream)),
+                None => RoundResult::WouldBlock(HandshakeMachine {
+                    state: HandshakeState::Flushing,
+                    ..self
+                }),
+            }),
         }
     }
 }
@@ -128,6 +138,8 @@ enum HandshakeState {
     Reading(ReadBuffer, AttackCheck),
     /// Sending data to the peer.
     Writing(Cursor<Vec<u8>>),
+    /// Flushing data to ensure that all intermediately buffered contents reach their destination.
+    Flushing,
 }
 
 /// Attack mitigation. Contains counters needed to prevent DoS attacks