From cbd0067cb005311de5082093c2a4c50a57bc88a5 Mon Sep 17 00:00:00 2001
From: MeexReay <meexreay@gmail.com>
Date: Sat, 14 Sep 2024 01:08:28 +0300
Subject: [PATCH] keep alive config settings

---
 conf.yml               |  3 +++
 src/flowgate/config.rs |  2 ++
 src/flowgate/server.rs | 10 +++++++++-
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/conf.yml b/conf.yml
index 7005125..49a92a3 100644
--- a/conf.yml
+++ b/conf.yml
@@ -6,9 +6,12 @@ sites:
   #   host: localhost:8080
   #   ssl_cert: "/path/to/public/certificate.txt"
   #   ssl_key: "/path/to/private/key.txt"
+  #   support_keep_alive: true
 
   # - domain: sub.example.com                        # Domain with no SSL
   #   host: localhost:8081
+  #   support_keep_alive: true
 
   - domain: localhost
     host: localhost:8080
+    support_keep_alive: false
diff --git a/src/flowgate/config.rs b/src/flowgate/config.rs
index 0bc1716..e1c5355 100644
--- a/src/flowgate/config.rs
+++ b/src/flowgate/config.rs
@@ -9,6 +9,7 @@ pub struct SiteConfig {
     pub domain: String,
     pub host: String,
     pub ssl: Option<SslCert>,
+    pub support_keep_alive: bool
 }
 
 impl SiteConfig {
@@ -53,6 +54,7 @@ impl Config {
                 domain: s.get("domain")?.as_str()?.to_string(),
                 host: s.get("host")?.as_str()?.to_string(),
                 ssl: cert,
+                support_keep_alive: s.get("support_keep_alive").map(|o| o.as_bool().unwrap()).unwrap_or(false)
             };
 
             sites.push(site);
diff --git a/src/flowgate/server.rs b/src/flowgate/server.rs
index cd2c9cc..1e2cc1f 100644
--- a/src/flowgate/server.rs
+++ b/src/flowgate/server.rs
@@ -165,7 +165,7 @@ impl FlowgateServer {
         }
 
         let site = site?.clone();
-        let mut site_stream = site.connect()?;
+        let mut site_stream = site.clone().connect()?;
 
         site_stream.write((addr.to_string() + "\n" + reqst).as_bytes()).ok()?;
 
@@ -195,6 +195,10 @@ impl FlowgateServer {
 
         if keep_alive {
             loop {
+                if !site.clone().support_keep_alive {
+                    site_stream.shutdown(Shutdown::Both).ok()?;
+                }
+
                 let mut reqst_data: Vec<u8> = vec![0; 4096];
 
                 stream.read(&mut reqst_data).ok()?;
@@ -220,6 +224,10 @@ impl FlowgateServer {
                     }
                 }
 
+                if !site.clone().support_keep_alive {
+                    site_stream = site.clone().connect()?
+                }
+
                 site_stream.write((addr.to_string() + "\n" + reqst).as_bytes()).ok()?;
 
                 if content_length != 0 && content_length > body.len() {