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
open CCFun
open Interface'
type opcode =
| Continuation
| Text
| Binary
| Reserved_non_control of int
| Close
| Ping
| Pong
| Reserved_control of int [@@deriving show]
type masking_key = char * char * char * char [@@deriving show]
type data = [ `Masked of masking_key * string | `Plain of string ] [@@deriving show]
type t =
{ fin: bool
; rsvd: int
; opcode: opcode
; data : data
} [@@deriving show]
let opcode_of_int = function
| 0 -> Continuation
| 1 -> Text
| 2 -> Binary
| x when x >= 3 && x <= 7 -> Reserved_non_control x
| 8 -> Close
| 9 -> Ping
| 10 -> Pong
| x -> Reserved_control x
let int_of_opcode = function
| Continuation -> 0
| Text -> 1
| Binary -> 2
| Close -> 8
| Ping -> 9
| Pong -> 10
| Reserved_non_control x
| Reserved_control x -> x
let guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
let to_bitstring = Bitstring.bitstring_of_string
let perform_mask key data =
String.mapi
(fun i c -> (int_of_char c) lxor (List.nth key (i mod 4)) |> char_of_int)
data
let unmask (`Masked ((b0, b1, b2, b3), data)) =
let key = List.map int_of_char [b0; b1; b2; b3] in
let unmasked_data = perform_mask key data in
`Plain unmasked_data
let mask (`Plain data) =
let key = CCList.init 4 (fun _ -> Random.int 256) in
let masked_data = perform_mask key data in
let key = List.map char_of_int key in
let b0 = List.nth key 0 in
let b1 = List.nth key 1 in
let b2 = List.nth key 2 in
let b3 = List.nth key 3 in
`Masked ((b0, b1, b2, b3), masked_data)
let of_string ~mask:m opcode s =
let data = match m with
| true -> mask (`Plain s)
| false -> `Plain s in
{ fin = true
; rsvd = 0
; opcode = opcode
; data
}
let close ~mask:m status =
let%bitstring data = {| status : 2*8 |} in
let s = data |> Bitstring.string_of_bitstring in
let data = match m with
| true -> mask (`Plain s)
| false -> `Plain s in
{ fin = true
; rsvd = 0
; opcode = Close
; data
}
module Make(Io: Io.S)(M_result : Monad_result.S with type E.t = string and type +'a u = 'a Io.M.t) = struct
open M_result
let parse_fin_rsvd_opcode = function%bitstring
| {| fin : 1
; rsvd : 3
; opcode : 4 |} ->
(fin, rsvd, opcode |> opcode_of_int) |> return
| {| _ |} ->
fail "Malformed_frame"
let parse_mask_len ic = function%bitstring
| {| mask : 1
; len1 : 7 |} ->
(match len1 with
| x when x <= 125 ->
return (mask, Int64.of_int len1)
| 126 ->
Io.read_exactly ic 2 |> lift >>= flip lift_opt (fail "Malformed_frame")
>|= (fun s ->
let len = (s.[1] |> int_of_char) + ((s.[0] |> int_of_char) lsl 8) in
mask, Int64.of_int len)
| 127 ->
Io.read_int64 ic |> lift >>= flip lift_opt (fail "Malformed_frame")
>>= (fun len -> return (mask, len))
| _ ->
fail "Malformed_frame")
| {| _ |} ->
fail "Malformed_frame"
let parse_plain_data len = function%bitstring
| {| data : len lsl 3 : string |} -> `Plain data |> return
let parse_masked_data len = function%bitstring
| {| key : 4*8 : string; data : len lsl 3 : string |} ->
let _list_key = [key.[0]; key.[1]; key.[2]; key.[3]] |> List.map int_of_char in
`Masked
( (key.[0] , key.[1] , key.[2] , key.[3])
, data ) |> return
let read_frame ic =
Io.read_exactly ic 1 |> lift >>= (flip lift_opt (fail "Malformed_frame"))
>|= to_bitstring
>>= parse_fin_rsvd_opcode
>>= fun (fin, rsvd, opcode) ->
Io.read_exactly ic 1 |> lift >>= (flip lift_opt (fail "Malformed_frame"))
>|= to_bitstring
>>= parse_mask_len ic
>>= fun (mask, len) ->
let intlen = Int64.to_int len in
(if mask then
Io.read_exactly ic (4+intlen) |> lift >>= (flip lift_opt (fail "Malformed_frame"))
>|= to_bitstring
>>= parse_masked_data intlen
else
Io.read_exactly ic intlen |> lift >>= (flip lift_opt (fail "Malformed_frame"))
>|= to_bitstring
>>= parse_plain_data intlen)
>|= fun data ->
{ fin
; rsvd
; opcode
; data = data
}
let write_frame oc frame =
let masked, mask_data, data_len = match frame.data with
| `Plain s -> false, s, String.length s
| `Masked ((b0, b1, b2, b3), s) ->
let len = String.length s in
let b = Buffer.create (4 + len) in
List.iter (Buffer.add_char b) [b0; b1; b2; b3];
true, Buffer.contents b, len in
let f = match data_len with
| x when x <= 125 ->
let%bitstring f =
{| frame.fin : 1
; frame.rsvd : 3
; frame.opcode |> int_of_opcode : 4
; masked : 1
; data_len : 7
; mask_data : data_len lsl 3 : string
|} in f
| x when x <= 65535 ->
let%bitstring f =
{| frame.fin : 1
; frame.rsvd : 3
; frame.opcode |> int_of_opcode : 4
; masked : 1
; 126 : 7
; data_len : 2*8
; mask_data : data_len lsl 3 : string
|} in f
| _ ->
let%bitstring f = {| frame.fin : 1
; frame.rsvd : 3
; frame.opcode |> int_of_opcode : 4
; masked : 1
; 127 : 7
; Int64.of_int data_len : 8*8
; mask_data : data_len lsl 3 : string
|} in f in
f |> Bitstring.string_of_bitstring |> Io.write oc
end