package websocket

  1. Overview
  2. Docs

Source file websocket.ml

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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
(*
 * 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

let b64_encoded_sha1sum s = Base64.encode_exn (Sha1.sha_1 s)
let websocket_uuid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

module Rng = struct
  let init () len = Mirage_crypto_rng.generate len
end

module Frame = struct
  module Opcode = struct
    type t =
      | Continuation
      | Text
      | Binary
      | Close
      | Ping
      | Pong
      | Ctrl of int
      | Nonctrl of int

    let to_string = function
      | Continuation -> "continuation"
      | Text -> "text"
      | Binary -> "binary"
      | Close -> "close"
      | Ping -> "ping"
      | Pong -> "pong"
      | Ctrl i -> "ctrl " ^ string_of_int i
      | Nonctrl i -> "nonctrl " ^ string_of_int i

    let pp ppf t = Format.fprintf ppf "%s" (to_string t)

    let of_enum = function
      | i when i < 0 || i > 0xf -> invalid_arg "Frame.Opcode.of_enum"
      | 0 -> Continuation
      | 1 -> Text
      | 2 -> Binary
      | 8 -> Close
      | 9 -> Ping
      | 10 -> Pong
      | i when i < 8 -> Nonctrl i
      | i -> Ctrl i

    let to_enum = function
      | Continuation -> 0
      | Text -> 1
      | Binary -> 2
      | Close -> 8
      | Ping -> 9
      | Pong -> 10
      | Ctrl i -> i
      | Nonctrl i -> i

    let is_ctrl opcode = to_enum opcode > 7
  end

  type t = {
    opcode : Opcode.t;
    extension : int;
    final : bool;
    content : string;
  }

  let pp ppf { opcode; extension; final; content } =
    Format.fprintf ppf "[%a (0x%x) (final=%b) %s]" Opcode.pp opcode extension
      final content

  let show t = Format.asprintf "%a" pp t

  let create ?(opcode = Opcode.Text) ?(extension = 0) ?(final = true)
      ?(content = "") () =
    { opcode; extension; final; content }

  let of_bytes ?opcode ?extension ?final content =
    let content = Bytes.to_string content in
    create ?opcode ?extension ?final ~content ()

  let close code =
    let content = Bytes.create 2 in
    EndianBytes.BigEndian.set_int16 content 0 code;
    of_bytes ~opcode:Opcode.Close content
end

let xor mask msg =
  for i = 0 to Bytes.length msg - 1 do
    (* masking msg to send *)
    Bytes.set msg i
      Char.(to_int mask.[i mod 4] lxor to_int (Bytes.get msg i) |> of_byte)
  done

let is_bit_set idx v = (v lsr idx) land 1 = 1
let set_bit v idx b = if b then v lor (1 lsl idx) else v land lnot (1 lsl idx)
let int_value shift len v = (v lsr shift) land ((1 lsl len) - 1)

let upgrade_present hs =
  Cohttp.Header.get_multi hs "connection" |> fun hs ->
  List.map (String.cuts ~sep:",") hs |> fun hs ->
  List.flatten hs |> fun hs ->
  List.map String.(fun h -> h |> String.Ascii.lowercase |> trim) hs
  |> List.mem "upgrade"

exception Protocol_error of string

let proto_error x = Format.kasprintf (fun x -> raise (Protocol_error x)) x

let check_origin ?(origin_mandatory = false) ~hosts =
  let pred origin_host =
    List.exists (fun h -> String.Ascii.lowercase h = origin_host) hosts
  in
  fun request ->
    let headers = request.Cohttp.Request.headers in
    match Cohttp.Header.get headers "origin" with
    | None -> not origin_mandatory
    | Some origin -> (
        let origin = Uri.of_string origin in
        match Uri.host origin with
        | None -> false
        | Some host ->
            (* host is already lowercased by Uri *)
            pred host)

let check_origin_with_host request =
  let headers = request.Cohttp.Request.headers in
  let host = Cohttp.Header.get headers "host" in
  match host with
  | None -> failwith "Missing host header" (* mandatory in http/1.1 *)
  | Some host ->
      (* remove port *)
      let hostname =
        match String.cut ~sep:":" host with None -> host | Some (h, _) -> h
      in
      check_origin ~hosts:[ hostname ] request

module type S = sig
  module IO : Cohttp.S.IO

  type mode = Client of (int -> string) | Server

  val make_read_frame :
    ?buf:Buffer.t -> mode:mode -> IO.ic -> IO.oc -> unit -> Frame.t IO.t

  val write_frame_to_buf : mode:mode -> Buffer.t -> Frame.t -> unit

  module Request :
    Cohttp.S.Http_io
      with type t = Cohttp.Request.t
       and type 'a IO.t = 'a IO.t
       and type IO.ic = IO.ic
       and type IO.oc = IO.oc

  module Response :
    Cohttp.S.Http_io
      with type t = Cohttp.Response.t
       and type 'a IO.t = 'a IO.t
       and type IO.ic = IO.ic
       and type IO.oc = IO.oc

  module Connected_client : sig
    type t

    val create :
      ?read_buf:Buffer.t ->
      ?write_buf:Buffer.t ->
      Cohttp.Request.t ->
      Conduit.endp ->
      IO.ic ->
      IO.oc ->
      t

    val make_standard : t -> t
    val send : t -> Frame.t -> unit IO.t
    val send_multiple : t -> Frame.t list -> unit IO.t
    val recv : t -> Frame.t IO.t
    val http_request : t -> Cohttp.Request.t
    val source : t -> Conduit.endp
  end
end

module Make (IO : Cohttp.S.IO) = struct
  open IO
  module IO = IO

  type mode = Client of (int -> string) | Server

  let is_client mode = mode <> Server

  let rec read_exactly ic remaining buf =
    read ic remaining >>= function
    | "" -> return None
    | s ->
        let recv_len = String.length s in
        Buffer.add_string buf s;
        if remaining - recv_len <= 0 then return @@ Some (Buffer.contents buf)
        else read_exactly ic (remaining - recv_len) buf

  let read_uint16 ic buf =
    read_exactly ic 2 buf >>= fun s ->
    match s with
    | None -> return None
    | Some s -> return @@ Some (EndianString.BigEndian.get_uint16 s 0)

  let read_int64 ic buf =
    read_exactly ic 8 buf >>= fun s ->
    match s with
    | None -> return None
    | Some s ->
        return @@ Some (Int64.to_int @@ EndianString.BigEndian.get_int64 s 0)

  let write_frame_to_buf ~mode buf fr =
    let scratch = Bytes.create 8 in
    let open Frame in
    let content = Bytes.of_string fr.content in
    let len = Bytes.length content in
    let opcode = Opcode.to_enum fr.opcode in
    let payload_len =
      match len with
      | n when n < 126 -> len
      | n when n < 1 lsl 16 -> 126
      | _ -> 127
    in
    let hdr = set_bit 0 15 fr.final in
    (* We do not support extensions for now *)
    let hdr = hdr lor (opcode lsl 8) in
    let hdr = set_bit hdr 7 (is_client mode) in
    let hdr = hdr lor payload_len in
    (* Payload len is guaranteed to fit in 7 bits *)
    EndianBytes.BigEndian.set_int16 scratch 0 hdr;
    Buffer.add_subbytes buf scratch 0 2;
    (match len with
    | n when n < 126 -> ()
    | n when n < 1 lsl 16 ->
        EndianBytes.BigEndian.set_int16 scratch 0 n;
        Buffer.add_subbytes buf scratch 0 2
    | n ->
        EndianBytes.BigEndian.set_int64 scratch 0 Int64.(of_int n);
        Buffer.add_subbytes buf scratch 0 8);
    (match mode with
    | Server -> ()
    | Client random_string ->
        let mask = random_string 4 in
        Buffer.add_string buf mask;
        if len > 0 then xor mask content);
    Buffer.add_bytes buf content

  let close_with_code mode buf oc code =
    Buffer.clear buf;
    write_frame_to_buf ~mode buf @@ Frame.close code;
    write oc @@ Buffer.contents buf

  let read_frame ic oc buf mode hdr =
    let hdr_part1 = EndianString.BigEndian.get_int8 hdr 0 in
    let hdr_part2 = EndianString.BigEndian.get_int8 hdr 1 in
    let final = is_bit_set 7 hdr_part1 in
    let extension = int_value 4 3 hdr_part1 in
    let opcode = int_value 0 4 hdr_part1 in
    let frame_masked = is_bit_set 7 hdr_part2 in
    let length = int_value 0 7 hdr_part2 in
    let opcode = Frame.Opcode.of_enum opcode in
    Buffer.clear buf;
    (match length with
    | i when i < 126 -> return i
    | 126 -> (
        read_uint16 ic buf >>= function
        | Some i -> return i
        | None -> return @@ -1)
    | 127 -> (
        read_int64 ic buf >>= function
        | Some i -> return i
        | None -> return @@ -1)
    | _ -> return @@ -1)
    >>= fun payload_len ->
    if payload_len = -1 then proto_error "payload len = %d" length
    else if extension <> 0 then
      close_with_code mode buf oc 1002 >>= fun () ->
      proto_error "unsupported extension"
    else if Frame.Opcode.is_ctrl opcode && payload_len > 125 then
      close_with_code mode buf oc 1002 >>= fun () ->
      proto_error "control frame too big"
    else
      (if frame_masked then (
       Buffer.clear buf;
       read_exactly ic 4 buf >>= function
       | None -> proto_error "could not read mask"
       | Some mask -> return mask)
      else return String.empty)
      >>= fun mask ->
      if payload_len = 0 then
        return @@ Frame.create ~opcode ~extension ~final ()
      else (
        Buffer.clear buf;
        read_exactly ic payload_len buf >>= fun payload ->
        match payload with
        | None -> proto_error "could not read payload (len=%d)" payload_len
        | Some payload ->
            let payload = Bytes.of_string payload in
            if frame_masked then xor mask payload;
            let frame = Frame.of_bytes ~opcode ~extension ~final payload in
            return frame)

  let make_read_frame ?(buf = Buffer.create 128) ~mode ic oc () =
    Buffer.clear buf;
    read_exactly ic 2 buf >>= function
    | None -> raise End_of_file
    | Some hdr -> read_frame ic oc buf mode hdr

  module Request = Cohttp.Request.Make (IO)
  module Response = Cohttp.Response.Make (IO)

  module Connected_client = struct
    type t = {
      buffer : Buffer.t;
      endp : Conduit.endp;
      ic : Request.IO.ic;
      oc : Request.IO.oc;
      http_request : Cohttp.Request.t;
      standard_frame_replies : bool;
      read_frame : unit -> Frame.t IO.t;
    }

    let source { endp; _ } = endp

    let create ?read_buf ?(write_buf = Buffer.create 128) http_request endp ic
        oc =
      let read_frame = make_read_frame ?buf:read_buf ~mode:Server ic oc in
      {
        buffer = write_buf;
        endp;
        ic;
        oc;
        http_request;
        standard_frame_replies = false;
        read_frame;
      }

    let send { buffer; oc; _ } frame =
      Buffer.clear buffer;
      write_frame_to_buf ~mode:Server buffer frame;
      IO.write oc @@ Buffer.contents buffer

    let send_multiple { buffer; oc; _ } frames =
      Buffer.clear buffer;
      List.iter (write_frame_to_buf ~mode:Server buffer) frames;
      IO.write oc @@ Buffer.contents buffer

    let standard_recv t =
      t.read_frame () >>= fun fr ->
      match fr.Frame.opcode with
      | Frame.Opcode.Ping ->
          send t @@ Frame.create ~opcode:Frame.Opcode.Pong () >>= fun () ->
          return fr
      | Frame.Opcode.Close ->
          (* Immediately echo and pass this last message to the user *)
          (if String.length fr.Frame.content >= 2 then
           send t
           @@ Frame.create ~opcode:Frame.Opcode.Close
                ~content:
                  String.(
                    sub ~start:0 ~stop:2 fr.Frame.content |> Sub.to_string)
                ()
          else send t @@ Frame.close 1000)
          >>= fun () -> return fr
      | _ -> return fr

    let recv t =
      if t.standard_frame_replies then standard_recv t else t.read_frame ()

    let http_request { http_request; _ } = http_request
    let make_standard t = { t with standard_frame_replies = true }
  end
end
OCaml

Innovation. Community. Security.