unikernel/duniverse/paf-le-chien/test/test_alpn.ml

@@ -0,0 +1,212 @@
+open Lwt.Infix
+
+let ( <.> ) f g x = f (g x)
+let ( >>? ) = Lwt_result.bind
+
+let reporter ppf =
+  let report src level ~over k msgf =
+    let k _ =
+      over () ;
+      k () in
+    let with_metadata header _tags k ppf fmt =
+      Format.kfprintf k ppf
+        ("%a[%a]: " ^^ fmt ^^ "\n%!")
+        Logs_fmt.pp_header (level, header)
+        Fmt.(styled `Magenta string)
+        (Logs.Src.name src) in
+    msgf @@ fun ?header ?tags fmt -> with_metadata header tags k ppf fmt in
+  { Logs.report }
+
+let () = Fmt_tty.setup_std_outputs ~style_renderer:`Ansi_tty ~utf_8:true ()
+let () = Logs.set_reporter (reporter Fmt.stderr)
+let () = Logs.set_level ~all:true (Some Logs.Debug)
+let () = Mirage_crypto_rng_unix.use_default ()
+
+module P = Paf_mirage.Make (Tcpip_stack_socket.V4V6.TCP)
+
+let unix_stack () =
+  Tcpip_stack_socket.V4V6.UDP.connect ~ipv4_only:false ~ipv6_only:false
+    Ipaddr.V4.Prefix.global None
+  >>= fun udpv4 ->
+  Tcpip_stack_socket.V4V6.TCP.connect ~ipv4_only:false ~ipv6_only:false
+    Ipaddr.V4.Prefix.global None
+  >>= fun tcpv4 -> Tcpip_stack_socket.V4V6.connect udpv4 tcpv4
+
+let load_file filename =
+  let ic = open_in filename in
+  let ln = in_channel_length ic in
+  let rs = Bytes.create ln in
+  really_input ic rs 0 ln ;
+  close_in ic ;
+  Bytes.unsafe_to_string rs
+
+let tls =
+  let cert = load_file "server.pem" in
+  let key = load_file "server.key" in
+  match
+    (X509.Certificate.decode_pem_multiple cert, X509.Private_key.decode_pem key)
+  with
+  | Ok certs, Ok (`RSA key) ->
+      Result.get_ok
+        (Tls.Config.server ~alpn_protocols:[ "http/1.1"; "h2" ]
+           ~certificates:(`Single (certs, `RSA key))
+           ())
+  | _ -> invalid_arg "Invalid certificate or key"
+
+let alpn_of_tls_connection (_, flow) =
+  match P.TLS.epoch flow with
+  | Ok { Tls.Core.alpn_protocol; _ } ->
+      Fmt.epr ">>> alpn_protocol (server side): %a.\n%!"
+        Fmt.(option string)
+        alpn_protocol ;
+      alpn_protocol
+  | Error _ -> None
+
+let peer_of_tls_connection ((ipaddr, port), _) =
+  Fmt.str "%a:%d" Ipaddr.pp ipaddr port
+
+let injection =
+  let module R = (val Mimic.repr P.tls_protocol) in
+  fun (_, flow) -> R.T flow
+
+let port =
+  let v = ref 9999 in
+  fun () ->
+    incr v ;
+    !v
+
+let service handler () =
+  let info =
+    {
+      Alpn.alpn = alpn_of_tls_connection;
+      Alpn.peer = peer_of_tls_connection;
+      Alpn.injection;
+    } in
+  let handshake flow =
+    let edn = P.TCP.dst flow in
+    P.TLS.server_of_flow tls flow >>= function
+    | Ok flow -> Lwt.return_ok (edn, flow)
+    | Error err ->
+        Lwt.return_error (`Msg (Fmt.str "%a" P.TLS.pp_write_error err))
+  and close = P.close in
+  Alpn.service info handler handshake P.accept close
+
+module R = (val Mimic.repr P.tls_protocol)
+
+type version = HTTP_1_1 | HTTP_2_0
+
+let error_handler _ _protocol ?request:_ _error _response = ()
+
+let request_handler : type reqd headers request response ro wo.
+    _ ->
+    _ ->
+    _ ->
+    _ ->
+    reqd ->
+    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
+    unit =
+ fun wk_request wk _flow _edn _reqd -> function
+  | Alpn.HTTP_1_1 (module Reqd) ->
+      Lwt.wakeup_later wk_request HTTP_1_1 ;
+      Lwt.wakeup_later wk ()
+  | Alpn.H2 (module Reqd) ->
+      Lwt.wakeup_later wk_request HTTP_2_0 ;
+      Lwt.wakeup_later wk ()
+
+let server_handler wk_request wk =
+  {
+    Alpn.error = error_handler;
+    Alpn.request =
+      (fun flow edn reqd protocol ->
+        request_handler wk_request wk flow edn reqd protocol);
+  }
+
+let client ~ctx handler req =
+  Mimic.resolve ctx >>= function
+  | Error err -> Alcotest.failf "%a" Mimic.pp_error err
+  | Ok (R.T v as flow) -> (
+      let alpn =
+        match P.TLS.epoch v with
+        | Ok { Tls.Core.alpn_protocol; _ } -> alpn_protocol
+        | Error _ -> None in
+      Alpn.run ?alpn handler () req flow >>= function
+      | Ok body -> Lwt.return body
+      | Error err -> Alcotest.failf "%a" Mimic.pp_error err)
+  | Ok flow -> (
+      Alpn.run handler () req flow >>= function
+      | Ok body -> Lwt.return body
+      | Error err -> Alcotest.failf "%a" Mimic.pp_error err)
+
+let ctx_with_tls stack ~port tls =
+  let ipaddr = Ipaddr_unix.of_inet_addr Unix.inet_addr_loopback in
+  Mimic.add P.tls_edn (None, tls, stack, ipaddr, port) Mimic.empty
+
+let authenticator ?ip:_ ~host:_ _ = Ok None
+let apply v f = f v
+
+let fake_client_handler =
+  {
+    Alpn.error = (fun _ _protocol _error -> ());
+    Alpn.response = (fun _flow _edn _response _body _protocol -> ());
+  }
+
+let test01 =
+  Alcotest_lwt.test_case "http/1.1" `Quick @@ fun _sw () ->
+  let port = port () in
+  let stop = Lwt_switch.create () in
+  let th, wk = Lwt.wait () in
+  let request, wk_request = Lwt.wait () in
+  let service = service (server_handler wk_request wk) () in
+  let tls =
+    Result.get_ok
+      (Tls.Config.client ~authenticator ~alpn_protocols:[ "http/1.1" ] ()) in
+  let req = `V1 (H1.Request.create `GET "/") in
+  Lwt.both
+    ( unix_stack () >|= Tcpip_stack_socket.V4V6.tcp >>= fun stack ->
+      P.init ~port stack >>= fun t ->
+      P.serve ~stop service t |> fun (`Initialized th) ->
+      let ctx = ctx_with_tls stack ~port tls in
+      Lwt.both (client ~ctx fake_client_handler req) th )
+    (th >>= fun () -> Lwt_switch.turn_off stop)
+  >>= fun ((body, ()), ()) ->
+  request >>= fun request ->
+  match (request, body) with
+  | HTTP_1_1, Alpn.Response_HTTP_1_1 _ ->
+      Alcotest.(check pass) "http/1.1" () () ;
+      Lwt.return_unit
+  | _ -> Alcotest.failf "Unexpected version of HTTP"
+
+let close_body = function
+  | Alpn.Response_HTTP_1_1 _ as response -> response
+  | Alpn.Response_H2 (body, _) as response ->
+      H2.Body.Writer.close body ;
+      response
+
+let test02 =
+  Alcotest_lwt.test_case "h2" `Quick @@ fun _sw () ->
+  let port = port () in
+  let stop = Lwt_switch.create () in
+  let th, wk = Lwt.wait () in
+  let request, wk_request = Lwt.wait () in
+  let service = service (server_handler wk_request wk) () in
+  let tls =
+    Result.get_ok (Tls.Config.client ~authenticator ~alpn_protocols:[ "h2" ] ())
+  in
+  let req = `V2 (H2.Request.create ~scheme:"https" `GET "/") in
+  Lwt.both
+    ( unix_stack () >|= Tcpip_stack_socket.V4V6.tcp >>= fun stack ->
+      P.init ~port stack >>= fun t ->
+      P.serve ~stop service t |> fun (`Initialized th) ->
+      let ctx = ctx_with_tls stack ~port tls in
+      Lwt.both (client ~ctx fake_client_handler req >|= close_body) th )
+    (th >>= fun () -> Lwt_switch.turn_off stop)
+  >>= fun ((body, ()), ()) ->
+  request >>= fun request ->
+  match (request, body) with
+  | HTTP_2_0, Alpn.Response_H2 _ ->
+      Alcotest.(check pass) "h2" () () ;
+      Lwt.return_unit
+  | _ -> Alcotest.failf "Unexpected version of HTTP"
+
+let test () = Alcotest_lwt.run "alpn" [ ("alpn", [ test01; test02 ]) ]
+let () = Lwt_main.run (test ())