hcs / lib / h1_client.ml

(** HTTP/1.1 Client implementation using h1.

    This module provides HTTP/1.1 client functionality built on Eio with
    connection pooling for high performance. *)

(** {1 Configuration} *)

type config = {
  (* Timeouts *)
  connect_timeout : float;  (** Connection timeout in seconds. Default: 30.0 *)
  read_timeout : float;  (** Read timeout in seconds. Default: 30.0 *)
  write_timeout : float;  (** Write timeout in seconds. Default: 30.0 *)
  (* Behavior *)
  follow_redirects : int option;
      (** Max redirects to follow. None = don't follow. Default: Some 10 *)
  (* Buffers *)
  buffer_size : int;  (** Read buffer size. Default: 16384 *)
  max_response_body : int64 option;
      (** Max response body size. None = unlimited. Default: None *)
  (* TLS *)
  tls : Tls_config.Client.t;  (** TLS configuration *)
  (* Headers *)
  default_headers : (string * string) list;
      (** Headers to add to every request *)
  (* Pool settings *)
  pool : Pool.config;  (** Connection pool configuration *)
}
(** Client configuration *)

let default_config =
  {
    connect_timeout = 30.0;
    read_timeout = 30.0;
    write_timeout = 30.0;
    follow_redirects = Some 10;
    buffer_size = 16384;
    max_response_body = None;
    tls = Tls_config.Client.default;
    default_headers = [ ("User-Agent", "hcs/" ^ Build_info.version) ];
    pool = Pool.default_config;
  }

(** {2 Config builders} *)

let with_timeout timeout config =
  { config with connect_timeout = timeout; read_timeout = timeout }

let with_connect_timeout timeout config =
  { config with connect_timeout = timeout }

let with_read_timeout timeout config = { config with read_timeout = timeout }
let with_write_timeout timeout config = { config with write_timeout = timeout }

let with_redirects max_redirects config =
  { config with follow_redirects = Some max_redirects }

let without_redirects config = { config with follow_redirects = None }
let with_buffer_size size config = { config with buffer_size = size }

let with_max_response_body max_size config =
  { config with max_response_body = Some max_size }

let with_tls tls config = { config with tls }
let with_insecure_tls config = { config with tls = Tls_config.Client.insecure }

let with_default_header name value config =
  { config with default_headers = (name, value) :: config.default_headers }

let with_default_headers headers config =
  { config with default_headers = headers @ config.default_headers }

let with_pool_config pool config = { config with pool }

let with_max_connections max_conn config =
  {
    config with
    pool = { config.pool with max_connections_per_host = max_conn };
  }

let with_idle_timeout timeout config =
  { config with pool = { config.pool with idle_timeout = timeout } }

(** Error type for client operations *)
type error =
  | Connection_failed of string
  | Tls_error of string
  | Timeout
  | Invalid_response of string
  | Too_many_redirects

type response = { status : H1.Status.t; headers : H1.Headers.t; body : string }
(** Response type *)

(** {1 Connection type for pooling} *)

type conn = {
  flow : Eio.Flow.two_way_ty Eio.Std.r;
  read_buffer : Bigstringaf.t;
  mutable alive : bool;
}
(** Internal connection state *)

(** {1 Client type} *)

type t = {
  config : config;
  pool : conn Pool.t;
  mutex : Eio.Mutex.t;
  (* Functions that capture the Eio environment *)
  connect : host:string -> port:int -> Eio.Flow.two_way_ty Eio.Std.r;
  resolve : string -> bool; (* returns true if host can be resolved *)
  now : unit -> float;
  with_timeout : 'a. float -> (unit -> 'a) -> 'a option;
}
(** HTTP/1.1 client with connection pooling. Uses closures to capture Eio
    environment. *)

(** {1 Internal helpers} *)

(** Write all IOVecs to the flow *)
let write_iovecs flow iovecs =
  let cstructs =
    List.map
      (fun iov ->
        Cstruct.of_bigarray ~off:iov.Httpun_types.IOVec.off
          ~len:iov.Httpun_types.IOVec.len iov.Httpun_types.IOVec.buffer)
      iovecs
  in
  Eio.Flow.write flow cstructs

