refactor: track proxy in config for whether this is a proxy connection

anmonteiro · Nov 30, 2024 · 5ff647d · 5ff647d
1 parent 7de12df
commit 5ff647d
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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"