websocket.ml

(*
ocamlbuild \
        -pkg containers \
        -pkg lwt \
        -pkg yojson \
        -pkg conduit.lwt-unix \
        -pkg nocrypto \
        -pkg nocrypto.unix \
        -pkg websocket.lwt \
        -cflags "-w A-4-40-41-42-44" \
        -cflags "-warn-error A-26-39-32-27-45-48-52-58"\
        -no-hygiene \
        -use-ocamlfind \
        -j 3 \
        -tag strict_sequence \
        -tag safe_string \
        -tag thread \
        -tag bin_annot \
        -tag short_paths \
        -tag 'color(always)' \
        -tag 'optimize(3)' \
        websocket.native --
*)

open Lwt.Infix

module Json = Yojson.Safe
module Frame = Websocket_lwt.Frame

let make_frame content = Frame.create ~content ~final:true ~opcode:Frame.Opcode.Text ()

let main () =
  let random_generator = Nocrypto.Rng.create (module Nocrypto.Rng.Generators.Fortuna) in
  Nocrypto_entropy_unix.initialize ();
  Nocrypto_entropy_unix.reseed random_generator;
  let ctx = Conduit_lwt_unix.default_ctx in
  let mode = `TCP (`Port 3000) in
  let random_string = Rng.nocrypto ~g:random_generator in
  let clients = CCVector.create () in
  let msg_count = ref 0 in
  let rec print () =
    Lwt_io.printf "%d %d\n" !msg_count (CCVector.length clients) >>= fun () ->
    Lwt_unix.sleep 0.3 >>= fun () ->
    print ()
  in
  Lwt.ignore_result (print ());
  let handle _id _http_request recv send =
    CCVector.push clients send;
    let rec loop () =
      recv () >>= fun frame ->
      msg_count := !msg_count + 1;
      let request = frame.Frame.content in
      match Json.from_string request with
      | `Assoc ["type", `String "broadcast"; "payload", payload] ->
        let broadcast_response = Json.to_string
            (`Assoc ["type", `String "broadcast"; "payload", payload]) in
        let broadcast_frame = make_frame broadcast_response in
        CCVector.iter begin fun send ->
          Lwt.ignore_result (send broadcast_frame)
        end clients;
        let content = Json.to_string
            (`Assoc ["type", `String "broadcastResult"; "payload", payload]) in
        let frame = make_frame content in
        Lwt.ignore_result (send frame);
        loop ()
      | `Assoc ["type", `String "echo"; "payload", payload] ->
        let content = Json.to_string
            (`Assoc ["type", `String "echo"; "payload", payload]) in
        let frame = make_frame content in
        Lwt.ignore_result (send frame);
        loop ()
      | _ ->
        assert false
    in
    loop ()
  in
  Websocket_lwt.establish_server ~ctx ~mode ~random_string handle

let () = Lwt_main.run (main ())