(** Read from flow into bigstring buffer *)
let read_into_bigstring flow buf ~off ~len =
  let cs = Cstruct.of_bigarray ~off ~len buf in
  try
    let n = Eio.Flow.single_read flow cs in
    `Ok n
  with End_of_file -> `Eof

(** Close a connection safely *)
let close_conn c =
  if c.alive then begin
    c.alive <- false;
    (* Use Obj.magic to close - the flow supports close even if type doesn't show it *)
    try Eio.Flow.close (Obj.magic c.flow) with _ -> ()
  end

(** Perform an HTTP/1.1 request on a connected flow *)
let do_request ?(request_body = "") conn req =
  let flow = conn.flow in
  (* Set up response handling *)
  let response_received = Eio.Promise.create () in
  let body_buffer = Buffer.create 4096 in
  let resolved = ref false in
  let keep_alive = ref true in

  let resolve_once result =
    if not !resolved then begin
      resolved := true;
      Eio.Promise.resolve (snd response_received) result
    end
  in

  let response_handler resp body_reader =
    (* Check Connection header for keep-alive *)
    (match H1.Headers.get resp.H1.Response.headers "connection" with
    | Some v when String.lowercase_ascii v = "close" -> keep_alive := false
    | _ -> ());

    let rec read_body () =
      H1.Body.Reader.schedule_read body_reader
        ~on_eof:(fun () ->
          let body = Buffer.contents body_buffer in
          resolve_once
            (Ok
               ( {
                   status = resp.H1.Response.status;
                   headers = resp.headers;
                   body;
                 },
                 !keep_alive )))
        ~on_read:(fun buf ~off ~len ->
          Buffer.add_string body_buffer (Bigstringaf.substring buf ~off ~len);
          read_body ())
    in
    read_body ()
  in

  let error_handler err =
    keep_alive := false;
    conn.alive <- false;
    let msg =
      match err with
      | `Malformed_response s -> s
      | `Invalid_response_body_length _ -> "Invalid response body length"
      | `Exn exn -> Printexc.to_string exn
    in
    resolve_once (Error (Invalid_response msg))
  in

  (* Create the client connection *)
  let body_writer, h1_conn =
    H1.Client_connection.request req ~error_handler ~response_handler
  in

  (* Write request body if provided, then close *)
  if String.length request_body > 0 then begin
    H1.Body.Writer.write_string body_writer request_body;
    H1.Body.Writer.flush body_writer (fun () -> ())
  end;
  H1.Body.Writer.close body_writer;

  (* Buffer for reading - track unconsumed bytes between reads *)
  let read_buffer = conn.read_buffer in
  let read_buffer_size = Bigstringaf.length read_buffer in
  let unconsumed = ref 0 in

  (* Connection loop - handle both read and write operations *)
  let rec loop () =
    (* First, handle any pending writes *)
    let write_done =
      match H1.Client_connection.next_write_operation h1_conn with
      | `Write iovecs ->
          write_iovecs flow iovecs;
          let len =
            List.fold_left
              (fun acc iov -> acc + iov.Httpun_types.IOVec.len)
              0 iovecs
          in
          H1.Client_connection.report_write_result h1_conn (`Ok len);
          false
      | `Yield -> true
      | `Close _ -> true
    in

    (* Then handle reads *)
    let read_done =
      match H1.Client_connection.next_read_operation h1_conn with
      | `Read -> (
          let available = read_buffer_size - !unconsumed in
          match
            read_into_bigstring flow read_buffer ~off:!unconsumed ~len:available
          with
          | `Ok n ->
              let total = !unconsumed + n in
              let consumed =
                H1.Client_connection.read h1_conn read_buffer ~off:0 ~len:total
              in
              (* Shift unconsumed bytes to start of buffer *)
              let remaining = total - consumed in
              if remaining > 0 && consumed > 0 then
                Bigstringaf.blit read_buffer ~src_off:consumed read_buffer
                  ~dst_off:0 ~len:remaining;
              unconsumed := remaining;
              false
          | `Eof ->
              conn.alive <- false;
              let _ =
                H1.Client_connection.read_eof h1_conn read_buffer ~off:0
                  ~len:!unconsumed
              in
              true)
      | `Close -> true
    in

    (* Continue until both read and write are done *)
    if not (write_done && read_done) then loop ()
  in

  (try loop () with _ -> conn.alive <- false);
  Eio.Promise.await (fst response_received)

let create_connection t ~host ~port ~is_https =
  if not (t.resolve host) then
    Error (Connection_failed ("Cannot resolve host: " ^ host))
  else
    try
      let tcp_flow = t.connect ~host ~port in
      let flow =
        if is_https then
          match Tls_config.Client.to_tls_config t.config.tls ~host with
          | Error msg -> raise (Failure msg)
          | Ok tls_config ->
              let host_domain =
                match Domain_name.of_string host with
                | Ok dn -> (
                    match Domain_name.host dn with
                    | Ok h -> Some h
                    | Error _ -> None)
                | Error _ -> None
              in
              let tls_flow =
                Tls_eio.client_of_flow tls_config ?host:host_domain
                  (Obj.magic tcp_flow)
              in
              (Obj.magic tls_flow : Eio.Flow.two_way_ty Eio.Std.r)
        else tcp_flow
      in
      Ok
        {
          flow;
          read_buffer = Bigstringaf.create t.config.buffer_size;
          alive = true;
        }
    with
    | Tls_eio.Tls_failure failure ->
        Error (Tls_error (Tls_config.failure_to_string failure))
    | Failure msg -> Error (Tls_error msg)
    | exn -> Error (Connection_failed (Printexc.to_string exn))

(** Acquire a connection from pool or create new one *)
let acquire_connection t ~host ~port ~is_https =
  let key = Pool.make_key ~host ~port ~tls:is_https in
  Eio.Mutex.use_rw ~protect:true t.mutex (fun () ->
      (* Try to get idle connection from pool *)
      match Pool.acquire t.pool key ~now:(t.now ()) with
      | Some conn when conn.alive -> Ok conn
      | Some conn ->
          (* Dead connection, remove and create new *)
          Pool.remove t.pool key conn;
          create_connection t ~host ~port ~is_https
      | None ->
          (* No idle connection, create new *)
          create_connection t ~host ~port ~is_https)

(** Release connection back to pool *)
let release_connection t ~host ~port ~is_https conn ~keep_alive =
  let key = Pool.make_key ~host ~port ~tls:is_https in
  Eio.Mutex.use_rw ~protect:true t.mutex (fun () ->
      if keep_alive && conn.alive then begin
        (* Try to add back to pool *)
        if not (Pool.add t.pool key conn ~now:(t.now ())) then
          (* Pool full, just mark for release *)
          Pool.release t.pool key conn ~now:(t.now ())
      end
      else begin
        (* Connection not reusable *)
        Pool.remove t.pool key conn;
        close_conn conn
      end)

(** {1 Public API} *)

(** Create a new HTTP/1.1 client with connection pooling *)
let create ~sw ~net ~clock ?(config = default_config) () =
  let connect ~host ~port =
    let addrs = Eio.Net.getaddrinfo_stream net host in
    match addrs with
    | [] -> failwith ("Cannot resolve host: " ^ host)
    | addr_info :: _ ->
        let addr =
          match addr_info with
          | `Tcp (ip, _) -> `Tcp (ip, port)
          | `Unix p -> `Unix p
        in
        let flow = Eio.Net.connect ~sw net addr in
        (flow :> Eio.Flow.two_way_ty Eio.Std.r)
  in
  let resolve host =
    match Eio.Net.getaddrinfo_stream net host with
    | _ :: _ -> true
    | [] -> false
  in
  let now () = Eio.Time.now clock in
  let with_timeout timeout f =
    try Some (Eio.Time.with_timeout_exn clock timeout f)
    with Eio.Time.Timeout -> None
  in
  {
    config;
    pool = Pool.create ~config:config.pool ();
    mutex = Eio.Mutex.create ();
    connect;
    resolve;
    now;
    with_timeout;
  }

(** Close all connections and cleanup *)
let close t =
  Eio.Mutex.use_rw ~protect:true t.mutex (fun () ->
      Pool.close_all t.pool ~close:close_conn)

let merge_headers ~default_headers ~headers =
  let key name = String.lowercase_ascii name in
  let overridden = Hashtbl.create 8 in
  List.iter
    (fun (name, _value) -> Hashtbl.replace overridden (key name) ())
    headers;
  let default_headers =
    List.filter
      (fun (name, _value) -> not (Hashtbl.mem overridden (key name)))
      default_headers
  in
  default_headers @ headers

(** Perform a GET request using pooled connections *)
let get t url =
  let uri = Uri.of_string url in
  let scheme = Uri.scheme uri |> Option.value ~default:"http" in
  let is_https = String.equal scheme "https" in
  let host = Uri.host uri |> Option.value ~default:"localhost" in
  let default_port = if is_https then 443 else 80 in
  let port = Uri.port uri |> Option.value ~default:default_port in
  let path = Uri.path_and_query uri in
  let path = if path = "" then "/" else path in

  let total_timeout = t.config.connect_timeout +. t.config.read_timeout in
  match
    t.with_timeout total_timeout (fun () ->
        match acquire_connection t ~host ~port ~is_https with
        | Error e -> Error e
        | Ok conn -> (
            let base_headers =
              [ ("Host", host); ("Connection", "keep-alive") ]
            in
            let headers =
              if t.config.default_headers = [] then base_headers
              else if t.config.default_headers = default_config.default_headers
              then t.config.default_headers @ base_headers
              else
                merge_headers ~default_headers:t.config.default_headers
                  ~headers:base_headers
            in
            let req =
              H1.Request.create ~headers:(H1.Headers.of_list headers) `GET path
            in
            match do_request conn req with
            | Ok (resp, keep_alive) ->
                release_connection t ~host ~port ~is_https conn ~keep_alive;
                Ok resp
            | Error e ->
                release_connection t ~host ~port ~is_https conn
                  ~keep_alive:false;
                Error e))
  with
  | Some result -> result
  | None -> Error Timeout

