Source file encoder.ml

(*----------------------------------------------------------------------------
 *  Copyright (c) 2019 António Nuno Monteiro
 *
 *  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 copyright holder nor the names of its
 *  contributors may be used to endorse or promote products derived from this
 *  software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND 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 COPYRIGHT HOLDER 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.
 *---------------------------------------------------------------------------*)

open Types

module HeaderFieldsTbl = Hashtbl.Make (struct
    type t = string

    let equal = String.equal
    let hash s = Hashtbl.hash s
  end)

module ValueMap = Map.Make (String)

type t =
  { table : Dynamic_table.t
        (* We maintain a lookup table of header fields to their indexes in the
         * dynamic table. The format is name -> (value -> index) *)
  ; lookup_table : int ValueMap.t HeaderFieldsTbl.t
  ; mutable next_seq : int
  }

module BinaryFormat = struct
  (* From RFC7541§6.2.3. Literal Header Field Never Indexed
   *   A literal header field never-indexed representation starts with the
   *   '0001' 4-bit pattern. *)
  let never_indexed = 0b0001_0000, 4

  (* From RFC7541§6.2.2: Literal Header Field without Indexing
   *   A literal header field without indexing representation starts with the
   *   '0000' 4-bit pattern. *)
  let without_indexing = 0b0000_0000, 4

  (* From RFC7541§6.2.1: Literal Header Field with Incremental Indexing
   *   A literal header field with incremental indexing representation starts
   *   with the '01' 2-bit pattern. *)
  let incremental_indexing = 0b0100_0000, 6

  (* From RFC7541§6.1: Indexed Header Field Representation
   *   An indexed header field starts with the '1' 1-bit pattern, followed by
   *   the index of the matching header field, represented as an integer with
   *   a 7-bit prefix (see Section 5.1). *)
  let indexed = 0b1000_0000, 7
  let[@inline] is_indexed = function 128 -> true | _ -> false
end

let create =
  let on_evict lookup_table (name, value) =
    let map = HeaderFieldsTbl.find lookup_table name in
    if ValueMap.cardinal map = 1
    then HeaderFieldsTbl.remove lookup_table name
    else
      let map = ValueMap.remove value map in
      HeaderFieldsTbl.replace lookup_table name map
  in
  fun capacity ->
    let lookup_table = HeaderFieldsTbl.create capacity in
    { table = Dynamic_table.create ~on_evict:(on_evict lookup_table) capacity
    ; lookup_table
    ; next_seq = 0
    }

let add ({ table; lookup_table; next_seq } as encoder) entry =
  let name, value = entry in
  Dynamic_table.add table entry;
  let map =
    match HeaderFieldsTbl.find_opt lookup_table name with
    | Some map -> ValueMap.add value next_seq map
    | None -> ValueMap.singleton value next_seq
  in
  encoder.next_seq <- next_seq + 1;
  HeaderFieldsTbl.replace lookup_table name map

