File: upgrade_connection.ml

package info (click to toggle)
ocaml-websocket 2.17-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 288 kB
  • sloc: ml: 1,611; makefile: 7
file content (151 lines) | stat: -rw-r--r-- 4,922 bytes parent folder | download
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)