1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
|
open Lwt.Infix
open Websocket
let src = Logs.Src.create "websocket.upgrade_connection"
let lwt_reporter () =
let buf_fmt ~like =
let b = Buffer.create 512 in
( Fmt.with_buffer ~like b,
fun () ->
let m = Buffer.contents b in
Buffer.reset b;
m )
in
let app, app_flush = buf_fmt ~like:Fmt.stdout in
let dst, dst_flush = buf_fmt ~like:Fmt.stderr in
let reporter = Logs_fmt.reporter ~app ~dst () in
let report src level ~over k msgf =
let k () =
let write () =
match level with
| Logs.App -> Lwt_io.write Lwt_io.stdout (app_flush ())
| _ -> Lwt_io.write Lwt_io.stderr (dst_flush ())
in
let unblock () =
over ();
Lwt.return_unit
in
Lwt.finalize write unblock |> Lwt.ignore_result;
k ()
in
reporter.Logs.report src level ~over:(fun () -> ()) k msgf
in
{ Logs.report }
let handler (_, conn) req body =
let open Frame in
Logs_lwt.app ~src (fun m ->
m "[CONN] %a" Sexplib.Sexp.pp (Cohttp.Connection.sexp_of_t conn))
>>= fun _ ->
let uri = Cohttp.Request.uri req in
match Uri.path uri with
| "/" ->
Logs_lwt.app ~src (fun m -> m "[PATH] /") >>= fun () ->
Cohttp_lwt_unix.Server.respond_string ~status:`OK
~body:
{|
<html>
<head>
<meta charset="utf-8">
<script src="//code.jquery.com/jquery-1.11.3.min.js"></script>
<script>
$(window).on('load', function(){
ws = new WebSocket('ws://localhost:7777/ws');
ws.onmessage = function(x) {
console.log(x.data);
var m = "<- Pong " + parseInt((x.data.substring(8)) - 1);
$('#msg').html("<p>" + x.data + "</p><p>" + m + "</p>");
ws.send(m);
};
});
</script>
</head>
<body>
<div id='msg'></div>
</body>
</html>
|}
()
>|= fun resp -> `Response resp
| "/ws" ->
Logs_lwt.app ~src (fun m -> m "[PATH] /ws") >>= fun () ->
Cohttp_lwt.Body.drain_body body >>= fun () ->
Websocket_cohttp_lwt.upgrade_connection req (fun { opcode; content; _ } ->
match opcode with
| Opcode.Close -> Logs.app ~src (fun m -> m "[RECV] CLOSE")
| _ -> Logs.app ~src (fun m -> m "[RECV] %s" content))
>>= fun (resp, frames_out_fn) ->
(* send a message to the client every second *)
let num_ref = ref 10 in
let rec go () =
if !num_ref = 0 then Logs_lwt.app ~src (fun m -> m "[INFO] Test done")
else
let msg = Printf.sprintf "-> Ping %d" !num_ref in
Logs_lwt.app ~src (fun m -> m "[SEND] %s" msg) >>= fun () ->
Lwt.wrap1 frames_out_fn @@ Some (Frame.create ~content:msg ())
>>= fun () ->
decr num_ref;
Lwt_unix.sleep 1. >>= go
in
Lwt.async go;
Lwt.return resp
| _ ->
Logs_lwt.app ~src (fun m -> m "[PATH] Catch-all") >>= fun () ->
Cohttp_lwt_unix.Server.respond_string ~status:`Not_found
~body:(Sexplib.Sexp.to_string_hum (Cohttp.Request.sexp_of_t req))
()
>|= fun resp -> `Response resp
let start_server ?tls port =
let conn_closed (_, c) =
Logs.app ~src (fun m ->
m "[SERV] connection %a closed" Sexplib.Sexp.pp
(Cohttp.Connection.sexp_of_t c))
in
Logs_lwt.app ~src (fun m -> m "[SERV] Listening for HTTP on port %d" port)
>>= fun () ->
let mode =
match tls with
| None ->
Logs.app ~src (fun m -> m "TCP mode selected");
`TCP (`Port port)
| Some (cert, key) ->
Logs.app ~src (fun m -> m "TLS mode selected");
`TLS (`Crt_file_path cert, `Key_file_path key, `No_password, `Port port)
in
Cohttp_lwt_unix.Server.create ~mode
(Cohttp_lwt_unix.Server.make_response_action ~callback:handler ~conn_closed
())
let () =
Logs.(set_reporter (lwt_reporter ()));
let port = ref 7777 in
let cert = ref "" in
let key = ref "" in
let speclist =
Arg.align
[
("-cert", Arg.Set_string cert, " cert file");
("-key", Arg.Set_string key, " key file");
( "-v",
Arg.Unit (fun () -> Logs.set_level (Some Info)),
" Set loglevel to info" );
( "-vv",
Arg.Unit (fun () -> Logs.set_level (Some Debug)),
" Set loglevel to debug" );
]
in
let anon_fun s =
match int_of_string_opt s with
| None -> invalid_arg "argument must be a port number"
| Some p -> port := p
in
let usage_msg = "Usage: " ^ Sys.argv.(0) ^ " <options> port\nOptions are:" in
Arg.parse speclist anon_fun usage_msg;
let tls =
match (!cert, !key) with
| "", _ | _, "" -> None
| cert, key -> Some (cert, key)
in
Lwt_main.run (start_server ?tls !port)
|