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.

Sign up for GitHub

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: mask all client frames #69

Merged
merged 2 commits into from
Aug 3, 2024
Merged

fix: mask all client frames #69

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5b56860
fix: mask all client frames
anmonteiro Aug 1, 2024
ddbbeea
fix test
anmonteiro Aug 1, 2024
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
39 changes: 25 additions & 14 deletions lib/serialize.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,25 @@
let serialize_headers ?mask faraday ~is_fin ~opcode ~payload_length =
type mode =
[ `Client of unit -> int32
| `Server
]

let mask mode =
match mode with
| `Client m -> Some (m ())
| `Server -> None


let serialize_headers ~mode faraday ~is_fin ~opcode ~payload_length =
let opcode = Websocket.Opcode.to_int opcode in
let is_fin = if is_fin then 1 lsl 7 else 0 in
let mask = mask mode in
let is_mask =
match mask with
| None -> 0
| Some _ -> 1 lsl 7
in
Faraday.write_uint8 faraday (is_fin lor opcode);
if payload_length <= 125 then
if payload_length <= 125 then
Faraday.write_uint8 faraday (is_mask lor payload_length)
else if payload_length <= 0xffff then begin
Faraday.write_uint8 faraday (is_mask lor 126);
Expand All @@ -16,28 +28,27 @@ let serialize_headers ?mask faraday ~is_fin ~opcode ~payload_length =
Faraday.write_uint8 faraday (is_mask lor 127);
Faraday.BE.write_uint64 faraday (Int64.of_int payload_length);
end;
begin match mask with
| None -> ()
| Some mask -> Faraday.BE.write_uint32 faraday mask
end
Option.iter (Faraday.BE.write_uint32 faraday) mask;
mask
;;

let serialize_control ?mask faraday ~opcode =
let serialize_control ~mode faraday ~opcode =
let opcode = (opcode :> Websocket.Opcode.t) in
serialize_headers faraday ?mask ~is_fin:true ~opcode ~payload_length:0
let _mask: int32 option =
serialize_headers faraday ~mode ~is_fin:true ~opcode ~payload_length:0
in
()

let schedule_serialize ?mask faraday ~is_fin ~opcode ~payload ~src_off ~off ~len =
serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length:len;
begin match mask with
let schedule_serialize ~mode faraday ~is_fin ~opcode ~payload ~src_off ~off ~len =
begin match serialize_headers faraday ~mode ~is_fin ~opcode ~payload_length:len with
| None -> ()
| Some mask -> Websocket.Frame.apply_mask mask payload ~src_off ~off ~len
end;
Faraday.schedule_bigstring faraday payload ~off ~len;
;;

let serialize_bytes ?mask faraday ~is_fin ~opcode ~payload ~src_off ~off ~len =
serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length:len;
begin match mask with
let serialize_bytes ~mode faraday ~is_fin ~opcode ~payload ~src_off ~off ~len =
begin match serialize_headers faraday ~mode ~is_fin ~opcode ~payload_length:len with
| None -> ()
| Some mask -> Websocket.Frame.apply_mask_bytes mask payload ~src_off ~off ~len
end;
Expand Down
29 changes: 8 additions & 21 deletions lib/wsd.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,7 @@ module IOVec = Httpun.IOVec

type error = [ `Exn of exn ]

type mode =
[ `Client of unit -> int32
| `Server
]
type mode = Serialize.mode

type t =
{ faraday : Faraday.t
Expand All @@ -26,11 +23,6 @@ let create ~error_handler mode =
; error_code = `Ok
}

let mask t =
match t.mode with
| `Client m -> Some (m ())
| `Server -> None

let is_closed t =
Faraday.is_closed t.faraday

Expand All @@ -47,22 +39,20 @@ let wakeup t =
Optional_thunk.call_if_some f

let schedule t ?(is_fin=true) ~kind payload ~off ~len =
let mask = mask t in
Serialize.schedule_serialize
t.faraday
(* TODO: is_fin *)
?mask
~mode:t.mode
~is_fin
~opcode:(kind :> Websocket.Opcode.t)
~src_off:0
~payload ~off ~len;
wakeup t

let send_bytes t ?(is_fin=true) ~kind payload ~off ~len =
let mask = mask t in
Serialize.serialize_bytes
t.faraday
?mask
~mode:t.mode
~is_fin
~opcode:(kind :> Websocket.Opcode.t)
~payload
Expand All @@ -73,12 +63,11 @@ let send_bytes t ?(is_fin=true) ~kind payload ~off ~len =

let send_ping ?application_data t =
begin match application_data with
| None -> Serialize.serialize_control t.faraday ~opcode:`Ping
| None -> Serialize.serialize_control ~mode:t.mode t.faraday ~opcode:`Ping
| Some { IOVec.buffer; off; len } ->
let mask = mask t in
Serialize.schedule_serialize
t.faraday
?mask
~mode:t.mode
~is_fin:true
~opcode:`Ping
~src_off:0
Expand All @@ -90,12 +79,11 @@ let send_ping ?application_data t =

let send_pong ?application_data t =
begin match application_data with
| None -> Serialize.serialize_control t.faraday ~opcode:`Pong;
| None -> Serialize.serialize_control ~mode:t.mode t.faraday ~opcode:`Pong;
| Some { IOVec.buffer; off; len } ->
let mask = mask t in
Serialize.schedule_serialize
t.faraday
?mask
~mode:t.mode
~is_fin:true
~opcode:`Pong
~src_off:0
Expand All @@ -110,11 +98,10 @@ let flushed t f = Faraday.flush t.faraday f
let close ?code t =
begin match code with
| Some code ->
let mask = mask t in
let payload = Bytes.create 2 in
Bytes.set_uint16_be payload 0 (Websocket.Close_code.to_int code);
Serialize.serialize_bytes t.faraday
?mask
~mode:t.mode
~is_fin:true
~opcode:`Connection_close
~src_off:0
Expand Down
1 change: 1 addition & 0 deletions lib_test/test_httpun_ws.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ module Websocket = struct
let f = Faraday.create 0x100 in
Serialize.serialize_bytes
f
~mode:`Server
~is_fin
~opcode:`Text
~payload:(Bytes.of_string frame)
Expand Down