<commit_before>open Lwt.Infix

let src =
  let src = Logs.Src.create "Uwt" ~doc:"Host interface based on Uwt" in
  Logs.Src.set_level src (Some Logs.Info);
  src

module Log = (val Logs.src_log src : Logs.LOG)

let default_read_buffer_size = 65536

let log_exception_continue description f =
  Lwt.catch
    (fun () -> f ())
    (fun e ->
       Log.err (fun f -> f "%s: caught %s" description (Printexc.to_string e));
       Lwt.return ()
    )

let make_sockaddr (ip, port) =
  Unix.ADDR_INET (Unix.inet_addr_of_string @@ Ipaddr.to_string ip, port)

let string_of_address (dst, dst_port) =
  Ipaddr.to_string dst ^ ":" ^ (string_of_int dst_port)

let sockaddr_of_address (dst, dst_port) =
  Unix.ADDR_INET(Unix.inet_addr_of_string @@ Ipaddr.to_string dst, dst_port)

module Common = struct
  (** FLOW boilerplate *)

  type 'a io = 'a Lwt.t
  type buffer = Cstruct.t
  type error = [`Msg of string]
  type write_error = [Mirage_flow.write_error | error]
  let pp_error ppf (`Msg x) = Fmt.string ppf x

  let pp_write_error ppf = function
  | #Mirage_flow.write_error as e -> Mirage_flow.pp_write_error ppf e
  | #error as e                   -> pp_error ppf e

  let errorf fmt = Fmt.kstrf (fun s -> Lwt_result.fail (`Msg s)) fmt

  let ip_port_of_sockaddr sockaddr =
    try match sockaddr with
    | Unix.ADDR_INET(ip, port) ->
      Some (Ipaddr.of_string @@ Unix.string_of_inet_addr ip, port)
    | _ -> None
    with _ -> None
end

module Sockets = struct

  let max_connections = ref None

  let set_max_connections x = max_connections := x

  let next_connection_idx =
    let idx = ref 0 in
    fun () ->
      let next = !idx in
      incr idx;
      next

  exception Too_many_connections

  let connection_table = Hashtbl.create 511
  let get_num_connections () = Hashtbl.length connection_table

  let connections () =
    let xs = Hashtbl.fold (fun _ c acc -> c :: acc) connection_table [] in
    Vfs.File.ro_of_string (String.concat "\n" xs)

  let register_connection_no_limit description =
    let idx = next_connection_idx () in
    Hashtbl.replace connection_table idx description;
    idx

  let register_connection =
    let last_error_log = ref 0. in
    fun description -> match !max_connections with
    | Some m when Hashtbl.length connection_table >= m ->
      let now = Unix.gettimeofday () in
      if (now -. !last_error_log) > 30. then begin
        (* Avoid hammering the logging system *)
        Log.warn (fun f ->
            f "exceeded maximum number of forwarded connections (%d)" m);
        last_error_log := now;
      end;
      Lwt.fail Too_many_connections
    | _ ->
      let idx = register_connection_no_limit description in
      Lwt.return idx

  let register_connection_noexn description =
    Lwt.catch (fun () -> register_connection description >>= fun idx -> Lwt.return (Some idx)) (fun _ -> Lwt.return None)

  let deregister_connection idx =
    if not(Hashtbl.mem connection_table idx) then begin
      Log.warn (fun f -> f "deregistered connection %d more than once" idx)
    end;
    Hashtbl.remove connection_table idx

  module Datagram = struct
    type address = Ipaddr.t * int

    module Udp = struct
      include Common

      type flow = {
        idx: int option;
        label: string;
        description: string;
        mutable fd: Uwt.Udp.t option;
        read_buffer_size: int;
        mutable already_read: Cstruct.t option;
        sockaddr: Unix.sockaddr;
        address: address;
      }

      type address = Ipaddr.t * int

      let string_of_flow t = Fmt.strf "udp -> %s" (string_of_address t.address)

      let of_fd
          ?idx ?(read_buffer_size = Constants.max_udp_length)
          ?(already_read = None) ~description sockaddr address fd
        =
        let label = match fst address with
        | Ipaddr.V4 _ -> "UDPv4"
        | Ipaddr.V6 _ -> "UDPv6" in
        { idx; label; description; fd = Some fd; read_buffer_size; already_read;
          sockaddr; address }

      let connect ?read_buffer_size address =
        let description = "udp:" ^ (string_of_address address) in
        register_connection description
        >>= fun idx ->
        let label, fd, addr =
          try match fst @@ address with
          | Ipaddr.V4 _ -> "UDPv4", Uwt.Udp.init_ipv4_exn (), Unix.inet_addr_any
          | Ipaddr.V6 _ -> "UDPv6", Uwt.Udp.init_ipv6_exn (), Unix.inet6_addr_any
          with e -> deregister_connection idx; raise e
        in
        Lwt.catch (fun () ->
            let sockaddr = make_sockaddr address in
            let result = Uwt.Udp.bind fd ~addr:(Unix.ADDR_INET(addr, 0)) () in
            if not(Uwt.Int_result.is_ok result) then begin
              let error = Uwt.Int_result.to_error result in
              Log.err (fun f ->
                  f "Socket.%s.connect(%s): %s" label
                    (string_of_address address) (Uwt.strerror error));
              Lwt.fail (Unix.Unix_error(Uwt.to_unix_error error, "bind", ""))
            end else
              Lwt_result.return
                (of_fd ~idx ?read_buffer_size ~description sockaddr address fd)
          ) (fun e ->
            deregister_connection idx;
            log_exception_continue "Udp.connect Uwt.Udp.close_wait"
              (fun () -> Uwt.Udp.close_wait fd)
            >>= fun () ->
            errorf "Socket.%s.connect %s: caught %a" label description Fmt.exn e
          )

      let rec read t = match t.fd, t.already_read with
      | None, _ -> Lwt.return (Ok `Eof)
      | Some _, Some data when Cstruct.len data > 0 ->
        t.already_read <- Some (Cstruct.sub data 0 0); (* next read is `Eof *)
        Lwt.return (Ok (`Data data))
      | Some _, Some _ ->
        Lwt.return (Ok `Eof)
      | Some fd, None ->
        let buf = Cstruct.create t.read_buffer_size in
        Lwt.catch (fun () ->
            Uwt.Udp.recv_ba ~pos:buf.Cstruct.off ~len:buf.Cstruct.len
              ~buf:buf.Cstruct.buffer fd
            >>= fun recv ->
            if recv.Uwt.Udp.is_partial then begin
              Log.err (fun f ->
                  f "Socket.%s.recvfrom: dropping partial response (buffer \
                     was %d bytes)" t.label (Cstruct.len buf));
              read t
            end else
              Lwt.return (Ok (`Data (Cstruct.sub buf 0 recv.Uwt.Udp.recv_len)))
          ) (function
          | Unix.Unix_error(e, _, _) when Uwt.of_unix_error e = Uwt.ECANCELED ->
            (* happens on normal timeout *)
            Lwt.return (Ok `Eof)
          | e ->
            Log.err (fun f ->
                f "Socket.%s.recvfrom: %s caught %s returning Eof"
                  t.label
                  (string_of_flow t)
                  (Printexc.to_string e)
              );
            Lwt.return (Ok `Eof)
          )

      let write t buf = match t.fd with
      | None -> Lwt.return (Error `Closed)
      | Some fd ->
        Lwt.catch (fun () ->
            Uwt.Udp.send_ba ~pos:buf.Cstruct.off ~len:buf.Cstruct.len
              ~buf:buf.Cstruct.buffer fd t.sockaddr
            >>= fun () -><commit_msg>Datagram.Udp.read: correct debug strings<commit_after>open Lwt.Infix

let src =
  let src = Logs.Src.create "Uwt" ~doc:"Host interface based on Uwt" in
  Logs.Src.set_level src (Some Logs.Info);
  src

module Log = (val Logs.src_log src : Logs.LOG)

let default_read_buffer_size = 65536

let log_exception_continue description f =
  Lwt.catch
    (fun () -> f ())
    (fun e ->
       Log.err (fun f -> f "%s: caught %s" description (Printexc.to_string e));
       Lwt.return ()
    )

let make_sockaddr (ip, port) =
  Unix.ADDR_INET (Unix.inet_addr_of_string @@ Ipaddr.to_string ip, port)

let string_of_address (dst, dst_port) =
  Ipaddr.to_string dst ^ ":" ^ (string_of_int dst_port)

let sockaddr_of_address (dst, dst_port) =
  Unix.ADDR_INET(Unix.inet_addr_of_string @@ Ipaddr.to_string dst, dst_port)

module Common = struct
  (** FLOW boilerplate *)

  type 'a io = 'a Lwt.t
  type buffer = Cstruct.t
  type error = [`Msg of string]
  type write_error = [Mirage_flow.write_error | error]
  let pp_error ppf (`Msg x) = Fmt.string ppf x

  let pp_write_error ppf = function
  | #Mirage_flow.write_error as e -> Mirage_flow.pp_write_error ppf e
  | #error as e                   -> pp_error ppf e

  let errorf fmt = Fmt.kstrf (fun s -> Lwt_result.fail (`Msg s)) fmt

  let ip_port_of_sockaddr sockaddr =
    try match sockaddr with
    | Unix.ADDR_INET(ip, port) ->
      Some (Ipaddr.of_string @@ Unix.string_of_inet_addr ip, port)
    | _ -> None
    with _ -> None
