Source file websocket_lwt_unix.ml

(*
 * Copyright (c) 2012-2018 Vincent Bernardoff <vb@luminar.eu.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *)

open Astring
open Websocket
open Lwt.Infix
include Websocket.Make (Cohttp_lwt_unix.IO)

let section = Lwt_log.Section.make "websocket_lwt_unix"

exception HTTP_Error of string

let http_error msg = Lwt.fail (HTTP_Error msg)
let protocol_error msg = Lwt.fail (Protocol_error msg)

let set_tcp_nodelay flow =
  let open Conduit_lwt_unix in
  match flow with
  | TCP {fd; _} -> Lwt_unix.setsockopt fd Lwt_unix.TCP_NODELAY true
  | _ -> ()

let fail_unless eq f = if not eq then f () else Lwt.return_unit
let fail_if eq f = if eq then f () else Lwt.return_unit

let drain_handshake req ic oc nonce =
  Request.write (fun _writer -> Lwt.return ()) req oc
  >>= fun () ->
  Response.read ic
  >>= (function
        | `Ok r -> Lwt.return r
        | `Eof -> Lwt.fail End_of_file
        | `Invalid s -> Lwt.fail @@ Failure s )
  >>= fun response ->
  let open Cohttp in
  let status = Response.status response in
  let headers = Response.headers response in
  fail_if
    Code.(is_error @@ code_of_status status)
    (fun () -> http_error Code.(string_of_status status))
  >>= fun () ->
  fail_unless
    (Response.version response = `HTTP_1_1)
    (fun () -> protocol_error "wrong http version")
  >>= fun () ->
  fail_unless
    (status = `Switching_protocols)
    (fun () -> protocol_error "wrong status")
  >>= fun () ->
  ( match Header.get headers "upgrade" with
  | Some a when String.Ascii.lowercase a = "websocket" -> Lwt.return_unit
  | _ -> protocol_error "wrong upgrade" )
  >>= fun () ->
  fail_unless (upgrade_present headers) (fun () ->
      protocol_error "upgrade header not present" )
  >>= fun () ->
  match Header.get headers "sec-websocket-accept" with
  | Some accept when accept = b64_encoded_sha1sum (nonce ^ websocket_uuid) ->
      Lwt.return_unit
  | _ -> protocol_error "wrong accept"

let connect ctx client url nonce extra_headers =
  let open Cohttp in
  let headers =
    Header.add_list extra_headers
      [ ("Upgrade", "websocket"); ("Connection", "Upgrade");
        ("Sec-WebSocket-Key", nonce); ("Sec-WebSocket-Version", "13") ] in
  let req = Request.make ~headers url in
  Conduit_lwt_unix.connect ~ctx client
  >>= fun (flow, ic, oc) ->
  set_tcp_nodelay flow ;
  Lwt.catch
    (fun () -> drain_handshake req ic oc nonce)
    (fun exn -> Lwt_io.close ic >>= fun () -> Lwt.fail exn)
  >>= fun () ->
  Lwt_log.info_f ~section "Connected to %s" (Uri.to_string url)
  >>= fun () -> Lwt.return (ic, oc)

type conn =
  { read_frame: unit -> Frame.t Lwt.t;
    write_frame: Websocket.Frame.t -> unit Lwt.t;
    oc: Lwt_io.output_channel }

let read {read_frame; _} = read_frame ()
let write {write_frame; _} frame = write_frame frame
let close_transport {oc; _} = Lwt_io.close oc

let connect ?(extra_headers = Cohttp.Header.init ())
    ?(random_string = Websocket.Rng.init ())
    ?(ctx = Lazy.force Conduit_lwt_unix.default_ctx) ?buf client url =
  let nonce = Base64.encode_exn (random_string 16) in
  connect ctx client url nonce extra_headers
  >|= fun (ic, oc) ->
  let read_frame = make_read_frame ?buf ~mode:(Client random_string) ic oc in
  let read_frame () =
    Lwt.catch read_frame (fun exn ->
        Lwt.async (fun () -> Lwt_io.close ic) ;
        Lwt.fail exn ) in
  let buf = Buffer.create 128 in
  let write_frame frame =
    Buffer.clear buf ;
    Lwt.wrap2 (write_frame_to_buf ~mode:(Client random_string)) buf frame
    >>= fun () ->
    Lwt.catch
      (fun () -> Lwt_io.write oc (Buffer.contents buf))
      (fun exn ->
        Lwt.async (fun () -> Lwt_io.close oc) ;
        Lwt.fail exn ) in
  {read_frame; write_frame; oc}

let write_failed_response oc =
  let body = "403 Forbidden" in
  let body_len = String.length body |> Int64.of_int in
  let response =
    Cohttp.Response.make ~status:`Forbidden
      ~encoding:(Cohttp.Transfer.Fixed body_len) () in
  let open Response in
  write ~flush:true (fun writer -> write_body writer body) response oc

let server_fun ?read_buf ?write_buf check_request flow ic oc react =
  let read = function
    | `Ok r -> Lwt.return r
    | `Eof ->
        (* Remote endpoint closed connection. No further action necessary here. *)
        Lwt_log.info ~section "Remote endpoint closed connection"
        >>= fun () -> Lwt.fail End_of_file
    | `Invalid reason ->
        Lwt_log.info_f ~section "Invalid input from remote endpoint: %s" reason
        >>= fun () -> Lwt.fail @@ HTTP_Error reason in
  Request.read ic >>= read
  >>= fun request ->
  let meth = Cohttp.Request.meth request in
  let version = Cohttp.Request.version request in
  let headers = Cohttp.Request.headers request in
  let key = Cohttp.Header.get headers "sec-websocket-key" in
  ( match
      ( version,
        meth,
        Cohttp.Header.get headers "upgrade",
        key,
        upgrade_present headers,
        check_request request )
    with
  | `HTTP_1_1, `GET, Some up, Some key, true, true
    when String.Ascii.lowercase up = "websocket" ->
      Lwt.return key
  | _ ->
      write_failed_response oc
      >>= fun () -> Lwt.fail (Protocol_error "Bad headers") )
  >>= fun key ->
  let hash = key ^ websocket_uuid |> b64_encoded_sha1sum in
  let response_headers =
    Cohttp.Header.of_list
      [ ("Upgrade", "websocket"); ("Connection", "Upgrade");
        ("Sec-WebSocket-Accept", hash) ] in
  let response =
    Cohttp.Response.make ~status:`Switching_protocols
      ~encoding:Cohttp.Transfer.Unknown ~headers:response_headers () in
  Response.write (fun _writer -> Lwt.return_unit) response oc
  >>= fun () ->
  let client = Connected_client.create ?read_buf ?write_buf request flow ic oc in
  react client

let establish_server ?read_buf ?write_buf ?timeout ?stop
    ?(on_exn = fun exn -> !Lwt.async_exception_hook exn)
    ?(check_request = check_origin_with_host)
    ?(ctx = Lazy.force Conduit_lwt_unix.default_ctx) ~mode react =
  let module C = Cohttp in
  Conduit_lwt_unix.serve ~on_exn ?timeout ?stop ~ctx ~mode (fun flow ic oc ->
      set_tcp_nodelay flow ;
      Lwt.catch
        (fun () ->
          server_fun ?read_buf ?write_buf check_request
            (Conduit_lwt_unix.endp_of_flow flow)
            ic oc react )
        (function
          | End_of_file -> Lwt.return_unit
          | HTTP_Error _ -> Lwt.return_unit
          | exn -> Lwt.fail exn ) )

let mk_frame_stream recv =
  let f () =
    recv ()
    >>= fun fr ->
    match fr.Frame.opcode with
    | Frame.Opcode.Close -> Lwt.return_none
    | _ -> Lwt.return (Some fr) in
  Lwt_stream.from f

let establish_standard_server ?read_buf ?write_buf ?timeout ?stop ?on_exn
    ?check_request ?(ctx = Lazy.force Conduit_lwt_unix.default_ctx) ~mode react
    =
  let f client = react (Connected_client.make_standard client) in
  establish_server ?read_buf ?write_buf ?timeout ?stop ?on_exn ?check_request
    ~ctx ~mode f