From 80cd29232def079e3dcec2b6c4bf58484aae67aa Mon Sep 17 00:00:00 2001
From: Ishkhan Nazaryan <105867377+ishkhan42@users.noreply.github.com>
Date: Thu, 7 Sep 2023 13:45:39 +0000
Subject: [PATCH] Fix: Decrypt TLS requests fully Add: Handle `Connection:
 close` header

---
 src/headers/automata.hpp   |  8 ++++----
 src/headers/connection.hpp | 18 +++++++++++++++---
 src/headers/containers.hpp |  3 +++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/headers/automata.hpp b/src/headers/automata.hpp
index 878b761..988e62e 100644
--- a/src/headers/automata.hpp
+++ b/src/headers/automata.hpp
@@ -165,10 +165,10 @@ void automata_t::operator()() noexcept {
         connection.pipes.mark_submitted_outputs(completed_result);
         if (!connection.pipes.has_remaining_outputs()) {
             connection.exchanges++;
-            // if (connection.exchanges >= server.max_lifetime_exchanges) TODO Why?
-            //     return close_gracefully();
-            // else
-            return receive_next();
+            if (connection.must_close())
+                return close_gracefully();
+            else
+                return receive_next();
         } else {
             connection.pipes.prepare_more_outputs();
             return send_next();
diff --git a/src/headers/connection.hpp b/src/headers/connection.hpp
index 22f8eb9..2a03a79 100644
--- a/src/headers/connection.hpp
+++ b/src/headers/connection.hpp
@@ -72,10 +72,15 @@ struct connection_t {
     bool expired() const noexcept {
         return std::chrono::high_resolution_clock::now().time_since_epoch().count() - last_active_ns >
                max_inactive_duration_ns_k;
-    };
+    }
 
     bool is_ready() const noexcept { return tls_ctx == nullptr || ptls_handshake_is_complete(tls_ctx); }
 
+    bool must_close() const noexcept {
+        auto conn = protocol.get_header("Connection");
+        return conn == "Close" || conn == "close";
+    }
+
     bool prepare_step() noexcept {
         if (is_ready())
             return true;
@@ -114,10 +119,17 @@ struct connection_t {
             return;
 
         work_buf.off = 0;
+        int res = 0;
         size_t in_len = pipes.input_span().size();
-        int res = ptls_receive(tls_ctx, &work_buf, pipes.input_span().data(), &in_len);
+        void const* input = pipes.input_span().data();
+        while (in_len != 0 && res != -1) {
+            size_t consumed = in_len;
+            res = ptls_receive(tls_ctx, &work_buf, input, &consumed);
+            in_len -= consumed;
+            input += consumed;
+        }
         if (res != -1) {
-            pipes.release_inputs();
+            pipes.release_current_input();
             std::memcpy(pipes.next_input_address(), work_buf.base, work_buf.off);
             pipes.absorb_input(work_buf.off);
         }
diff --git a/src/headers/containers.hpp b/src/headers/containers.hpp
index 6976c68..9818f51 100644
--- a/src/headers/containers.hpp
+++ b/src/headers/containers.hpp
@@ -218,6 +218,9 @@ class exchange_pipes_t {
         input_.dynamic.reset();
         input_.embedded_used = 0;
     }
+
+    void release_current_input() noexcept { input_.embedded_used = 0; }
+
     void release_outputs() noexcept {
         output_.dynamic.reset();
         output_.embedded_used = 0;