From c4edb915c6a55ade414eadea4f9b18a26bc27c32 Mon Sep 17 00:00:00 2001 From: Antonio Nuno Monteiro Date: Wed, 31 Jul 2024 22:57:59 -0700 Subject: [PATCH] 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)