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

@@ -0,0 +1,193 @@
+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 apply v f = f v
+let sigpipe = 13
+let () = Mirage_crypto_rng_unix.use_default ()
+let () = Printexc.record_backtrace true
+
+(*
+let () = Fmt_tty.setup_std_outputs ~style_renderer:`Ansi_tty ~utf_8:true ()
+let () = Logs.set_reporter (reporter Fmt.stdout)
+let () = Logs.set_level ~all:true (Some Logs.Debug)
+*)
+
+let () = Sys.set_signal sigpipe Sys.Signal_ignore
+let src = Logs.Src.create "simple-server"
+
+module Log = (val Logs.src_log src : Logs.LOG)
+module P = Paf_mirage.Make (Tcpip_stack_socket.V4V6.TCP)
+module Ke = Ke.Rke
+
+let getline queue =
+  let exists ~predicate queue =
+    let pos = ref 0 and res = ref (-1) in
+    Ke.iter
+      (fun chr ->
+        if predicate chr then res := !pos ;
+        incr pos)
+      queue ;
+    if !res = -1 then None else Some !res in
+  let blit src src_off dst dst_off len =
+    Bigstringaf.blit_to_bytes src ~src_off dst ~dst_off ~len in
+  match exists ~predicate:(( = ) '\n') queue with
+  | Some pos ->
+      let tmp = Bytes.create pos in
+      Ke.N.keep_exn queue ~blit ~length:Bytes.length ~off:0 ~len:pos tmp ;
+      Ke.N.shift_exn queue (pos + 1) ;
+      Some (Bytes.unsafe_to_string tmp)
+  | None -> None
+
+let http_large filename ?shutdown:_ (_ip, _port) ic oc =
+  let open H1 in
+  Body.Reader.close ic ;
+  let ic = open_in filename in
+  let tp = Bytes.create 0x1000 in
+  let rec go () =
+    match input ic tp 0 (Bytes.length tp) with
+    | 0 -> Body.Writer.close oc
+    | len ->
+        Body.Writer.write_string oc (Bytes.sub_string tp 0 len) ;
+        go ()
+    | exception End_of_file -> Body.Writer.close oc in
+  go () ;
+  close_in ic
+
+let http_ping_pong ?shutdown:_ (_ip, _port) ic oc =
+  let open H1 in
+  let open Lwt.Infix in
+  let closed = ref false and queue = Ke.create ~capacity:0x1000 Bigarray.char in
+
+  let blit src src_off dst dst_off len =
+    Bigstringaf.blit src ~src_off dst ~dst_off ~len in
+  let on_eof () = closed := true in
+  let rec on_read buf ~off ~len =
+    Ke.N.push queue ~blit ~length:Bigstringaf.length buf ~off ~len ;
+    Body.Reader.schedule_read ic ~on_eof ~on_read in
+  Body.Reader.schedule_read ic ~on_eof ~on_read ;
+  let rec go () =
+    match (!closed, getline queue) with
+    | false, None -> Lwt.pause () >>= go
+    | false, Some "ping" ->
+        Body.Writer.write_string oc "pong\n" ;
+        go ()
+    | false, Some "pong" ->
+        Body.Writer.write_string oc "ping\n" ;
+        go ()
+    | false, Some _line ->
+        Body.Writer.close oc ;
+        Lwt.return_unit
+    | true, _ ->
+        Body.Writer.close oc ;
+        Lwt.return_unit in
+  Lwt.async go
+
+let request_handler large ?shutdown _flow (ip, port) reqd =
+  let open H1 in
+  let request = Reqd.request reqd in
+  match request.Request.target with
+  | "/" ->
+      let headers = Headers.of_list [ ("transfer-encoding", "chunked") ] in
+      let response = Response.create ~headers `OK in
+      let oc = Reqd.respond_with_streaming reqd response in
+      http_ping_pong ?shutdown (ip, port) (Reqd.request_body reqd) oc
+  | "/ping" ->
+      let headers = Headers.of_list [ ("content-length", "4") ] in
+      let response = Response.create ~headers `OK in
+      Reqd.respond_with_string reqd response "pong" ;
+      Option.iter (apply ()) shutdown
+  | "/pong" ->
+      let headers = Headers.of_list [ ("content-length", "4") ] in
+      let response = Response.create ~headers `OK in
+      Reqd.respond_with_string reqd response "ping" ;
+      Option.iter (apply ()) shutdown
+  | "/large" ->
+      let headers = Headers.of_list [ ("transfer-encoding", "chunked") ] in
+      let response = Response.create ~headers `OK in
+      let oc = Reqd.respond_with_streaming reqd response in
+      http_large large ?shutdown (ip, port) (Reqd.request_body reqd) oc
+  | _ -> assert false
+
+let error_handler _ ?request:_ error _respond =
+  match error with
+  | `Exn _exn -> Printexc.print_backtrace stderr
+  | `Bad_gateway -> Fmt.epr "Got a bad gateway error.\n%!"
+  | `Bad_request -> Fmt.epr "Got a bad request error.\n%!"
+  | `Internal_server_error -> Fmt.epr "Got an internal server error.\n%!"
+
+let ( <.> ) f g x = f (g x)
+
+open Lwt.Infix
+
+let ( >>? ) x f =
+  x >>= function Ok x -> f x | Error _ as err -> Lwt.return err
+
+let fd_8080 = Unix.openfile "lock.8080" Unix.[ O_CREAT; O_RDWR ] 0o644
+let () = at_exit (fun () -> try Unix.close fd_8080 with _exn -> ())
+let fd_4343 = Unix.openfile "lock.4343" Unix.[ O_CREAT; O_RDWR ] 0o644
+let () = at_exit (fun () -> try Unix.close fd_4343 with _exn -> ())
+let unlock fd = Unix.lockf fd Unix.F_ULOCK 0
+
+let server_http large stack =
+  P.init ~port:8080 stack >>= fun service ->
+  let http = P.http_service ~error_handler (request_handler large) in
+  let (`Initialized th) = P.serve http service in
+  unlock fd_8080 ;
+  Log.debug (fun m -> m "HTTP server initialized.") ;
+  th
+
+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 server_https cert key large stack =
+  let cert = load_file cert in
+  let key = load_file key in
+  match
+    (X509.Certificate.decode_pem_multiple cert, X509.Private_key.decode_pem key)
+  with
+  | Ok certs, Ok (`RSA key) ->
+      let tls =
+        Result.get_ok
+          (Tls.Config.server ~certificates:(`Single (certs, `RSA key)) ()) in
+      P.init ~port:4343 stack >>= fun service ->
+      let https = P.https_service ~tls ~error_handler (request_handler large) in
+      let (`Initialized th) = P.serve https service in
+      unlock fd_4343 ;
+      th
+  | _ -> invalid_arg "Invalid certificate or key"
+
+let 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 run_http large = stack >|= Tcpip_stack_socket.V4V6.tcp >>= server_http large
+
+let run_https cert key large =
+  stack >|= Tcpip_stack_socket.V4V6.tcp >>= server_https cert key large
+
+let () =
+  match Sys.argv with
+  | [| _; "--with-tls"; cert; key; large |] ->
+      Lwt_main.run (run_https cert key large)
+  | [| _; large |] -> Lwt_main.run (run_http large)
+  | _ -> Fmt.epr "%s [--with-tls cert key] large\n%!" Sys.argv.(0)