From 8632cdce725a93a346adafa39ba41e0b811effc8 Mon Sep 17 00:00:00 2001
From: Antonio Nuno Monteiro <anmonteiro@gmail.com>
Date: Fri, 29 Nov 2024 18:26:34 -0800
Subject: [PATCH] refactor: track `proxy` in config for whether this is a proxy
 connection (#145)

---
 lib/config.ml            | 8 ++++++--
 lib/httpun.mli           | 1 +
 lib/reqd.ml              | 5 ++++-
 lib/server_connection.ml | 7 +++++--
 4 files changed, 16 insertions(+), 5 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 3b6f0c0..6dd4d5c 100644
--- a/lib/reqd.ml
+++ b/lib/reqd.ml
@@ -73,6 +73,7 @@ type t =
   ; mutable persistent      : bool
   ; mutable response_state  : Response_state.t
   ; mutable error_code      : [`Ok | error ]
+  ; proxy: bool
   }
 
 let create
@@ -80,6 +81,7 @@ let create
   ~reader
   ~writer
   ~response_body_buffer
+  ~proxy
   request
   request_body =
   { request
@@ -91,6 +93,7 @@ let create
   ; persistent              = Request.persistent_connection request
   ; response_state          = Waiting
   ; error_code              = `Ok
+  ; proxy
   }
 
 let request { request; _ } = request
@@ -148,7 +151,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 cafa21e..cbac9c2 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
@@ -124,6 +125,7 @@ let create ?(config=Config.default) ?(error_handler=default_error_handler) reque
         ~reader:(Lazy.force reader)
         ~writer
         ~response_body_buffer
+        ~proxy:config.proxy
         request
         request_body
     in
@@ -140,6 +142,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
@@ -197,7 +200,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"