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

anmonteiro · Nov 30, 2024 · ec902ce · ec902ce
1 parent bd5ea5a
commit ec902ce
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 5 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,13 +73,15 @@ type t =
   ; mutable persistent      : bool
   ; mutable response_state  : Response_state.t
   ; mutable error_code      : [`Ok | error ]
+  ; proxy: bool
   }
 
 let create
   ~error_handler
   ~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
@@ -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"