(** Perform a POST request using pooled connections *)
let post t url ~body:request_body =
  let uri = Uri.of_string url in
  let scheme = Uri.scheme uri |> Option.value ~default:"http" in
  let is_https = String.equal scheme "https" in
  let host = Uri.host uri |> Option.value ~default:"localhost" in
  let default_port = if is_https then 443 else 80 in
  let port = Uri.port uri |> Option.value ~default:default_port in
  let path = Uri.path_and_query uri in
  let path = if path = "" then "/" else path in

  let total_timeout = t.config.connect_timeout +. t.config.read_timeout in
  match
    t.with_timeout total_timeout (fun () ->
        match acquire_connection t ~host ~port ~is_https with
        | Error e -> Error e
        | Ok conn -> (
            let content_length = String.length request_body in
            let base_headers =
              [
                ("Host", host);
                ("Connection", "keep-alive");
                ("Content-Length", string_of_int content_length);
              ]
            in
            let headers =
              if t.config.default_headers = [] then base_headers
              else if t.config.default_headers = default_config.default_headers
              then t.config.default_headers @ base_headers
              else
                merge_headers ~default_headers:t.config.default_headers
                  ~headers:base_headers
            in
            let req =
              H1.Request.create ~headers:(H1.Headers.of_list headers) `POST path
            in
            match do_request ~request_body conn req with
            | Ok (resp, keep_alive) ->
                release_connection t ~host ~port ~is_https conn ~keep_alive;
                Ok resp
            | Error e ->
                release_connection t ~host ~port ~is_https conn
                  ~keep_alive:false;
                Error e))
  with
  | Some result -> result
  | None -> Error Timeout

(** {1 Backward-compatible stateless API} *)

(** Perform a GET request (creates temporary client, no pooling benefit) *)
let get' ~sw ~net ~clock ?(config = default_config) url =
  let t = create ~sw ~net ~clock ~config () in
  let result = get t url in
  close t;
  result

(** Perform a POST request (creates temporary client, no pooling benefit) *)
let post' ~sw ~net ~clock ?(config = default_config) url ~body =
  let t = create ~sw ~net ~clock ~config () in
  let result = post t url ~body in
  close t;
  result