From b06518846f0a7b0d366b3c7da49b342be4c2663c Mon Sep 17 00:00:00 2001
From: MeexReay <meexreay@gmail.com>
Date: Fri, 20 Sep 2024 14:51:50 +0300
Subject: [PATCH] close stream after response send

---
 src/flowgate.rs           |  4 +++-
 src/flowgate/closeable.rs | 21 +++++++++++++++++++++
 src/flowgate/server.rs    |  7 ++++---
 src/flowgate/ssl_cert.rs  |  2 +-
 4 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 src/flowgate/closeable.rs

diff --git a/src/flowgate.rs b/src/flowgate.rs
index d4fd4c1..805eab7 100644
--- a/src/flowgate.rs
+++ b/src/flowgate.rs
@@ -1,7 +1,9 @@
 pub mod config;
 pub mod server;
 pub mod ssl_cert;
+pub mod closeable;
 
 pub use config::*;
 pub use server::*;
-pub use ssl_cert::*;
\ No newline at end of file
+pub use ssl_cert::*;
+pub use closeable::*;
\ No newline at end of file
diff --git a/src/flowgate/closeable.rs b/src/flowgate/closeable.rs
new file mode 100644
index 0000000..ae4d812
--- /dev/null
+++ b/src/flowgate/closeable.rs
@@ -0,0 +1,21 @@
+use std::net::{Shutdown, TcpStream};
+
+#[cfg(feature = "use-openssl")]
+use openssl::ssl::SslStream;
+
+pub trait Closeable {
+    fn close(&self);
+}
+
+impl Closeable for TcpStream {
+    fn close(&self) {
+        let _ = self.shutdown(Shutdown::Both);
+    }
+}
+
+#[cfg(feature = "use-openssl")]
+impl<T: Closeable> Closeable for SslStream<T> {
+    fn close(&self) {
+        self.get_ref().close();
+    }
+}
diff --git a/src/flowgate/server.rs b/src/flowgate/server.rs
index 14748c5..4b1fa92 100644
--- a/src/flowgate/server.rs
+++ b/src/flowgate/server.rs
@@ -3,7 +3,7 @@ use std::{io::{Read, Write}, net::{Shutdown, SocketAddr, TcpListener}, sync::Arc
 use log::info;
 use threadpool::ThreadPool;
 
-use super::Config;
+use super::{Closeable, Config};
 
 pub struct FlowgateServer {
     config: Arc<Config>,
@@ -179,7 +179,7 @@ impl FlowgateServer {
 
     pub fn accept_stream(
         config: Arc<Config>, 
-        stream: &mut (impl Read + Write), 
+        stream: &mut (impl Read + Write + Closeable), 
         addr: SocketAddr,
         https: bool
     ) -> Option<()> {
@@ -314,7 +314,8 @@ impl FlowgateServer {
             }
         }
 
-        site_stream.shutdown(Shutdown::Both).ok()?;
+        site_stream.close();
+        stream.close();
 
         Some(())
     }
diff --git a/src/flowgate/ssl_cert.rs b/src/flowgate/ssl_cert.rs
index 2a43cc3..0cef70a 100644
--- a/src/flowgate/ssl_cert.rs
+++ b/src/flowgate/ssl_cert.rs
@@ -108,4 +108,4 @@ impl AdoptedConnection {
     }
 }
 
-// TODO: implement Read and Write to AdoptedConnection
\ No newline at end of file
+// TODO: implement Read and Write and Closeable to AdoptedConnection
\ No newline at end of file