unikernel/duniverse/ocaml-caqti/caqti-mirage/lib/caqti_mirage.ml

@@ -0,0 +1,260 @@
+(* Copyright (C) 2022--2024  Petter A. Urkedal <paurkedal@gmail.com>
+ *
+ * This library is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or (at your
+ * option) any later version, with the LGPL-3.0 Linking Exception.
+ *
+ * This library is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * and the LGPL-3.0 Linking Exception along with this library.  If not, see
+ * <http://www.gnu.org/licenses/> and <https://spdx.org>, respectively.
+ *)
+
+open Lwt.Infix
+open Caqti_platform
+
+module type SOCKET_OPS =
+  Caqti_platform.System_sig.SOCKET_OPS with type 'a fiber := 'a Lwt.t
+
+module Make
+  (STACK : Tcpip.Stack.V4V6)
+  (DNS : Dns_client_mirage.S) =
+struct
+  module TCP = STACK.TCP
+  module TLS = Tls_mirage.Make (TCP)
+  module TCP_channel = Mirage_channel.Make (TCP)
+  module TLS_channel = Mirage_channel.Make (TLS)
+
+  module System_core = struct
+    include Caqti_lwt.System_core
+
+    type stdenv = {
+      stack: STACK.t;
+      dns: DNS.t;
+    }
+  end
+
+  module Alarm = struct
+    type t = {cancel: unit -> unit}
+
+    let schedule ~sw ~stdenv:_ t f =
+      let t_now = Mtime_clock.now () in
+      let dt_ns =
+        if Mtime.is_later t ~than:t_now then 0L else
+        Mtime.Span.to_uint64_ns (Mtime.span t t_now)
+      in
+      let task = Mirage_sleep.ns dt_ns >|= f in
+      let hook =
+        Caqti_lwt.Switch.on_release_cancellable sw
+          (fun () -> Lwt.cancel task; Lwt.return_unit)
+      in
+      {cancel = (fun () -> Caqti_lwt.Switch.remove_hook hook; Lwt.cancel task)}
+
+    let unschedule alarm = alarm.cancel ()
+  end
+
+  module Pool = Caqti_platform.Pool.Make (System_core) (Alarm)
+
+  module System = struct
+    include System_core
+    module Pool = Pool
+
+    module Net = struct
+
+      module Sockaddr = struct
+        type t = [`Tcp of Ipaddr.t * int | `Unix of string]
+        let unix s = `Unix s
+        let tcp (host, port) = `Tcp (host, port)
+      end
+
+      let getaddrinfo_ipv4 dns host port =
+        let extract (_, ips) =
+          Ipaddr.V4.Set.elements ips
+            |> List.map (fun ip -> `Tcp (Ipaddr.V4 ip, port))
+        in
+        DNS.getaddrinfo dns Dns.Rr_map.A host >|= Result.map extract
+
+      let getaddrinfo_ipv6 dns host port =
+        let extract (_, ips) =
+          Ipaddr.V6.Set.elements ips
+            |> List.map (fun ip -> `Tcp (Ipaddr.V6 ip, port))
+        in
+        DNS.getaddrinfo dns Dns.Rr_map.Aaaa host >|= Result.map extract
+
+      let getaddrinfo ~stdenv:{stack; dns} host port =
+        let laddrs = STACK.IP.configured_ips (STACK.ip stack) in
+        (match
+          List.exists Ipaddr.(function V4 _ -> true | V6 _ -> false) laddrs,
+          List.exists Ipaddr.(function V4 _ -> false | V6 _ -> true) laddrs
+         with
+         | true, true ->
+            getaddrinfo_ipv4 dns host port >>= fun r4 ->
+            getaddrinfo_ipv6 dns host port >|= fun r6 ->
+            (match r4, r6 with
+             | Ok addrs4, Ok addrs6 -> Ok (addrs4 @ addrs6)
+             | Ok addrs, Error _ | Error _, Ok addrs -> Ok addrs
+             | Error (`Msg msg4), Error (`Msg msg6) ->
+                if String.equal msg4 msg6 then Error (`Msg msg4) else
+                Error (`Msg ("IPv4: " ^ msg4 ^ " IPv6: " ^ msg6)))
+         | true, false -> getaddrinfo_ipv4 dns host port
+         | false, true -> getaddrinfo_ipv6 dns host port
+         | false, false ->
+            Lwt.return (Error (`Msg "No IP address assigned to host.")))
+
+      let convert_io_exception = function
+       | Failure msg -> Some (Caqti_error.Msg msg) (* Channel.S.error *)
+       | _ -> None
+
+      module Make_stream_ops (Channel : Mirage_channel.S) = struct
+        type t = Channel.t
+
+        let output_char channel c =
+          Channel.write_char channel c;
+          Lwt.return_unit
+
+        let output_string channel s =
+          Channel.write_string channel s 0 (String.length s);
+          Lwt.return_unit
+
+        let flush channel =
+          Channel.flush channel >>= function
+           | Ok () -> Lwt.return_unit
+           | Error err ->
+              Lwt.fail_with (Format.asprintf "%a" Channel.pp_write_error err)
+
+        let input_char channel =
+          Channel.read_char channel >>= function
+           | Ok (`Data c) -> Lwt.return c
+           | Ok `Eof -> Lwt.fail End_of_file
+           | Error err ->
+              Lwt.fail_with (Format.asprintf "%a" Channel.pp_error err)
+
+        let really_input channel buf off len =
+          Channel.read_exactly ~len channel >>= function
+           | Ok (`Data bufs) ->
+              let content = Cstruct.copyv bufs in
+              Bytes.blit_string content 0 buf off len;
+              Lwt.return_unit
+           | Ok `Eof -> Lwt.fail End_of_file
+           | Error err ->
+              Lwt.fail_with (Format.asprintf "%a" Channel.pp_error err)
+
+        let close channel =
+          Channel.close channel >>= function
+           | Ok () -> Lwt.return_unit
+           | Error err ->
+              Lwt.fail_with (Format.asprintf "%a" Channel.pp_write_error err)
+      end
+
+      module TCP_stream_ops = Make_stream_ops (TCP_channel)
+      module TLS_stream_ops = Make_stream_ops (TLS_channel)
+
+      module Socket = struct
+        type t = V : {
+          tcp_flow: TCP_channel.flow option;
+          ops: (module SOCKET_OPS with type t = 'a);
+          channel: 'a;
+        } -> t
+
+        let output_char (V {ops = (module Ops); channel; _}) =
+          Ops.output_char channel
+        let output_string (V {ops = (module Ops); channel; _}) =
+          Ops.output_string channel
+        let flush (V {ops = (module Ops); channel; _}) =
+          Ops.flush channel
+        let input_char (V {ops = (module Ops); channel; _}) =
+          Ops.input_char channel
+        let really_input (V {ops = (module Ops); channel; _}) =
+          Ops.really_input channel
+        let close (V {ops = (module Ops); channel; _}) =
+          Ops.close channel
+      end
+
+      type tcp_flow = TCP_channel.flow
+      type tls_flow = Tls_flow : {
+        ops: (module SOCKET_OPS with type t = 'a);
+        channel: 'a;
+      } -> tls_flow
+
+      let connect_tcp ~sw:_ ~stdenv:{stack; _} sockaddr =
+        (match sockaddr with
+         | `Unix _ ->
+            Lwt.return_error
+              (Caqti_error.Msg "Unix sockets are not available under MirageOS.")
+         | `Tcp (ipaddr, port) ->
+            TCP.create_connection (STACK.tcp stack) (ipaddr, port) >|=
+            (function
+             | Ok flow ->
+                let channel = TCP_channel.create flow in
+                Ok (Socket.V {
+                  tcp_flow = Some flow;
+                  ops = (module TCP_stream_ops);
+                  channel;
+                })
+             | Error err ->
+                let msg = Format.asprintf "%a" TCP.pp_error err in
+                Error (Caqti_error.Msg msg)))
+
+      let tcp_flow_of_socket (Socket.V {tcp_flow; _}) = tcp_flow
+
+      let socket_of_tls_flow ~sw:_ (Tls_flow {ops; channel}) =
+        Socket.V {tcp_flow = None; ops; channel}
+
+      module type TLS_PROVIDER = Caqti_platform.System_sig.TLS_PROVIDER
+        with type 'a fiber := 'a Lwt.t
+         and type tcp_flow := tcp_flow
+         and type tls_flow := tls_flow
+
+      module Tls_provider = struct
+        type tls_config = Tls.Config.client
+
+        let tls_config_key = Caqti_tls.Config.client
+
+        let start_tls ~config ?host flow =
+          TLS.client_of_flow config ?host flow >|=
+          (function
+           | Ok tls_flow ->
+              Ok (Tls_flow {
+                ops = (module TLS_stream_ops);
+                channel = TLS_channel.create tls_flow;
+              })
+           | Error err ->
+              let msg = Format.asprintf "%a" TLS.pp_write_error err in
+              Error (Caqti_error.Msg msg))
+      end
+
+      let tls_providers_r : (module TLS_PROVIDER) list ref =
+        ref [(module Tls_provider : TLS_PROVIDER)]
+
+      let tls_providers _ = !tls_providers_r
+
+      let register_tls_provider p = tls_providers_r := p :: !tls_providers_r
+
+    end
+  end
+
+  module Loader = Caqti_platform.Driver_loader.Make (System)
+
+  include Connector.Make (System) (Pool) (Loader)
+
+  let connect
+        ?subst ?env ?config ?tweaks_version ?(sw = Caqti_lwt.Switch.eternal)
+        stack dns uri =
+    connect ?subst ?env ?config ?tweaks_version ~sw ~stdenv:{stack; dns} uri
+
+  let with_connection ?subst ?env ?config ?tweaks_version stack dns uri f =
+    with_connection ?subst ?env ?config ?tweaks_version ~stdenv:{stack; dns} uri f
+
+  let connect_pool
+        ?pool_config ?post_connect ?subst ?env ?config ?tweaks_version
+        ?(sw = Caqti_lwt.Switch.eternal) stack dns uri =
+    connect_pool
+      ?pool_config ?post_connect ?subst ?env ?config ?tweaks_version
+      ~sw ~stdenv:{stack; dns} uri
+
+end