mte/unikernel/duniverse/ocaml-h1/lib/websocket.ml
2025-11-11 02:07:51 +01:00

656 lines
21 KiB
OCaml

(*----------------------------------------------------------------------------
Copyright (c) 2018 Inhabited Type LLC.
Copyright (c) 2025 Robur Cooperative
Copyright (c) 2025 Swrup <swrup@protonmail.com>
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
3. Neither the name of the author nor the names of his contributors
may be used to endorse or promote products derived from this software
without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE CONTRIBUTORS ``AS IS'' AND ANY EXPRESS
OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
----------------------------------------------------------------------------
Modified by Swrup <swrup@protonmail.com> *)
open Httpun_types
module H1_client_connection = Client_connection.Oneshot
module Opcode = struct
type standard_non_control = [ `Continuation | `Text | `Binary ]
type standard_control = [ `Connection_close | `Ping | `Pong ]
type standard = [ standard_non_control | standard_control ]
type t = [ standard | `Other of int ]
let code = function
| `Continuation -> 0x0
| `Text -> 0x1
| `Binary -> 0x2
| `Connection_close -> 0x8
| `Ping -> 0x9
| `Pong -> 0xa
| `Other code -> code
let code_table : t array =
[|
`Continuation;
`Text;
`Binary;
`Other 0x3;
`Other 0x4;
`Other 0x5;
`Other 0x6;
`Other 0x7;
`Connection_close;
`Ping;
`Other 0xb;
`Other 0xc;
`Other 0xd;
`Other 0xe;
`Other 0xf;
|]
let unsafe_of_code code = Array.unsafe_get code_table code
let of_code code =
if code > 0xf then None else Some (Array.unsafe_get code_table code)
let of_code_exn code =
if code > 0xf then
failwith "Opcode.of_code_exn: value can't fit in four bits";
Array.unsafe_get code_table code
let to_int = code
let of_int = of_code
let of_int_exn = of_code_exn
let pp_hum fmt = function
| `Continuation -> Format.fprintf fmt "`Continuation"
| `Text -> Format.fprintf fmt "`Text"
| `Binary -> Format.fprintf fmt "`Binary"
| `Connection_close -> Format.fprintf fmt "`Connection_close"
| `Ping -> Format.fprintf fmt "`Ping"
| `Pong -> Format.fprintf fmt "`Pong"
| `Other code -> Format.fprintf fmt "`Other %#x" code
end
module Close_code = struct
type standard =
[ `Normal_closure
| `Going_away
| `Protocol_error
| `Unsupported_data
| `No_status_rcvd
| `Abnormal_closure
| `Invalid_frame_payload_data
| `Policy_violation
| `Message_too_big
| `Mandatory_ext
| `Internal_server_error
| `TLS_handshake ]
type t = [ standard | `Other of int ]
let code = function
| `Normal_closure -> 1000
| `Going_away -> 1001
| `Protocol_error -> 1002
| `Unsupported_data -> 1003
| `No_status_rcvd -> 1005
| `Abnormal_closure -> 1006
| `Invalid_frame_payload_data -> 1007
| `Policy_violation -> 1008
| `Message_too_big -> 1009
| `Mandatory_ext -> 1010
| `Internal_server_error -> 1011
| `TLS_handshake -> 1015
| `Other code -> code
let code_table : t array =
[|
`Normal_closure;
`Going_away;
`Protocol_error;
`Unsupported_data;
`Other 1004;
`No_status_rcvd;
`Abnormal_closure;
`Invalid_frame_payload_data;
`Policy_violation;
`Message_too_big;
`Mandatory_ext;
`Internal_server_error;
`Other 1012;
`Other 1013;
`Other 1014;
`TLS_handshake;
|]
let unsafe_of_code code = Array.unsafe_get code_table code
let of_code code =
if code > 0xffff || code < 1000 then None
else if code < 1016 then Some (unsafe_of_code (code land 0b1111))
else Some (`Other code)
let of_code_exn code =
if code > 0xffff then
failwith "Close_code.of_code_exn: value can't fit in two bytes";
if code < 1000 then
failwith "Close_code.of_code_exn: value in invalid range 0-999";
if code < 1016 then unsafe_of_code (code land 0b1111) else `Other code
let to_int = code
let of_int = of_code
let of_int_exn = of_code_exn
end
module Frame = struct
type t = Bstr.t
let is_fin t =
let bits = Bstr.unsafe_get t 0 |> Char.code in
bits land (1 lsl 7) = 1 lsl 7
let rsv t =
let bits = Bstr.unsafe_get t 0 |> Char.code in
(bits lsr 4) land 0b0111
let opcode t =
let bits = Bstr.unsafe_get t 0 |> Char.code in
bits land 0b1111 |> Opcode.unsafe_of_code
let payload_length_of_offset t off =
let bits = Bstr.unsafe_get t (off + 1) |> Char.code in
let length = bits land 0b01111111 in
if length = 126 then Bstr.get_int16_be t (off + 2)
else if
(* This is technically unsafe, but if somebody's asking us to read 2^63
* bytes, then we're already screwd. *)
length = 127
then Bstr.get_int64_be t (off + 2) |> Int64.to_int
else length
let payload_length t = payload_length_of_offset t 0
let has_mask t =
let bits = Bstr.unsafe_get t 1 |> Char.code in
bits land (1 lsl 7) = 1 lsl 7
let mask t =
if not (has_mask t) then None
else
Some
(let bits = Bstr.unsafe_get t 1 |> Char.code in
if bits = 254 then Bstr.get_int32_be t 4
else if bits = 255 then Bstr.get_int32_be t 10
else Bstr.get_int32_be t 2)
let mask_exn t =
let bits = Bstr.unsafe_get t 1 |> Char.code in
if bits = 254 then Bstr.get_int32_be t 4
else if bits = 255 then Bstr.get_int32_be t 10
else if bits >= 127 then Bstr.get_int32_be t 2
else failwith "Frame.mask_exn: no mask present"
let payload_offset_of_bits bits =
let initial_offset = 2 in
let mask_offset = (bits land (1 lsl 7)) lsr (7 - 2) in
let length_offset =
let length = bits land 0b01111111 in
if length < 126 then 0 else 2 lsl (length land 0b1) lsl 2
in
initial_offset + mask_offset + length_offset
let payload_offset t =
let bits = Bstr.unsafe_get t 1 |> Char.code in
payload_offset_of_bits bits
let payload_view t =
let len = payload_length t in
let off = payload_offset t in
Bstr.sub t ~off ~len
let length_of_offset t off =
let bits = Bstr.unsafe_get t (off + 1) |> Char.code in
let payload_offset = payload_offset_of_bits bits in
let payload_length = payload_length_of_offset t off in
payload_offset + payload_length
let length t = length_of_offset t 0
let apply_mask mask bs ~off ~len =
for i = off to off + len - 1 do
let j = (i - off) mod 4 in
let c = Bstr.unsafe_get bs i |> Char.code in
let c =
c lxor Int32.(logand (shift_right mask (8 * (3 - j))) 0xffl |> to_int)
in
Bstr.unsafe_set bs i (Char.unsafe_chr c)
done
let apply_mask_bytes mask bs ~off ~len =
for i = off to off + len - 1 do
let j = (i - off) mod 4 in
let c = Bytes.unsafe_get bs i |> Char.code in
let c =
c lxor Int32.(logand (shift_right mask (8 * (3 - j))) 0xffl |> to_int)
in
Bytes.unsafe_set bs i (Char.unsafe_chr c)
done
let unmask_inplace t =
if has_mask t then
let mask = mask_exn t in
let len = payload_length t in
let off = payload_offset t in
apply_mask mask t ~off ~len
let mask_inplace = unmask_inplace
let parse =
let open Angstrom in
Unsafe.peek 2 (fun bs ~off ~len:_ -> length_of_offset bs off) >>= fun len ->
Unsafe.take len Bstr.sub
let serialize_headers faraday ~mask ~is_fin ~opcode ~payload_length =
let opcode = Opcode.to_int opcode in
let is_fin = if is_fin then 1 lsl 7 else 0 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
Faraday.write_uint8 faraday (is_mask lor payload_length)
else if payload_length <= 0xffff then (
Faraday.write_uint8 faraday (is_mask lor 126);
Faraday.BE.write_uint16 faraday payload_length)
else (
Faraday.write_uint8 faraday (is_mask lor 127);
Faraday.BE.write_uint64 faraday (Int64.of_int payload_length));
match mask with
| None -> ()
| Some mask -> Faraday.BE.write_uint32 faraday mask
let serialize_control faraday ~mask ~opcode =
let opcode = (opcode :> Opcode.t) in
serialize_headers faraday ~mask ~is_fin:true ~opcode ~payload_length:0
let schedule_serialize faraday ~mask ~is_fin ~opcode ~payload ~off ~len =
serialize_headers faraday ~mask ~is_fin ~opcode ~payload_length:len;
(match mask with
| None -> ()
| Some mask -> apply_mask mask payload ~off ~len);
Faraday.schedule_bigstring faraday payload ~off ~len
let serialize_bytes faraday ~mask ~is_fin ~opcode ~payload ~off ~len =
serialize_headers faraday ~mask ~is_fin ~opcode ~payload_length:len;
(match mask with
| None -> ()
| Some mask -> apply_mask_bytes mask payload ~off ~len);
Faraday.write_bytes faraday payload ~off ~len
let schedule_serialize_bytes faraday ~mask ~is_fin ~opcode ~payload ~off ~len
=
serialize_headers faraday ~mask ~is_fin ~opcode ~payload_length:len;
(match mask with
| None -> ()
| Some mask -> apply_mask_bytes mask payload ~off ~len);
Faraday.write_bytes faraday payload ~off ~len
end
type frame_handler =
opcode:Opcode.t -> is_fin:bool -> Bstr.t -> off:int -> len:int -> unit
type input_handlers = { frame_handler : frame_handler; eof : unit -> unit }
module Wsd = struct
type mode = [ `Client of unit -> int32 | `Server ]
type t = {
faraday : Faraday.t;
mode : mode;
mutable when_ready_to_write : unit -> unit;
}
let default_ready_to_write = Sys.opaque_identity (fun () -> ())
let create mode =
{
faraday = Faraday.create 0x1000;
mode;
when_ready_to_write = default_ready_to_write;
}
let mask t = match t.mode with `Client m -> Some (m ()) | `Server -> None
let ready_to_write t =
let callback = t.when_ready_to_write in
t.when_ready_to_write <- default_ready_to_write;
callback ()
let schedule t ~kind ~is_fin payload ~off ~len =
let opcode :> Opcode.t = kind in
let mask = mask t in
Frame.schedule_serialize t.faraday ~mask ~is_fin ~opcode ~payload
~off ~len;
ready_to_write t
let send_bytes t ~kind ~is_fin payload ~off ~len =
let opcode :> Opcode.t = kind in
let mask = mask t in
Frame.schedule_serialize_bytes t.faraday ~mask ~is_fin ~opcode
~payload ~off ~len;
ready_to_write t
let send_ping t =
Frame.serialize_control t.faraday ~mask:None ~opcode:`Ping;
ready_to_write t
let send_pong t =
Frame.serialize_control t.faraday ~mask:None ~opcode:`Pong;
ready_to_write t
let flushed t f = Faraday.flush t.faraday f
let close t =
Frame.serialize_control t.faraday ~mask:None
~opcode:`Connection_close;
Faraday.close t.faraday;
ready_to_write t
let next t =
match Faraday.operation t.faraday with
| `Close -> `Close 0 (* XXX(andreas): should track unwritten bytes *)
| `Yield -> `Yield
| `Writev iovecs -> `Write iovecs
let report_result t result =
match result with
| `Closed -> close t
| `Ok len -> Faraday.shift t.faraday len
let is_closed t = Faraday.is_closed t.faraday
let when_ready_to_write t callback =
if not (t.when_ready_to_write == default_ready_to_write) then
failwith
"Wsd.when_ready_to_write: only one callback can be registered at a time"
else if is_closed t then callback ()
else t.when_ready_to_write <- callback
end
module Reader = struct
module AU = Angstrom.Unbuffered
type 'error parse_state =
| Done
| Fail of 'error
| Partial of (Bstr.t -> off:int -> len:int -> AU.more -> unit AU.state)
type 'error t = {
parser : unit Angstrom.t;
mutable parse_state : 'error parse_state;
mutable closed : bool;
}
let create frame_handler =
let parser =
let open Angstrom in
Frame.parse >>| fun frame ->
let is_fin = Frame.is_fin frame in
let opcode = Frame.opcode frame in
Frame.unmask_inplace frame;
let off = Frame.payload_offset frame in
let len = Frame.payload_length frame in
frame_handler ~opcode ~is_fin frame ~off ~len
in
{ parser; parse_state = Done; closed = false }
let transition t state =
match state with
| AU.Done (consumed, ()) | AU.Fail ((0 as consumed), _, _) ->
t.parse_state <- Done;
consumed
| AU.Fail (consumed, marks, msg) ->
t.parse_state <- Fail (`Parse (marks, msg));
consumed
| AU.Partial { committed; continue } ->
t.parse_state <- Partial continue;
committed
and start t state =
match state with
| AU.Done _ -> failwith "Websocket.Reader.unable to start parser"
| AU.Fail (0, marks, msg) -> t.parse_state <- Fail (`Parse (marks, msg))
| AU.Partial { committed = 0; continue } -> t.parse_state <- Partial continue
| _ -> assert false
let next t =
match t.parse_state with
| Done -> if t.closed then `Close else `Read
| Fail _ -> `Close
| Partial _ -> `Read
let rec read_with_more t bs ~off ~len more =
let consumed =
match t.parse_state with
| Fail _ -> 0
| Done ->
start t (AU.parse t.parser);
read_with_more t bs ~off ~len more
| Partial continue -> transition t (continue bs more ~off ~len)
in
(match more with Complete -> t.closed <- true | Incomplete -> ());
consumed
end
module Connection = struct
type t = {
wsd : Wsd.t;
reader : [ `Parse of string list * string ] Reader.t;
eof : unit -> unit;
}
let create ~mode ~websocket_handler =
let wsd = Wsd.create mode in
let { frame_handler; eof } = websocket_handler wsd in
{ wsd; reader = Reader.create frame_handler; eof }
let next_read_operation t = Reader.next t.reader
let next_write_operation t = Wsd.next t.wsd
let read t bs ~off ~len = Reader.read_with_more t.reader bs ~off ~len Incomplete
let read_eof t bs ~off ~len =
let len = Reader.read_with_more t.reader bs ~off ~len Complete in
t.eof ();
len
let is_closed t = Wsd.is_closed t.wsd
let close t = Wsd.close t.wsd
let yield_writer t k =
if is_closed t then (
close t;
k ())
else Wsd.when_ready_to_write t.wsd k
let report_write_result t result = Wsd.report_result t.wsd result
end
module Handshake = struct
let compute_accept ~sha1 nonce = sha1 (nonce ^ "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
let get_nonce request = Headers.get request.Request.headers "sec-websocket-key"
let server_headers ~sha1 ~nonce =
Headers.of_list
[ ("connection", "upgrade"); ("upgrade", "websocket")
; ("sec-websocket-accept", compute_accept ~sha1 nonce ) ]
let is_valid_accept_headers ~sha1 ~nonce headers =
let sec_websocket_accept = Headers.get headers "sec-websocket-accept" in
let upgrade = Headers.get headers "upgrade" |> Option.map String.lowercase_ascii in
let connection = Headers.get headers "connection" |> Option.map String.lowercase_ascii in
(sec_websocket_accept = Some (compute_accept ~sha1 nonce))
&& (upgrade = Some "websocket")
&& (connection = Some "upgrade")
end
module Client_handshake = struct
type t = { connection : H1_client_connection.t; body : Body.Writer.t }
(* assumes [nonce] is base64 encoded *)
let create ~nonce ~host ~port ~resource ~error_handler ~response_handler =
let headers = Headers.of_list
[ ("upgrade", "websocket");
("connection", "upgrade");
("host", String.concat ":" [ host; string_of_int port ]);
("sec-websocket-version", "13");
("sec-websocket-key", nonce); ]
in
let body, connection =
H1_client_connection.request
(Request.create ~headers `GET resource)
~error_handler ~response_handler
in
{ connection; body }
let next_read_operation t =
H1_client_connection.next_read_operation t.connection
let next_write_operation t =
H1_client_connection.next_write_operation t.connection
let read t = H1_client_connection.read t.connection
let report_write_result t =
H1_client_connection.report_write_result t.connection
let yield_writer t = H1_client_connection.yield_writer t.connection
let close t = Body.Writer.close t.body
end
module Client_connection = struct
type state =
| Uninitialized
| Handshake of Client_handshake.t
| Websocket of Connection.t
type t = state ref
type error =
[ H1_client_connection.error
| `Handshake_failure of Response.t * Body.Reader.t ]
let handshake_exn t =
match !t with
| Handshake handshake -> handshake
| Uninitialized | Websocket _ -> assert false
let create ~nonce ~host ~port ~resource ~sha1 ~error_handler ~websocket_handler
=
let t = ref Uninitialized in
let nonce = Base64.encode_exn nonce in
let response_handler response response_body =
match response.Response.status with
| `Switching_protocols when Handshake.is_valid_accept_headers ~sha1 ~nonce response.headers ->
Body.Reader.close response_body;
let handshake = handshake_exn t in
t :=
Websocket
(Connection.create
~mode:(`Client (fun () -> Random.int32 Int32.max_int))
~websocket_handler);
Client_handshake.close handshake
| _ -> error_handler (`Handshake_failure (response, response_body))
in
let handshake =
let error_handler = (error_handler :> H1_client_connection.error_handler) in
Client_handshake.create ~nonce ~host ~port ~resource ~error_handler
~response_handler
in
t := Handshake handshake;
t
let next_read_operation t =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.next_read_operation handshake
| Websocket websocket -> Connection.next_read_operation websocket
let read t bs ~off ~len =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.read handshake bs ~off ~len
| Websocket websocket -> Connection.read websocket bs ~off ~len
let read_eof t bs ~off ~len =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.read handshake bs ~off ~len
| Websocket websocket -> Connection.read_eof websocket bs ~off ~len
let next_write_operation t =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.next_write_operation handshake
| Websocket websocket -> Connection.next_write_operation websocket
let report_write_result t result =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.report_write_result handshake result
| Websocket websocket -> Connection.report_write_result websocket result
let yield_writer t f =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.yield_writer handshake f
| Websocket websocket -> Connection.yield_writer websocket f
let close t =
match !t with
| Uninitialized -> assert false
| Handshake handshake -> Client_handshake.close handshake
| Websocket websocket -> Connection.close websocket
end
module Server_connection = struct
type t = Connection.t
type error = [ `Exn of exn ]
let create ~websocket_handler =
let t = Connection.create ~mode:`Server ~websocket_handler in
t
let next_read_operation = Connection.next_read_operation
let next_write_operation = Connection.next_write_operation
let read t bs ~off ~len = Connection.read t bs ~off ~len
let read_eof t bs ~off ~len = Connection.read_eof t bs ~off ~len
let report_write_result t result = Connection.report_write_result t result
let yield_writer t f = Connection.yield_writer t f
let is_closed t = Connection.is_closed t
let close t = Connection.close t
end