From 5ff647d45eb41fc9f0520d9e18c041677d738a01 Mon Sep 17 00:00:00 2001
From: Antonio Nuno Monteiro <anmonteiro@gmail.com>
Date: Fri, 29 Nov 2024 17:58:42 -0800
Subject: [PATCH] refactor: track `proxy` in config for whether this is a proxy
 connection

---
 lib/config.ml            | 8 ++++++--
 lib/httpun.mli           | 1 +
 lib/reqd.ml              | 6 ++++--
 lib/server_connection.ml | 8 +++++---
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/lib/config.ml b/lib/config.ml
index d417040..b817638 100644
--- a/lib/config.ml
+++ b/lib/config.ml
@@ -2,10 +2,14 @@ type t =
   { read_buffer_size          : int
   ; request_body_buffer_size  : int
   ; response_buffer_size      : int
-  ; response_body_buffer_size : int }
+  ; response_body_buffer_size : int
+  ; proxy : bool
+  }
 
 let default =
   { read_buffer_size          = 0x1000
   ; request_body_buffer_size  = 0x1000
   ; response_buffer_size      = 0x400
-  ; response_body_buffer_size = 0x1000 }
+  ; response_body_buffer_size = 0x1000
+  ; proxy = false
+  }
diff --git a/lib/httpun.mli b/lib/httpun.mli
index 628c8cf..bf34543 100644
--- a/lib/httpun.mli
+++ b/lib/httpun.mli
@@ -265,6 +265,7 @@ module Config : sig
     ; request_body_buffer_size  : int (** Default is [4096] *)
     ; response_buffer_size      : int (** Default is [1024] *)
     ; response_body_buffer_size : int (** Default is [4096] *)
+    ; proxy : bool
     }
 
   val default : t
diff --git a/lib/reqd.ml b/lib/reqd.ml
index 8845bc5..2cda47c 100644
--- a/lib/reqd.ml
+++ b/lib/reqd.ml
@@ -73,9 +73,10 @@ type t =
   ; mutable persistent      : bool
   ; mutable response_state  : Response_state.t
   ; mutable error_code      : [`Ok | error ]
+  ; proxy: bool
   }
 
-let create error_handler request request_body reader writer response_body_buffer =
+let create error_handler request request_body reader writer response_body_buffer proxy =
   { request
   ; request_body
   ; reader
@@ -85,6 +86,7 @@ let create error_handler request request_body reader writer response_body_buffer
   ; persistent              = Request.persistent_connection request
   ; response_state          = Waiting
   ; error_code              = `Ok
+  ; proxy
   }
 
 let request { request; _ } = request
@@ -142,7 +144,7 @@ let unsafe_respond_with_streaming ~flush_headers_immediately t response =
   match t.response_state with
   | Waiting ->
     let encoding =
-      match Response.body_length ~request_method:t.request.meth response with
+      match Response.body_length ~proxy:t.proxy ~request_method:t.request.meth response with
       | `Fixed _ | `Close_delimited | `Chunked as encoding -> encoding
       | `Error (`Bad_gateway | `Internal_server_error) ->
         failwith "httpun.Reqd.respond_with_streaming: invalid response body length"
diff --git a/lib/server_connection.ml b/lib/server_connection.ml
index 3398517..ffa7f7b 100644
--- a/lib/server_connection.ml
+++ b/lib/server_connection.ml
@@ -52,7 +52,8 @@ type error_code =
     }
 
 type t =
-  { reader                 : Reader.request
+  { config : Config.t
+  ; reader                 : Reader.request
   ; writer                 : Writer.t
   ; response_body_buffer   : Bigstringaf.t
   ; request_handler        : request_handler
@@ -119,7 +120,7 @@ let create ?(config=Config.default) ?(error_handler=default_error_handler) reque
   let rec reader = lazy (Reader.request handler)
   and handler request request_body =
     let reqd =
-      Reqd.create error_handler request request_body (Lazy.force reader) writer response_body_buffer
+      Reqd.create error_handler request request_body (Lazy.force reader) writer response_body_buffer config.proxy
     in
     let call_handler = Queue.is_empty request_queue in
     Queue.push reqd request_queue;
@@ -134,6 +135,7 @@ let create ?(config=Config.default) ?(error_handler=default_error_handler) reque
     ; error_handler   = error_handler
     ; request_queue
     ; error_code = No_error
+    ; config
     }
   in
   Lazy.force t
@@ -191,7 +193,7 @@ let set_error_and_handle ?request t error =
               | None -> `GET
               | Some (request: Request.t) -> request.meth
             in
-            match Response.body_length ~request_method response with
+            match Response.body_length ~proxy:t.config.proxy ~request_method response with
             | `Fixed _ | `Close_delimited | `Chunked as encoding -> encoding
             | `Error (`Bad_gateway | `Internal_server_error) ->
               failwith "httpun.Server_connection.error_handler: invalid response body length"