unikernel/duniverse/ocaml-h1/lib/websocket.ml

@@ -0,0 +1,656 @@
+(*----------------------------------------------------------------------------
+    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