let encode =
  let find_token =
    let encode_missing_value encoder without_indexing token name value =
      (* This is a header field whose value we didn't find in the static
       * table after looping. We ended up here (name <> name') because we
       * looped to check whether the value was indexed in the static table.
       * We can still use the token index to encode the header name. *)
      let index = token + 1 in
      if without_indexing
      then
        (* From RFC7541§6.2.2: Literal Header Field without Indexing
         *   If the header field name matches the header field name of an entry
         *   stored in the static table or the dynamic table, the header field
         *   name can be represented using the index of that entry. *)
        BinaryFormat.without_indexing, index
      else (
        (* From RFC7541§6.2.1: Literal Header Field with Incremental Indexing
         *   A literal header field with incremental indexing representation
         *   results in appending a header field to the decoded header list and
         *   inserting it as a new entry into the dynamic table. *)
        add encoder (name, value);
        BinaryFormat.incremental_indexing, index)
    in
    fun encoder without_indexing token name value ->
      let rec loop i =
        if i >= Static_table.table_size
        then encode_missing_value encoder without_indexing token name value
        else
          let name', value' = Static_table.table.(i) in
          if name = name'
          then
            if value' = value
            then
              (* From RFC7541§6.1: Indexed Header Field Representation
               *   An indexed header field starts with the '1' 1-bit pattern,
               *   followed by the index of the matching header field. *)
              BinaryFormat.indexed, i + 1
            else
              (* Advance one token in the static table, as the next entry might have
               * a value that can fall into the above branch. We're guaranteed to
               * always get the first token (index) in the static table for `name`,
               * because that's what `Static_table.lookup_token` returns. *)
              loop (i + 1)
          else encode_missing_value encoder without_indexing token name value
      in
      loop token
  in
  let[@inline] seq_to_index next_seq seq =
    Static_table.table_size + next_seq - seq
  in
  let is_without_indexing =
    let module IntSet = Set.Make (Int) in
    let tokens_without_indexing =
      (* From RFC7541§6.2.2: Never-Indexed Literals
       *   Either form of header field name representation is followed by the
       *   header field value represented as a string literal (see Section 5.2).
       *
       * Note: we choose not to index the values of these fields as they would
       * vary immensely. This way, we save some additions / evictions from the
       * dynamic table. *)
      IntSet.of_list
        Static_table.TokenIndices.
          [ path
          ; age
          ; content_length
          ; etag
          ; if_modified_since
          ; if_none_match
          ; location
          ; set_cookie
          ]
    in
    fun [@inline] token ->
      token <> -1 && IntSet.mem token tokens_without_indexing
  in
  let[@inline] is_sensitive token value =
    token <> -1
    && (* From RFC7541§7.1.3: Never-Indexed Literals
        *   An encoder might also choose not to index values for header fields
        *   that are considered to be highly valuable or sensitive to recovery,
        *   such as the Cookie or Authorization header fields. *)
    Static_table.TokenIndices.(
      token == authorization || (token == cookie && String.length value < 20))
  in
  fun ({ lookup_table; next_seq; _ } as encoder) { name; value; sensitive } ->
    let token = Static_table.lookup_token_index name in
    let token_found_in_static_table = token <> -1 in
    if sensitive || is_sensitive token value
    then
      (* never indexed literal header field, find the index *)
      let index =
        if token_found_in_static_table
        then
          (* From RFC7541§6.2.2: Literal Header Field without Indexing
           *   If the header field name matches the header field name of an entry
           *   stored in the static table or the dynamic table, the header field
           *   name can be represented using the index of that entry. *)
          token + 1
        else
          match HeaderFieldsTbl.find_opt lookup_table name with
          | Some map ->
            let _, any_entry = ValueMap.choose map in
            seq_to_index next_seq any_entry
          | None ->
            (* From RFC7541§6.2.2: Literal Header Field without Indexing
             *   Otherwise, the header field name is represented as a string
             *   literal (see Section 5.2). A value 0 is used in place of the
             *   4-bit index, followed by the header field name. *)
            0
      in
      BinaryFormat.never_indexed, index
    else if token_found_in_static_table
    then
      (* Header name is represented in the static table. *)
      match HeaderFieldsTbl.find_opt lookup_table name with
      | Some map ->
        (* Header value is indexed in the dynamic table. *)
        (match ValueMap.find_opt value map with
        | Some seq ->
          (* From RFC7541§6.1: Indexed Header Field Representation
           *   An indexed header field representation identifies an entry in
           *   either the static table or the dynamic table (see Section 2.3). *)
          BinaryFormat.indexed, seq_to_index next_seq seq
        | None ->
          (* Header value is not indexed in the dynamic table. Check if it's an
           * entry in the static table or if we need to encode its value, (and
           * potentially name if the field is requested to be encoded without
           * indexing). *)
          let without_indexing = is_without_indexing token in
          find_token encoder without_indexing token name value)
      | None ->
        let without_indexing = is_without_indexing token in
        find_token encoder without_indexing token name value
    else
      match HeaderFieldsTbl.find_opt lookup_table name with
      | Some map ->
        (match ValueMap.find_opt value map with
        | Some seq -> BinaryFormat.indexed, seq_to_index next_seq seq
        | None ->
          let index = seq_to_index next_seq (snd (ValueMap.choose map)) in
          if is_without_indexing token
          then BinaryFormat.without_indexing, index
          else (
            (* From RFC7541§6.2.1
             *   A literal header field with incremental indexing representation
             *   results in appending a header field to the decoded header list
             *   and inserting it as a new entry into the dynamic table. *)
            add encoder (name, value);
            BinaryFormat.incremental_indexing, index))
      | None ->
        if is_without_indexing token
        then BinaryFormat.without_indexing, 0
        else (
          (* From RFC7541§6.2.1
           *   A literal header field with incremental indexing representation
           *   results in appending a header field to the decoded header list and
           *   inserting it as a new entry into the dynamic table. *)
          add encoder (name, value);
          BinaryFormat.incremental_indexing, 0)

let encode_int t prefix n i =
  let max_prefix = (1 lsl n) - 1 in
  if i < max_prefix
  then
    (* From RFC7541§5.1:
     *   If the integer value is small enough, i.e., strictly less than 2^N-1,
     *   it is encoded within the N-bit prefix. *)
    Faraday.write_uint8 t (prefix lor i)
  else
    (* From RFC7541§5.1:
     *   Otherwise, all the bits of the prefix are set to 1, and the value,
     *   decreased by 2^N-1, is encoded using a list of one or more octets. The
     *   most significant bit of each octet is used as a continuation flag: its
     *   value is set to 1 except for the last octet in the list. The remaining
     *   bits of the octets are used to encode the decreased value. *)
    let i = i - max_prefix in
    Faraday.write_uint8 t (prefix lor max_prefix);
    let rec loop i =
      if i >= 128
      then (
        Faraday.write_uint8 t (i land 127 lor 128);
        loop (i lsr 7))
      else Faraday.write_uint8 t i
    in
    loop i

let encode_header =
  let encode_string t s =
    let string_length = String.length s in
    let huffman_length = Huffman.encoded_length s in
    if huffman_length > string_length
    then (
      (* From RFC7541§5.2:
       *   The number of octets used to encode the string literal, encoded as an
       *   integer with a 7-bit prefix (see Section 5.1). *)
      encode_int t 0 7 string_length;
      (* From RFC7541§5.2:
       *   The encoded data of the string literal. If H is '0', then the encoded
       *   data is the raw octets of the string literal. If H is '1', then the
       *   encoded data is the Huffman encoding of the string literal. *)
      Faraday.write_string t s)
    else (
      (* From RFC7541§5.2:
       *   The number of octets used to encode the string literal, encoded as an
       *   integer with a 7-bit prefix (see Section 5.1). *)
      encode_int t 128 7 huffman_length;
      (* From RFC7541§5.2:
       *   The encoded data of the string literal. If H is '0', then the encoded
       *   data is the raw octets of the string literal. If H is '1', then the
       *   encoded data is the Huffman encoding of the string literal. *)
      Huffman.encode t s)
  in
  fun encoder t ({ name; value; _ } as header) ->
    let (prefix, prefix_length), index = encode encoder header in
    encode_int t prefix prefix_length index;
    match BinaryFormat.is_indexed prefix with
    | true -> ()
    | false ->
      if index == 0
      then
        (* From RFC7541§6.2.2: Literal Header Field without Indexing * If the
           header field name matches the header field name of an entry * stored
           in the static table or the dynamic table, the header field * name can
           be represented using the index of that entry. In this case, * the
           index of the entry is represented as an integer with a 4-bit * prefix
           (see Section 5.1). This value is always non-zero. * * Otherwise, the
           header field name is represented as a string literal * (see Section
           5.2). A value 0 is used in place of the 4-bit index, * followed by
           the header field name. *)
        encode_string t name;
      (* From RFC7541§6.2.2: Literal Header Field without Indexing
       *   Either form of header field name representation is followed by the
       *   header field value represented as a string literal (see
       *   Section 5.2). *)
      encode_string t value

let set_capacity { table; _ } new_capacity =
  Dynamic_table.set_capacity table new_capacity