From 8bebb3086fa2dad2ff7741b42dddf0d4530d7987 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8B=8F=E5=90=91=E5=A4=9C?= <fu050409@163.com>
Date: Sun, 24 Mar 2024 21:14:01 +0800
Subject: [PATCH] feat(gear): add flush after data wrote to stream

---
 src/exceptions.rs |  6 ++++--
 src/utils/gear.rs | 15 +++++++++++----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/exceptions.rs b/src/exceptions.rs
index b8db0d3..bcd5adb 100644
--- a/src/exceptions.rs
+++ b/src/exceptions.rs
@@ -1,10 +1,10 @@
 //! # Oblivion exception
 //! All exceptions to the Oblivion function return `OblivionException`.
+#[cfg(feature = "python")]
+use pyo3::prelude::*;
 use ring::error::Unspecified;
 use scrypt::errors::InvalidOutputLen;
 use thiserror::Error;
-#[cfg(feature = "python")]
-use pyo3::prelude::*;
 
 /// ## Oblivion exception iterator
 /// Use an iterator as the type of exception returned by a function.
@@ -55,6 +55,8 @@ pub enum OblivionException {
     EncryptError { error: Unspecified },
     #[error("Exception while decrypting: {error:?}")]
     DecryptError { error: Unspecified },
+    #[error("Failed to send some data: {message}")]
+    TCPWriteFailed { message: String },
 }
 
 #[cfg(feature = "python")]
diff --git a/src/utils/gear.rs b/src/utils/gear.rs
index 87b1f7f..4a421b6 100644
--- a/src/utils/gear.rs
+++ b/src/utils/gear.rs
@@ -35,7 +35,7 @@ impl AbsoluteNonceSequence {
 ///
 /// Used to abstract Oblivion's handling of transmitted data, wrapping all data type conversions.
 pub struct Socket {
-    tcp: TcpStream,
+    pub tcp: TcpStream,
 }
 
 impl Socket {
@@ -101,9 +101,16 @@ impl Socket {
     }
 
     pub async fn send(&mut self, data: &[u8]) -> Result<(), OblivionException> {
-        match self.tcp.write(data).await {
-            Ok(_) => Ok(()),
-            Err(_) => Err(OblivionException::UnexpectedDisconnection),
+        match self.tcp.write_all(&data).await {
+            Ok(_) => match self.tcp.flush().await {
+                Ok(_) => Ok(()),
+                Err(e) => Err(OblivionException::TCPWriteFailed {
+                    message: e.to_string(),
+                }),
+            },
+            Err(e) => Err(OblivionException::TCPWriteFailed {
+                message: e.to_string(),
+            }),
         }
     }