Expose async tls_handler

async/examples/test_server.ml

@@ -4,7 +4,7 @@ open! Async
 let server_cert = "./certificates/server.pem"
 let server_key = "./certificates/server.key"
 
-let serve_tls port handler =
+let serve_tls ~low_level port handler =
   let%bind certificate =
     Tls_async.X509_async.Certificate.of_pem_file server_cert |> Deferred.Or_error.ok_exn
   in
@@ -21,10 +21,18 @@ let serve_tls port handler =
   in
   let where_to_listen = Tcp.Where_to_listen.of_port port in
   let on_handler_error = `Ignore in
-  Tls_async.listen ~on_handler_error config where_to_listen handler
+  if low_level then
+    Tcp.Server.create
+      ~on_handler_error
+      where_to_listen
+      (fun sa ->
+        printf !"connection establised from %{Socket.Address.Inet} starting TLS\n" sa;
+        Tls_async.upgrade_server_handler ~config (handler sa))
+  else
+    Tls_async.listen ~on_handler_error config where_to_listen handler
 ;;
 
-let test_server port =
+let test_server ~low_level port =
   let handler (_ : Socket.Address.Inet.t) (_ : Tls_async.Session.t) rd wr =
     let pipe = Reader.pipe rd in
     let rec read_from_pipe () =
@@ -35,17 +43,18 @@ let test_server port =
     in
     read_from_pipe ()
   in
-  serve_tls port handler
+  serve_tls ~low_level port handler
 ;;
 
 let cmd =
   let open Command.Let_syntax in
   Command.async
     ~summary:"test server"
-    (let%map_open port = anon ("PORT" %: int) in
+    (let%map_open port = anon ("PORT" %: int)
+    and low_level = flag "-low-level" no_arg ~doc:"set up Tcp.server directly" in
      fun () ->
        let open Deferred.Let_syntax in
-       let%bind server = test_server port in
+       let%bind server = test_server ~low_level port in
        Tcp.Server.close_finished server)
 ;;
 

async/tls_async.ml

@@ -65,6 +65,25 @@ let upgrade_client_reader_writer_to_tls ?host config rw =
   upgrade_connection tls_session rw |> Deferred.ok
 ;;
 
+type 'a io_handler = Reader.t -> Writer.t -> 'a Deferred.t
+type 'a tls_handler = Session.t -> 'a io_handler
+
+let upgrade_server_handler ~config handle_client outer_reader outer_writer =
+  let%bind ( tls_session
+           , inner_reader
+           , inner_writer
+           , `Tls_closed_and_flushed_downstream inner_cafd )
+    =
+     upgrade_server_reader_writer_to_tls config (outer_reader, outer_writer)
+     |> Deferred.Or_error.ok_exn
+  in
+  Monitor.protect
+    (fun () -> handle_client tls_session inner_reader inner_writer)
+    ~finally:(fun () ->
+      Deferred.all_unit
+        [ Reader.close inner_reader; Writer.close inner_writer; inner_cafd ])
+;;
+
 let listen
       ?buffer_age_limit
       ?max_connections
@@ -76,21 +95,6 @@ let listen
       where_to_listen
       handle_client
   =
-  let tls_handler sock outer_reader outer_writer =
-    let%bind ( tls_session
-             , inner_reader
-             , inner_writer
-             , `Tls_closed_and_flushed_downstream inner_cafd )
-      =
-      upgrade_server_reader_writer_to_tls config (outer_reader, outer_writer)
-      |> Deferred.Or_error.ok_exn
-    in
-    Monitor.protect
-      (fun () -> handle_client sock tls_session inner_reader inner_writer)
-      ~finally:(fun () ->
-        Deferred.all_unit
-          [ Reader.close inner_reader; Writer.close inner_writer; inner_cafd ])
-  in
   Tcp.Server.create
     ?buffer_age_limit
     ?max_connections
@@ -99,7 +103,8 @@ let listen
     ?socket
     ~on_handler_error
     where_to_listen
-    tls_handler
+    (fun sock ->
+      upgrade_server_handler ~config (handle_client sock))
 ;;
 
 let connect

async/tls_async.mli

@@ -24,6 +24,23 @@ val listen
   -> ('address -> Session.t -> Reader.t -> Writer.t -> unit Deferred.t)
   -> ('address, 'listening_on) Tcp.Server.t Deferred.t
 
+type 'a io_handler = Reader.t -> Writer.t -> 'a Deferred.t
+type 'a tls_handler = Session.t -> 'a io_handler
+
+(** [upgrade_server_handler] is what [listen] calls to handle each client.
+    It is exposed so that low-level end-users of the library can use tls-async
+    inside of code that manages Tcp services directly.
+
+    The [tls_handler] argument will be called with the client Tls session,
+    reader and writer to be used for cleartext data.
+
+    The outer [reader] and [writer] will read encrypted data from and write
+    encrypted data to the connected socket. *)
+val upgrade_server_handler
+  : config:Tls.Config.server
+  -> 'a tls_handler
+  -> 'a io_handler
+
 (** [connect] behaves similarly to [Tcp.connect], exposing a cleartext reader and writer.
     Callers should ensure they close the [Writer.t] and wait for the [unit Deferred.t]
     returned by [`Closed_and_flushed_downstream] to completely shut down the TLS connection