Skip to content

Commit 3492a4a

Browse files
authored
fix: garbled data after committing frame header (#68)
* fix: garbled data after committing frame header * fix tests
1 parent c126014 commit 3492a4a

File tree

3 files changed

+69
-70
lines changed

3 files changed

+69
-70
lines changed

lib/parse.ml

+55-50
Original file line numberDiff line numberDiff line change
@@ -1,62 +1,67 @@
11
type t =
2-
{ headers: Bigstringaf.t
3-
; payload: Payload.t
2+
{ payload_length: int
3+
; is_fin: bool
4+
; mask: int32 option
5+
; payload: Payload.t
6+
; opcode: Websocket.Opcode.t
47
}
58

6-
let is_fin t =
7-
let bits = Bigstringaf.unsafe_get t.headers 0 |> Char.code in
9+
let is_fin headers =
10+
let bits = Bigstringaf.unsafe_get headers 0 |> Char.code in
811
bits land (1 lsl 7) = 1 lsl 7
912
;;
1013

11-
let rsv t =
14+
(* let rsv t =
1215
let bits = Bigstringaf.unsafe_get t.headers 0 |> Char.code in
1316
(bits lsr 4) land 0b0111
1417
;;
18+
*)
1519

16-
let opcode t =
17-
let bits = Bigstringaf.unsafe_get t.headers 0 |> Char.code in
20+
let opcode headers =
21+
let bits = Bigstringaf.unsafe_get headers 0 |> Char.code in
1822
bits land 0b1111 |> Websocket.Opcode.unsafe_of_code
1923
;;
2024

2125
let payload_length_of_headers headers =
2226
let bits = Bigstringaf.unsafe_get headers 1 |> Char.code in
2327
let length = bits land 0b01111111 in
24-
if length = 126 then Bigstringaf.unsafe_get_int16_be headers 2 else
25-
(* This is technically unsafe, but if somebody's asking us to read 2^63
26-
* bytes, then we're already screwed. *)
27-
if length = 127 then Bigstringaf.unsafe_get_int64_be headers 2 |> Int64.to_int else
28-
length
28+
if length <= 125 then
29+
(* From RFC6455§5.3:
30+
* The length of the "Payload data", in bytes: if 0-125, that is the
31+
* payload length. *)
32+
length
33+
else if length = 126 then
34+
(* From RFC6455§5.3:
35+
* If 126, the following 2 bytes interpreted as a 16-bit unsigned integer
36+
* are the payload length. *)
37+
Bigstringaf.unsafe_get_int16_be headers 2
38+
else begin
39+
assert (length = 127);
40+
(* This is technically unsafe, but if somebody's asking us to read 2^63
41+
* bytes, then we're already screwed. *)
42+
Bigstringaf.unsafe_get_int64_be headers 2 |> Int64.to_int
43+
end
2944
;;
3045

31-
let payload_length t = payload_length_of_headers t.headers
32-
33-
let has_mask t =
34-
let bits = Bigstringaf.unsafe_get t.headers 1 |> Char.code in
35-
bits land (1 lsl 7) = 1 lsl 7
36-
;;
46+
(* let payload_length t = payload_length_of_headers t.headers *)
3747

38-
let mask t =
39-
if not (has_mask t)
48+
let mask =
49+
let has_mask headers =
50+
let bits = Bigstringaf.unsafe_get headers 1 |> Char.code in
51+
(bits land 0b1000_0000) = 0b1000_0000
52+
in
53+
let mask_exn headers =
54+
let bits = Bigstringaf.unsafe_get headers 1 |> Char.code in
55+
if bits = 254 then Bigstringaf.unsafe_get_int32_be headers 4 else
56+
if bits = 255 then Bigstringaf.unsafe_get_int32_be headers 10 else
57+
if bits >= 127 then Bigstringaf.unsafe_get_int32_be headers 2 else
58+
failwith "Frame.mask_exn: no mask present"
59+
in
60+
fun headers ->
61+
if not (has_mask headers)
4062
then None
4163
else
42-
Some (
43-
let bits = Bigstringaf.unsafe_get t.headers 1 |> Char.code in
44-
if bits = 254 then Bigstringaf.unsafe_get_int32_be t.headers 4 else
45-
if bits = 255 then Bigstringaf.unsafe_get_int32_be t.headers 10 else
46-
Bigstringaf.unsafe_get_int32_be t.headers 2)
47-
;;
48-
49-
let mask_exn t =
50-
let bits = Bigstringaf.unsafe_get t.headers 1 |> Char.code in
51-
if bits = 254 then Bigstringaf.unsafe_get_int32_be t.headers 4 else
52-
if bits = 255 then Bigstringaf.unsafe_get_int32_be t.headers 10 else
53-
if bits >= 127 then Bigstringaf.unsafe_get_int32_be t.headers 2 else
54-
failwith "Frame.mask_exn: no mask present"
55-
;;
56-
57-
let length t =
58-
let payload_length = payload_length t in
59-
Bigstringaf.length t.headers + payload_length
64+
Some (mask_exn headers)
6065
;;
6166

6267
let payload_offset_of_bits bits =
@@ -85,25 +90,23 @@ let parse_headers =
8590
let payload_parser t =
8691
let open Angstrom in
8792
let unmask t bs ~src_off =
88-
match mask t with
93+
match t.mask with
8994
| None -> bs
9095
| Some mask ->
9196
Websocket.Frame.apply_mask mask bs ~src_off;
9297
bs
9398
in
9499
let finish payload =
95-
let open Angstrom in
96100
Payload.close payload;
97101
commit
98102
in
99103
let schedule_size ~src_off payload n =
100-
let open Angstrom in
101104
begin if Payload.is_closed payload
102105
then advance n
103106
else take_bigstring n >>| fun bs ->
104107
let faraday = Payload.unsafe_faraday payload in
105108
Faraday.schedule_bigstring faraday (unmask ~src_off t bs)
106-
end *> commit
109+
end <* commit
107110
in
108111
let read_exact =
109112
let rec read_exact src_off n =
@@ -117,24 +120,28 @@ let payload_parser t =
117120
available >>= fun m ->
118121
let m' = (min m n) in
119122
let n' = n - m' in
120-
schedule_size ~src_off t.payload m' >>= fun () -> read_exact (src_off + m') n'
123+
schedule_size ~src_off t.payload m'
124+
>>= fun () -> read_exact (src_off + m') n'
121125
in
122126
fun n -> read_exact 0 n
123127
in
124-
read_exact (payload_length t)
128+
read_exact t.payload_length
125129
>>= fun () -> finish t.payload
126130
;;
127131

128132
let frame ~buf =
129133
let open Angstrom in
130134
parse_headers
131135
>>| fun headers ->
132-
let len = payload_length_of_headers headers in
133-
let payload = match len with
136+
let payload_length = payload_length_of_headers headers
137+
and is_fin = is_fin headers
138+
and opcode = opcode headers
139+
and mask = mask headers in
140+
let payload = match payload_length with
134141
| 0 -> Payload.empty
135142
| _ -> Payload.create buf
136143
in
137-
{ headers; payload }
144+
{ is_fin; opcode; mask; payload_length; payload }
138145
;;
139146

140147
module Reader = struct
@@ -157,9 +164,7 @@ module Reader = struct
157164
skip_many
158165
(frame ~buf <* commit >>= fun frame ->
159166
let payload = frame.payload in
160-
let is_fin = is_fin frame in
161-
let opcode = opcode frame in
162-
let len = payload_length frame in
167+
let { is_fin; opcode; payload_length = len; _ } = frame in
163168
frame_handler ~opcode ~is_fin ~len payload;
164169
payload_parser frame)
165170
in

lib/websocket.ml

-1
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,6 @@ module Frame = struct
169169
in
170170
for i = off to off + len - 1 do
171171
let j = (i + src_off - off) mod 4 in
172-
(* let j = (i - off) mod 4 in *)
173172
let c = Bigstringaf.unsafe_get bs i |> Char.code in
174173
let c = c lxor Int32.(logand (shift_right mask (8 * (3 - j))) 0xffl |> to_int) in
175174
Bigstringaf.unsafe_set bs i (Char.unsafe_chr c)

lib_test/test_httpun_ws.ml

+14-19
Original file line numberDiff line numberDiff line change
@@ -35,18 +35,15 @@ module Websocket = struct
3535

3636
let test_parsing_ping_frame () =
3737
let frame = parse_frame "\137\128\000\000\046\216" in
38-
Alcotest.check Testable.opcode "opcode" `Ping (Parse.opcode frame);
39-
Alcotest.(check bool) "has mask" true (Parse.has_mask frame);
40-
Alcotest.(check int32) "mask" 11992l (Parse.mask_exn frame);
41-
Alcotest.(check int) "payload_length" (Parse.payload_length frame) 0;
42-
Alcotest.(check int) "length" (Parse.length frame) 6
38+
Alcotest.check Testable.opcode "opcode" `Ping frame.opcode;
39+
Alcotest.(check (option int32)) "mask" (Some 11992l) frame.mask;
40+
Alcotest.(check int) "payload_length" 0 frame.payload_length
4341

4442
let test_parsing_close_frame () =
4543
let frame = parse_frame "\136\000" in
46-
Alcotest.check Testable.opcode "opcode" `Connection_close (Parse.opcode frame);
47-
Alcotest.(check int) "payload_length" (Parse.payload_length frame) 0;
48-
Alcotest.(check int) "length" (Parse.length frame) 2;
49-
Alcotest.(check bool) "is_fin" true (Parse.is_fin frame)
44+
Alcotest.check Testable.opcode "opcode" `Connection_close frame.opcode;
45+
Alcotest.(check int) "payload_length" 0 frame.payload_length;
46+
Alcotest.(check bool) "is_fin" true frame.is_fin
5047

5148
let read_payload frame =
5249
let rev_payload_chunks = ref [] in
@@ -60,23 +57,21 @@ module Websocket = struct
6057

6158
let test_parsing_text_frame () =
6259
let frame = parse_frame "\129\139\086\057\046\216\103\011\029\236\099\015\025\224\111\009\036" in
63-
Alcotest.check Testable.opcode "opcode" `Text (Parse.opcode frame);
64-
Alcotest.(check bool) "has mask" true (Parse.has_mask frame);
65-
Alcotest.(check int32) "mask" 1446588120l (Parse.mask_exn frame);
66-
Alcotest.(check int) "payload_length" (Parse.payload_length frame) 11;
67-
Alcotest.(check int) "length" (Parse.length frame) 17;
60+
Alcotest.check Testable.opcode "opcode" `Text frame.opcode;
61+
Alcotest.(check (option int32)) "mask" (Some 1446588120l) frame.mask;
62+
Alcotest.(check int) "payload_length" 11 frame.payload_length;
6863
let rev_payload_chunks = read_payload frame in
69-
Alcotest.(check bool) "is_fin" true (Parse.is_fin frame);
64+
Alcotest.(check bool) "is_fin" true frame.is_fin;
7065
Alcotest.(check (list string)) "payload" ["1234567890\n"] rev_payload_chunks
7166

7267

7368
let test_parsing_fin_bit () =
7469
let frame = parse_frame (serialize_frame ~is_fin:false "hello") in
75-
Alcotest.check Testable.opcode "opcode" `Text (Parse.opcode frame);
76-
Alcotest.(check bool) "is_fin" false (Parse.is_fin frame);
70+
Alcotest.check Testable.opcode "opcode" `Text frame.opcode;
71+
Alcotest.(check bool) "is_fin" false frame.is_fin;
7772
let frame = parse_frame (serialize_frame ~is_fin:true "hello") in
78-
Alcotest.check Testable.opcode "opcode" `Text (Parse.opcode frame);
79-
Alcotest.(check bool) "is_fin" true (Parse.is_fin frame);
73+
Alcotest.check Testable.opcode "opcode" `Text frame.opcode;
74+
Alcotest.(check bool) "is_fin" true frame.is_fin;
8075
let rev_payload_chunks = read_payload frame in
8176
Alcotest.(check (list string)) "payload" ["hello"] rev_payload_chunks
8277

0 commit comments

Comments
 (0)