end

module Sockets = struct

  let max_connections = ref None

  let set_max_connections x = max_connections := x

  let next_connection_idx =
    let idx = ref 0 in
    fun () ->
      let next = !idx in
      incr idx;
      next

  exception Too_many_connections

  let connection_table = Hashtbl.create 511
  let get_num_connections () = Hashtbl.length connection_table

  let connections () =
    let xs = Hashtbl.fold (fun _ c acc -> c :: acc) connection_table [] in
    Vfs.File.ro_of_string (String.concat "\n" xs)

  let register_connection_no_limit description =
    let idx = next_connection_idx () in
    Hashtbl.replace connection_table idx description;
    idx

  let register_connection =
    let last_error_log = ref 0. in
    fun description -> match !max_connections with
    | Some m when Hashtbl.length connection_table >= m ->
      let now = Unix.gettimeofday () in
      if (now -. !last_error_log) > 30. then begin
        (* Avoid hammering the logging system *)
        Log.warn (fun f ->
            f "exceeded maximum number of forwarded connections (%d)" m);
        last_error_log := now;
      end;
      Lwt.fail Too_many_connections
    | _ ->
      let idx = register_connection_no_limit description in
      Lwt.return idx

  let register_connection_noexn description =
    Lwt.catch (fun () -> register_connection description >>= fun idx -> Lwt.return (Some idx)) (fun _ -> Lwt.return None)

  let deregister_connection idx =
    if not(Hashtbl.mem connection_table idx) then begin
      Log.warn (fun f -> f "deregistered connection %d more than once" idx)
    end;
    Hashtbl.remove connection_table idx

  module Datagram = struct
    type address = Ipaddr.t * int

    module Udp = struct
      include Common

      type flow = {
        idx: int option;
        label: string;
        description: string;
        mutable fd: Uwt.Udp.t option;
        read_buffer_size: int;
        mutable already_read: Cstruct.t option;
        sockaddr: Unix.sockaddr;
        address: address;
      }

      type address = Ipaddr.t * int

      let string_of_flow t = Fmt.strf "udp -> %s" (string_of_address t.address)

      let of_fd
          ?idx ?(read_buffer_size = Constants.max_udp_length)
          ?(already_read = None) ~description sockaddr address fd
        =
        let label = match fst address with
        | Ipaddr.V4 _ -> "UDPv4"
        | Ipaddr.V6 _ -> "UDPv6" in
        { idx; label; description; fd = Some fd; read_buffer_size; already_read;
          sockaddr; address }

      let connect ?read_buffer_size address =
        let description = "udp:" ^ (string_of_address address) in
        register_connection description
        >>= fun idx ->
        let label, fd, addr =
          try match fst @@ address with
          | Ipaddr.V4 _ -> "UDPv4", Uwt.Udp.init_ipv4_exn (), Unix.inet_addr_any
          | Ipaddr.V6 _ -> "UDPv6", Uwt.Udp.init_ipv6_exn (), Unix.inet6_addr_any
          with e -> deregister_connection idx; raise e
        in
        Lwt.catch (fun () ->
            let sockaddr = make_sockaddr address in
            let result = Uwt.Udp.bind fd ~addr:(Unix.ADDR_INET(addr, 0)) () in
            if not(Uwt.Int_result.is_ok result) then begin
              let error = Uwt.Int_result.to_error result in
              Log.err (fun f ->
                  f "Socket.%s.connect(%s): %s" label
                    (string_of_address address) (Uwt.strerror error));
              Lwt.fail (Unix.Unix_error(Uwt.to_unix_error error, "bind", ""))
            end else
              Lwt_result.return
                (of_fd ~idx ?read_buffer_size ~description sockaddr address fd)
          ) (fun e ->
            deregister_connection idx;
            log_exception_continue "Udp.connect Uwt.Udp.close_wait"
              (fun () -> Uwt.Udp.close_wait fd)
            >>= fun () ->
            errorf "Socket.%s.connect %s: caught %a" label description Fmt.exn e
          )

      let rec read t = match t.fd, t.already_read with
      | None, _ -> Lwt.return (Ok `Eof)
      | Some _, Some data when Cstruct.len data > 0 ->
        t.already_read <- Some (Cstruct.sub data 0 0); (* next read is `Eof *)
        Lwt.return (Ok (`Data data))
      | Some _, Some _ ->
        Lwt.return (Ok `Eof)
      | Some fd, None ->
        let buf = Cstruct.create t.read_buffer_size in
        Lwt.catch (fun () ->
            Uwt.Udp.recv_ba ~pos:buf.Cstruct.off ~len:buf.Cstruct.len
              ~buf:buf.Cstruct.buffer fd
            >>= fun recv ->
            if recv.Uwt.Udp.is_partial then begin
              Log.err (fun f ->
                  f "Socket.%s.read: dropping partial response (buffer \
                     was %d bytes)" t.label (Cstruct.len buf));
              read t
            end else
              Lwt.return (Ok (`Data (Cstruct.sub buf 0 recv.Uwt.Udp.recv_len)))
          ) (function
          | Unix.Unix_error(e, _, _) when Uwt.of_unix_error e = Uwt.ECANCELED ->
            (* happens on normal timeout *)
            Lwt.return (Ok `Eof)
          | e ->
            Log.err (fun f ->
                f "Socket.%s.read: %s caught %s returning Eof"
                  t.label
                  (string_of_flow t)
                  (Printexc.to_string e)
              );
            Lwt.return (Ok `Eof)
          )

      let write t buf = match t.fd with
      | None -> Lwt.return (Error `Closed)
      | Some fd ->
        Lwt.catch (fun () ->
            Uwt.Udp.send_ba ~pos:buf.Cstruct.off ~len:buf.Cstruct.len
              ~buf:buf.Cstruct.buffer fd t.sockaddr
            >>= fun () ->
{'fraction_non_alphanumeric': 0.1017921654742509, 'fraction_numerical': 0.003966848480555359, 'mean_word_length': 2.062472885032538, 'pattern_counts': {'":': 2, '<': 5, '<?xml version=': 0, '>': 119, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6497012', 'n_tokens_mistral': 4737, 'n_tokens_neox': 4550, 'n_words': 1269}