From c4edb915c6a55ade414eadea4f9b18a26bc27c32 Mon Sep 17 00:00:00 2001 From: Antonio Nuno Monteiro Date: Wed, 31 Jul 2024 22:57:59 -0700 Subject: [PATCH 1/2] fix: garbled data after committing frame header --- lib/parse.ml | 131 +++++++++++++++++++++++++---------------------- lib/websocket.ml | 1 - 2 files changed, 69 insertions(+), 63 deletions(-) diff --git a/lib/parse.ml b/lib/parse.ml index 52c33936..da690e81 100644 --- a/lib/parse.ml +++ b/lib/parse.ml @@ -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 = @@ -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 = @@ -117,25 +120,36 @@ 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 frame = let open Angstrom in - parse_headers - >>| fun headers -> - let len = payload_length_of_headers headers in - let payload = match len with - | 0 -> Payload.empty - | _ -> Payload.create buf - in - { headers; payload } -;; + let frame ~buf = + parse_headers + >>| fun headers -> + 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 + { is_fin; opcode; mask; payload_length; payload } + in + fun ~buf handler -> + frame ~buf <* commit >>= fun frame -> + let { is_fin; opcode; payload_length = len; _ } = frame in + handler ~opcode ~is_fin ~len frame.payload; + payload_parser frame + ;; module Reader = struct module AU = Angstrom.Unbuffered @@ -154,14 +168,7 @@ module Reader = struct let parser = let open Angstrom in let buf = Bigstringaf.create 0x1000 in - 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 - frame_handler ~opcode ~is_fin ~len payload; - payload_parser frame) + skip_many (frame ~buf frame_handler) in { parser ; parse_state = Done diff --git a/lib/websocket.ml b/lib/websocket.ml index eebb5b08..1c56a086 100644 --- a/lib/websocket.ml +++ b/lib/websocket.ml @@ -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) From e560c5c99ed2768f7c3752f856577e82d93a69fd Mon Sep 17 00:00:00 2001 From: Antonio Nuno Monteiro Date: Wed, 31 Jul 2024 23:07:45 -0700 Subject: [PATCH 2/2] fix tests --- lib/parse.ml | 40 ++++++++++++++++++-------------------- lib_test/test_httpun_ws.ml | 33 +++++++++++++------------------ 2 files changed, 33 insertions(+), 40 deletions(-) diff --git a/lib/parse.ml b/lib/parse.ml index da690e81..3a5016a2 100644 --- a/lib/parse.ml +++ b/lib/parse.ml @@ -129,27 +129,20 @@ let payload_parser t = >>= fun () -> finish t.payload ;; -let frame = +let frame ~buf = let open Angstrom in - let frame ~buf = - parse_headers - >>| fun headers -> - 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 - { is_fin; opcode; mask; payload_length; payload } - in - fun ~buf handler -> - frame ~buf <* commit >>= fun frame -> - let { is_fin; opcode; payload_length = len; _ } = frame in - handler ~opcode ~is_fin ~len frame.payload; - payload_parser frame - ;; + parse_headers + >>| fun headers -> + 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 + { is_fin; opcode; mask; payload_length; payload } +;; module Reader = struct module AU = Angstrom.Unbuffered @@ -168,7 +161,12 @@ module Reader = struct let parser = let open Angstrom in let buf = Bigstringaf.create 0x1000 in - skip_many (frame ~buf frame_handler) + skip_many + (frame ~buf <* commit >>= fun frame -> + let payload = frame.payload in + let { is_fin; opcode; payload_length = len; _ } = frame in + frame_handler ~opcode ~is_fin ~len payload; + payload_parser frame) in { parser ; parse_state = Done diff --git a/lib_test/test_httpun_ws.ml b/lib_test/test_httpun_ws.ml index a5db8ef0..309c5009 100644 --- a/lib_test/test_httpun_ws.ml +++ b/lib_test/test_httpun_ws.ml @@ -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 @@ -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