From 33ef6189d4d96b58a2848e5e3bfb2d904f8b1417 Mon Sep 17 00:00:00 2001
From: "Ross A. Baker" <ross@rossabaker.com>
Date: Thu, 24 Oct 2024 17:43:04 -0400
Subject: [PATCH] Refactor connection closing logic

---
 .../http4s/blaze/client/Http1Connection.scala |  7 ++---
 .../org/http4s/blaze/core/Http1Stage.scala    | 28 +++++++++++++++++++
 .../blaze/server/Http1ServerStage.scala       | 19 ++++++-------
 3 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/blaze-client/src/main/scala/org/http4s/blaze/client/Http1Connection.scala b/blaze-client/src/main/scala/org/http4s/blaze/client/Http1Connection.scala
index 919cae4db..5df817233 100644
--- a/blaze-client/src/main/scala/org/http4s/blaze/client/Http1Connection.scala
+++ b/blaze-client/src/main/scala/org/http4s/blaze/client/Http1Connection.scala
@@ -200,10 +200,9 @@ private final class Http1Connection[F[_]](
           if (userAgent.nonEmpty && !req.headers.contains[`User-Agent`])
             rr << userAgent.get << "\r\n"
 
-          val mustClose: Boolean = req.headers.get[HConnection] match {
-            case Some(conn) => checkCloseConnection(conn, rr)
-            case None => getHttpMinor(req) == 0
-          }
+          val mustClose: Boolean = checkRequestCloseConnection(req)
+          if (mustClose)
+            rr << "Connection: close\r\n"
 
           val writeRequest: F[Boolean] = getChunkEncoder(req, mustClose, rr)
             .write(rr, req.entity)
diff --git a/blaze-core/src/main/scala/org/http4s/blaze/core/Http1Stage.scala b/blaze-core/src/main/scala/org/http4s/blaze/core/Http1Stage.scala
index bb8b00aa4..a5acf4c59 100644
--- a/blaze-core/src/main/scala/org/http4s/blaze/core/Http1Stage.scala
+++ b/blaze-core/src/main/scala/org/http4s/blaze/core/Http1Stage.scala
@@ -26,6 +26,7 @@ import org.http4s.Entity.Empty
 import org.http4s.Header
 import org.http4s.Header.Raw
 import org.http4s.Headers
+import org.http4s.HttpVersion
 import org.http4s.InvalidBodyException
 import org.http4s.Method
 import org.http4s.Request
@@ -67,6 +68,7 @@ private[blaze] trait Http1Stage[F[_]] { self: TailStage[ByteBuffer] =>
   protected def contentComplete(): Boolean
 
   /** Check Connection header and add applicable headers to response */
+  @deprecated("Use checkRequestCloseConnection(Request) instead", "0.23.17")
   protected final def checkCloseConnection(conn: Connection, rr: StringWriter): Boolean =
     if (conn.hasKeepAlive) { // connection, look to the request
       logger.trace("Found Keep-Alive header")
@@ -83,6 +85,32 @@ private[blaze] trait Http1Stage[F[_]] { self: TailStage[ByteBuffer] =>
       true
     }
 
+  private[http4s] final def checkRequestCloseConnection(req: Request[F]): Boolean = {
+    val conn = req.headers.get[Connection]
+    if (conn.fold(false)(_.hasClose)) {
+      logger.trace(s"Closing ${conn} due to explicit close option in request's Connection header")
+      true
+    } else if (req.httpVersion >= HttpVersion.`HTTP/1.1`) {
+      logger.trace(s"Keeping ${conn} alive per default behavior of HTTP >= 1.1")
+      false
+    } else if (req.httpVersion == HttpVersion.`HTTP/1.0`) {
+      if (conn.fold(false)(_.hasKeepAlive)) {
+        logger.trace(
+          s"Keeping ${conn} alive due to explicit keep-alive option in request's Connection header"
+        )
+        false
+      } else {
+        logger.trace(s"Closing ${conn} per default behavior of HTTP/1.0")
+        true
+      }
+    } else {
+      // It would be strange to serve HTTP/0.x, but we need a value and
+      // this is the rule.
+      logger.trace(s"Closing ${conn} for HTTP < 1.0")
+      true
+    }
+  }
+
   /** Get the proper body encoder based on the request */
   protected final def getEncoder(
       req: Request[F],
diff --git a/blaze-server/src/main/scala/org/http4s/blaze/server/Http1ServerStage.scala b/blaze-server/src/main/scala/org/http4s/blaze/server/Http1ServerStage.scala
index f17124bf5..7c675b7d7 100644
--- a/blaze-server/src/main/scala/org/http4s/blaze/server/Http1ServerStage.scala
+++ b/blaze-server/src/main/scala/org/http4s/blaze/server/Http1ServerStage.scala
@@ -245,12 +245,7 @@ private[blaze] class Http1ServerStage[F[_]](
     // Need to decide which encoder and if to close on finish
     val closeOnFinish = respConn
       .map(_.hasClose)
-      .orElse {
-        req.headers.get[Connection].map(checkCloseConnection(_, rr))
-      }
-      .getOrElse(
-        parser.minorVersion() == 0
-      ) // Finally, if nobody specifies, http 1.0 defaults to close
+      .getOrElse(checkRequestCloseConnection(req))
 
     // choose a body encoder. Will add a Transfer-Encoding header if necessary
     val bodyEncoder: Http1Writer[F] =
@@ -274,10 +269,14 @@ private[blaze] class Http1ServerStage[F[_]](
             case _ => // nop
           }
 
-        // add KeepAlive to Http 1.0 responses if the header isn't already present
-        rr << (if (!closeOnFinish && parser.minorVersion() == 0 && respConn.isEmpty)
-                 "Connection: keep-alive\r\n\r\n"
-               else "\r\n")
+        closeOnFinish match {
+          case true if respConn.isEmpty =>
+            rr << "Connection: close\r\n\r\n"
+          case false if parser.minorVersion() == 0 && respConn.isEmpty =>
+            rr << "Connection: keep-alive\r\n\r\n"
+          case _ =>
+            rr << "\r\n"
+        }
 
         new BodylessWriter[F](this, closeOnFinish)
       } else