104 lines
3.8 KiB
OCaml
104 lines
3.8 KiB
OCaml
(* Copyright (C) 2023--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 Async_kernel
|
|
open Async_unix
|
|
open Core
|
|
open Tls_async
|
|
|
|
(* BEGIN IMPORT
|
|
*
|
|
* The following functions were copied from the tls_async.ml from the ocaml-tls
|
|
* project (https://github.com/mirleft/ocaml-tls) with minor changes.
|
|
* The original code is under BSD-2-Clause license. *)
|
|
let try_to_close t =
|
|
match%bind Session.close_tls t with
|
|
| Ok () -> return ()
|
|
| Error tls_close_error ->
|
|
Caqti_async.System.Log.err (fun m ->
|
|
m "Failed to close TLS session: %a" Error.pp tls_close_error)
|
|
;;
|
|
|
|
let pipe t =
|
|
let b_reader = Bytes.create 0x8000 in
|
|
let rec f_reader writer =
|
|
match%bind Session.read t b_reader with
|
|
| Ok 0 ->
|
|
Pipe.close writer;
|
|
return ()
|
|
| Ok len ->
|
|
let%bind () = Pipe.write writer (Stdlib.Bytes.sub_string b_reader 0 len) in
|
|
f_reader writer
|
|
| Error read_error ->
|
|
let%map () =
|
|
Caqti_async.System.Log.err (fun m ->
|
|
m "TLS read failed: %a" Error.pp read_error)
|
|
in
|
|
Pipe.close writer
|
|
in
|
|
let rec f_writer reader =
|
|
let%bind pipe_read = Pipe.read reader in
|
|
match pipe_read with
|
|
| `Ok s ->
|
|
(match%bind Session.writev t [ s ] with
|
|
| Ok () -> f_writer reader
|
|
| Error (_ : Error.t) -> try_to_close t)
|
|
| `Eof -> try_to_close t
|
|
in
|
|
Pipe.create_reader ~close_on_exception:false f_reader, Pipe.create_writer f_writer
|
|
;;
|
|
|
|
let upgrade_connection tls_session ((_ : Reader.t), outer_writer) =
|
|
let pipe_r, pipe_w = pipe tls_session in
|
|
let%bind inner_reader = Reader.of_pipe (Info.of_string "tls_reader") pipe_r in
|
|
let%map inner_writer, `Closed_and_flushed_downstream inner_cafd =
|
|
Writer.of_pipe (Info.of_string "tls_writer") pipe_w
|
|
in
|
|
Writer.set_raise_when_consumer_leaves inner_writer false;
|
|
let outer_cafd =
|
|
(* Ordering is important here to ensure no data is lost during the session shutdown *)
|
|
let%bind () = Writer.close_finished inner_writer in
|
|
let%bind () = inner_cafd in
|
|
let%bind () = try_to_close tls_session in
|
|
Writer.flushed outer_writer
|
|
in
|
|
tls_session, inner_reader, inner_writer, `Tls_closed_and_flushed_downstream outer_cafd
|
|
;;
|
|
(* END IMPORT *)
|
|
|
|
module Tls_provider = struct
|
|
type tls_config = Tls.Config.client
|
|
let tls_config_key = Caqti_tls.Config.client
|
|
|
|
let start_tls ~config ?host ((_outer_reader, outer_writer) as outer_rw) =
|
|
(match%bind Tls_async.Session.client_of_fd config ?host outer_rw with
|
|
| Error error ->
|
|
return (Error (Caqti_error.Msg (Error.to_string_hum error)))
|
|
| Ok session ->
|
|
let%map _, inner_reader, inner_writer,
|
|
`Tls_closed_and_flushed_downstream outer_cafd =
|
|
upgrade_connection session outer_rw
|
|
in
|
|
don't_wait_for begin
|
|
let%bind () = outer_cafd (* triggerd by closing inner_writer *) in
|
|
let%bind () = Writer.close outer_writer in
|
|
Reader.close inner_reader
|
|
end;
|
|
Ok (inner_reader, inner_writer))
|
|
end
|
|
|
|
let () = Caqti_async.System.Net.register_tls_provider (module Tls_provider)
|