Skip to content

Commit

Permalink
fix: garbled data after committing frame header (#68)
Browse files Browse the repository at this point in the history 
* fix: garbled data after committing frame header

* fix tests
  • Loading branch information
@anmonteiro
anmonteiro authored Aug 1, 2024
1 parent c126014 commit 3492a4a
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 70 deletions.
105 changes: 55 additions & 50 deletions lib/parse.ml
View file
Original file line number Diff line number Diff line change
@@ -1,62 +1,67 @@
type t =
{ headers: Bigstringaf.t
; payload: Payload.t
{ payload_length: int
; is_fin: bool
; mask: int32 option
; payload: Payload.t
; opcode: Websocket.Opcode.t
}

let is_fin t =
let bits = Bigstringaf.unsafe_get t.headers 0 |> Char.code in
let is_fin headers =
let bits = Bigstringaf.unsafe_get headers 0 |> Char.code in
bits land (1 lsl 7) = 1 lsl 7
;;

let rsv t =
(* let rsv t =
let bits = Bigstringaf.unsafe_get t.headers 0 |> Char.code in
(bits lsr 4) land 0b0111
;;
*)

let opcode t =
let bits = Bigstringaf.unsafe_get t.headers 0 |> Char.code in
let opcode headers =
let bits = Bigstringaf.unsafe_get headers 0 |> Char.code in
bits land 0b1111 |> Websocket.Opcode.unsafe_of_code
;;

let payload_length_of_headers headers =
let bits = Bigstringaf.unsafe_get headers 1 |> Char.code in
let length = bits land 0b01111111 in
if length = 126 then Bigstringaf.unsafe_get_int16_be headers 2 else
(* This is technically unsafe, but if somebody's asking us to read 2^63
* bytes, then we're already screwed. *)
if length = 127 then Bigstringaf.unsafe_get_int64_be headers 2 |> Int64.to_int else
length
if length <= 125 then
(* From RFC6455§5.3:
* The length of the "Payload data", in bytes: if 0-125, that is the
* payload length. *)
length
else if length = 126 then
(* From RFC6455§5.3:
* If 126, the following 2 bytes interpreted as a 16-bit unsigned integer
* are the payload length. *)
Bigstringaf.unsafe_get_int16_be headers 2
else begin
assert (length = 127);
(* This is technically unsafe, but if somebody's asking us to read 2^63
* bytes, then we're already screwed. *)
Bigstringaf.unsafe_get_int64_be headers 2 |> Int64.to_int
end
;;

let payload_length t = payload_length_of_headers t.headers

let has_mask t =
let bits = Bigstringaf.unsafe_get t.headers 1 |> Char.code in
bits land (1 lsl 7) = 1 lsl 7
;;
(* let payload_length t = payload_length_of_headers t.headers *)

let mask t =
if not (has_mask t)
let mask =
let has_mask headers =
let bits = Bigstringaf.unsafe_get headers 1 |> Char.code in
(bits land 0b1000_0000) = 0b1000_0000
in
let mask_exn headers =
let bits = Bigstringaf.unsafe_get headers 1 |> Char.code in
if bits = 254 then Bigstringaf.unsafe_get_int32_be headers 4 else
if bits = 255 then Bigstringaf.unsafe_get_int32_be headers 10 else
if bits >= 127 then Bigstringaf.unsafe_get_int32_be headers 2 else
failwith "Frame.mask_exn: no mask present"
in
fun headers ->
if not (has_mask headers)
then None
else
Some (
let bits = Bigstringaf.unsafe_get t.headers 1 |> Char.code in
if bits = 254 then Bigstringaf.unsafe_get_int32_be t.headers 4 else
if bits = 255 then Bigstringaf.unsafe_get_int32_be t.headers 10 else
Bigstringaf.unsafe_get_int32_be t.headers 2)
;;

let mask_exn t =
let bits = Bigstringaf.unsafe_get t.headers 1 |> Char.code in
if bits = 254 then Bigstringaf.unsafe_get_int32_be t.headers 4 else
if bits = 255 then Bigstringaf.unsafe_get_int32_be t.headers 10 else
if bits >= 127 then Bigstringaf.unsafe_get_int32_be t.headers 2 else
failwith "Frame.mask_exn: no mask present"
;;

let length t =
let payload_length = payload_length t in
Bigstringaf.length t.headers + payload_length
Some (mask_exn headers)
;;

let payload_offset_of_bits bits =
Expand Down Expand Up @@ -85,25 +90,23 @@ let parse_headers =
let payload_parser t =
let open Angstrom in
let unmask t bs ~src_off =
match mask t with
match t.mask with
| None -> bs
| Some mask ->
Websocket.Frame.apply_mask mask bs ~src_off;
bs
in
let finish payload =
let open Angstrom in
Payload.close payload;
commit
in
let schedule_size ~src_off payload n =
let open Angstrom in
begin if Payload.is_closed payload
then advance n
else take_bigstring n >>| fun bs ->
let faraday = Payload.unsafe_faraday payload in
Faraday.schedule_bigstring faraday (unmask ~src_off t bs)
end *> commit
end <* commit
in
let read_exact =
let rec read_exact src_off n =
Expand All @@ -117,24 +120,28 @@ let payload_parser t =
available >>= fun m ->
let m' = (min m n) in
let n' = n - m' in
schedule_size ~src_off t.payload m' >>= fun () -> read_exact (src_off + m') n'
schedule_size ~src_off t.payload m'
>>= fun () -> read_exact (src_off + m') n'
in
fun n -> read_exact 0 n
in
read_exact (payload_length t)
read_exact t.payload_length
>>= fun () -> finish t.payload
;;

let frame ~buf =
let open Angstrom in
parse_headers
>>| fun headers ->
let len = payload_length_of_headers headers in
let payload = match len with
let payload_length = payload_length_of_headers headers
and is_fin = is_fin headers
and opcode = opcode headers
and mask = mask headers in
let payload = match payload_length with
| 0 -> Payload.empty
| _ -> Payload.create buf
in
{ headers; payload }
{ is_fin; opcode; mask; payload_length; payload }
;;

module Reader = struct
Expand All @@ -157,9 +164,7 @@ module Reader = struct
skip_many
(frame ~buf <* commit >>= fun frame ->
let payload = frame.payload in
let is_fin = is_fin frame in
let opcode = opcode frame in
let len = payload_length frame in
let { is_fin; opcode; payload_length = len; _ } = frame in
frame_handler ~opcode ~is_fin ~len payload;
payload_parser frame)
in
Expand Down
1 change: 0 additions & 1 deletion lib/websocket.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,6 @@ module Frame = struct
in
for i = off to off + len - 1 do
let j = (i + src_off - off) mod 4 in
(* let j = (i - off) mod 4 in *)
let c = Bigstringaf.unsafe_get bs i |> Char.code in
let c = c lxor Int32.(logand (shift_right mask (8 * (3 - j))) 0xffl |> to_int) in
Bigstringaf.unsafe_set bs i (Char.unsafe_chr c)
Expand Down
33 changes: 14 additions & 19 deletions lib_test/test_httpun_ws.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,18 +35,15 @@ module Websocket = struct

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

let test_parsing_close_frame () =
let frame = parse_frame "\136\000" in
Alcotest.check Testable.opcode "opcode" `Connection_close (Parse.opcode frame);
Alcotest.(check int) "payload_length" (Parse.payload_length frame) 0;
Alcotest.(check int) "length" (Parse.length frame) 2;
Alcotest.(check bool) "is_fin" true (Parse.is_fin frame)
Alcotest.check Testable.opcode "opcode" `Connection_close frame.opcode;
Alcotest.(check int) "payload_length" 0 frame.payload_length;
Alcotest.(check bool) "is_fin" true frame.is_fin

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

let test_parsing_text_frame () =
let frame = parse_frame "\129\139\086\057\046\216\103\011\029\236\099\015\025\224\111\009\036" in
Alcotest.check Testable.opcode "opcode" `Text (Parse.opcode frame);
Alcotest.(check bool) "has mask" true (Parse.has_mask frame);
Alcotest.(check int32) "mask" 1446588120l (Parse.mask_exn frame);
Alcotest.(check int) "payload_length" (Parse.payload_length frame) 11;
Alcotest.(check int) "length" (Parse.length frame) 17;
Alcotest.check Testable.opcode "opcode" `Text frame.opcode;
Alcotest.(check (option int32)) "mask" (Some 1446588120l) frame.mask;
Alcotest.(check int) "payload_length" 11 frame.payload_length;
let rev_payload_chunks = read_payload frame in
Alcotest.(check bool) "is_fin" true (Parse.is_fin frame);
Alcotest.(check bool) "is_fin" true frame.is_fin;
Alcotest.(check (list string)) "payload" ["1234567890\n"] rev_payload_chunks


let test_parsing_fin_bit () =
let frame = parse_frame (serialize_frame ~is_fin:false "hello") in
Alcotest.check Testable.opcode "opcode" `Text (Parse.opcode frame);
Alcotest.(check bool) "is_fin" false (Parse.is_fin frame);
Alcotest.check Testable.opcode "opcode" `Text frame.opcode;
Alcotest.(check bool) "is_fin" false frame.is_fin;
let frame = parse_frame (serialize_frame ~is_fin:true "hello") in
Alcotest.check Testable.opcode "opcode" `Text (Parse.opcode frame);
Alcotest.(check bool) "is_fin" true (Parse.is_fin frame);
Alcotest.check Testable.opcode "opcode" `Text frame.opcode;
Alcotest.(check bool) "is_fin" true frame.is_fin;
let rev_payload_chunks = read_payload frame in
Alcotest.(check (list string)) "payload" ["hello"] rev_payload_chunks

Expand Down

0 comments on commit 3492a4a

Please sign in to comment.