feat ws: a bit of cleanup, expose masking primitive

This commit is contained in:
Simon Cruanes 2024-04-04 15:54:02 -04:00
parent 0014334010
commit dbd00259da
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
3 changed files with 22 additions and 8 deletions

View file

@ -192,7 +192,8 @@ module Reader = struct
type state =
| Begin (** At the beginning of a frame *)
| Reading_frame of { mutable remaining_bytes: int }
(** Currently reading the payload of a frame with [remaining_bytes] left to read *)
(** Currently reading the payload of a frame with [remaining_bytes]
left to read from the underlying [ic] *)
| Close
type t = {
@ -245,7 +246,7 @@ module Reader = struct
) else if len = 127 then (
IO.Input.really_input self.ic self.header_buf 0 8;
let len64 = Bytes.get_int64_be self.header_buf 0 in
if compare len64 (Int64.of_int max_fragment_size) > 0 then (
if Int64.compare len64 (Int64.of_int max_fragment_size) > 0 then (
Log.error (fun k ->
k "websocket: maximum frame fragment exceeded (%Ld > %d)" len64
max_fragment_size);
@ -267,14 +268,15 @@ module Reader = struct
self.header.payload_len self.header.mask);*)
()
external apply_masking_ : bytes -> bytes -> int -> int -> unit
external apply_masking_ : key:bytes -> buf:bytes -> int -> int -> unit
= "tiny_httpd_ws_apply_masking"
[@@noalloc]
(** Apply masking to the parsed data *)
let[@inline] apply_masking ~mask_key (buf : bytes) off len : unit =
assert (off >= 0 && off + len <= Bytes.length buf);
apply_masking_ mask_key buf off len
assert (
Bytes.length mask_key = 4 && off >= 0 && off + len <= Bytes.length buf);
apply_masking_ ~key:mask_key ~buf off len
let read_body_to_string (self : t) : string =
let len = self.header.payload_len in
@ -465,3 +467,7 @@ let add_route_handler ?accept ?(accept_ws_protocol = fun _ -> true)
end) in
let up : Server.upgrade_handler = (module M) in
Server.add_upgrade_handler ?accept server route up
module Private_ = struct
let apply_masking = Reader.apply_masking
end

View file

@ -20,3 +20,11 @@ val add_route_handler :
(** Add a route handler for a websocket endpoint.
@param accept_ws_protocol decides whether this endpoint accepts the websocket protocol
sent by the client. Default accepts everything. *)
(**/**)
module Private_ : sig
val apply_masking : mask_key:bytes -> bytes -> int -> int -> unit
end
(**/**)

View file

@ -13,9 +13,9 @@ CAMLprim value tiny_httpd_ws_apply_masking(value _mask_key, value _buf,
intnat len = Int_val(_len);
for (intnat i = 0; i < len; ++i) {
char c = buf[offset + i];
char c_m = mask_key[i & 0x3];
buf[offset + i] = c ^ c_m;
unsigned char c = buf[offset + i];
unsigned char c_m = mask_key[i & 0x3];
buf[offset + i] = (unsigned char)(c ^ c_m);
}
CAMLreturn(Val_unit);
}