Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: garbled data after committing frame header #68

Merged
merged 2 commits into from
Aug 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 55 additions & 50 deletions lib/parse.ml
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
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
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