Source file handshake_state.ml

type role = Initiator | Responder

type t = {
  suite : Suite.t;
  role : role;
  turn : role; (* whose turn it is to write *)
  sym : Symmetric_state.t;
  s : Dh.keypair option;
  e : Dh.keypair option;
  rs : string option;
  re : string option;
  psks : string list;
  patterns : Pattern.message_pattern list;
  finished : bool;
  h_final : string option;
  psk_mode : bool;
  one_way : bool;
}

let ( let* ) = Result.bind
let peer = function Initiator -> Responder | Responder -> Initiator
let hash st = st.suite.Suite.hash
let cipher st = st.suite.Suite.cipher
let dh st = st.suite.Suite.dh

let create ~protocol_name ~role ?prologue ?s ?e ?rs ?re ?(psks = []) () =
  let* suite = Suite.of_protocol_name protocol_name in
  let* pattern =
    Pattern.of_name (List.nth (String.split_on_char '_' protocol_name) 1)
  in
  let h = suite.Suite.hash in
  let sym0 = Symmetric_state.initialize_symmetric ~hash:h protocol_name in
  let sym1 =
    Symmetric_state.mix_hash sym0 ~hash:h (Option.value prologue ~default:"")
  in
  (* Pre-message keys are mixed initiator-first (spec §5.3). *)
  let mix_pre_keys sym role_pre is_local =
    List.fold_left
      (fun sym tok ->
        match tok with
        | Pattern.S -> (
            let key =
              if is_local then
                match s with Some kp -> Some kp.Dh.pub | None -> None
              else rs
            in
            match key with
            | Some pub -> Symmetric_state.mix_hash sym ~hash:h pub
            | None -> sym)
        | Pattern.E -> (
            let key =
              if is_local then
                match e with Some kp -> Some kp.Dh.pub | None -> None
              else re
            in
            match key with
            | Some pub -> Symmetric_state.mix_hash sym ~hash:h pub
            | None -> sym)
        | _ -> sym)
      sym role_pre
  in
  let sym2 =
    let sym_a =
      mix_pre_keys sym1 pattern.Pattern.pre.initiator (role = Initiator)
    in
    mix_pre_keys sym_a pattern.Pattern.pre.responder (role = Responder)
  in
  let psk_mode = List.exists (List.mem Pattern.PSK) pattern.Pattern.messages in
  let one_way = List.length pattern.Pattern.messages = 1 in
  let psk_tokens =
    List.fold_left
      (fun acc msg ->
        acc + List.length (List.filter (fun tok -> tok = Pattern.PSK) msg))
      0 pattern.Pattern.messages
  in
  let* () =
    if List.length psks <> psk_tokens then
      Error
        (Error.Bad_psk
           (Printf.sprintf "pattern requires %d, got %d" psk_tokens
              (List.length psks)))
    else if List.exists (fun psk -> String.length psk <> 32) psks then
      Error (Error.Bad_psk "pre-shared keys must be 32 bytes")
    else Ok ()
  in
  Ok
    {
      suite;
      role;
      turn = Initiator;
      sym = sym2;
      s;
      e;
      rs;
      re;
      psks;
      patterns = pattern.Pattern.messages;
      finished = false;
      h_final = None;
      psk_mode;
      one_way;
    }

let do_dh st local_kp remote_pub =
  (dh st).Dh.dh ~local:local_kp ~remote:remote_pub

let get_e st =
  match st.e with
  | Some kp -> Ok kp
  | None -> Error (Error.Invalid_state "missing local ephemeral key")

let get_s st =
  match st.s with
  | Some kp -> Ok kp
  | None -> Error (Error.Invalid_state "missing local static key")

let get_re st =
  match st.re with
  | Some k -> Ok k
  | None -> Error (Error.Invalid_state "missing remote ephemeral key")

let get_rs st =
  match st.rs with
  | Some k -> Ok k
  | None -> Error (Error.Invalid_state "missing remote static key")

let write_token st buf tok =
  let h = hash st in
  let c = cipher st in
  match tok with
  | Pattern.E ->
      let* e_kp =
        match st.e with Some kp -> Ok kp | None -> (dh st).Dh.generate ()
      in
      let pub = e_kp.Dh.pub in
      Buffer.add_string buf pub;
      let sym = Symmetric_state.mix_hash st.sym ~hash:h pub in
      let sym =
        if st.psk_mode then Symmetric_state.mix_key sym ~hash:h pub else sym
      in
      Ok ({ st with sym; e = Some e_kp }, buf)
  | Pattern.S ->
      let* s_kp = get_s st in
      let* ct, sym =
        Symmetric_state.encrypt_and_hash st.sym ~hash:h ~cipher:c s_kp.Dh.pub
      in
      Buffer.add_string buf ct;
      Ok ({ st with sym }, buf)
  | Pattern.EE ->
      let* e_kp = get_e st in
      let* re_pk = get_re st in
      let* shared = do_dh st e_kp re_pk in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok ({ st with sym }, buf)
  | Pattern.ES ->
      (* Initiator: DH(e, rs);  Responder: DH(s, re) *)
      let* shared =
        match st.role with
        | Initiator ->
            let* e_kp = get_e st in
            let* rs_pk = get_rs st in
            do_dh st e_kp rs_pk
        | Responder ->
            let* s_kp = get_s st in
            let* re_pk = get_re st in
            do_dh st s_kp re_pk
      in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok ({ st with sym }, buf)
  | Pattern.SE ->
      (* Initiator: DH(s, re);  Responder: DH(e, rs) *)
      let* shared =
        match st.role with
        | Initiator ->
            let* s_kp = get_s st in
            let* re_pk = get_re st in
            do_dh st s_kp re_pk
        | Responder ->
            let* e_kp = get_e st in
            let* rs_pk = get_rs st in
            do_dh st e_kp rs_pk
      in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok ({ st with sym }, buf)
  | Pattern.SS ->
      let* s_kp = get_s st in
      let* rs_pk = get_rs st in
      let* shared = do_dh st s_kp rs_pk in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok ({ st with sym }, buf)
  | Pattern.PSK -> (
      match st.psks with
      | [] -> Error (Error.Bad_psk "no pre-shared key left for PSK token")
      | psk :: psks' ->
          let sym = Symmetric_state.mix_key_and_hash st.sym ~hash:h psk in
          Ok ({ st with sym; psks = psks' }, buf))

let read_token st msg pos tok =
  let h = hash st in
  let c = cipher st in
  match tok with
  | Pattern.E ->
      let dhlen = (dh st).Dh.dhlen in
      if String.length msg - !pos < dhlen then
        Error (Error.Invalid_state "message too short for E token")
      else begin
        let pub = String.sub msg !pos dhlen in
        pos := !pos + dhlen;
        let sym = Symmetric_state.mix_hash st.sym ~hash:h pub in
        let sym =
          if st.psk_mode then Symmetric_state.mix_key sym ~hash:h pub else sym
        in
        Ok { st with sym; re = Some pub }
      end
  | Pattern.S ->
      let dhlen = (dh st).Dh.dhlen in
      let tag_len =
        if Cipher_state.has_key st.sym.Symmetric_state.cs then 16 else 0
      in
      let needed = dhlen + tag_len in
      if String.length msg - !pos < needed then
        Error (Error.Invalid_state "message too short for S token")
      else begin
        let ct = String.sub msg !pos needed in
        pos := !pos + needed;
        let* pub, sym =
          Symmetric_state.decrypt_and_hash st.sym ~hash:h ~cipher:c ct
        in
        Ok { st with sym; rs = Some pub }
      end
  | Pattern.EE ->
      let* re_pk = get_re st in
      let* e_kp = get_e st in
      let* shared = do_dh st e_kp re_pk in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok { st with sym }
  | Pattern.ES ->
      let* shared =
        match st.role with
        | Initiator ->
            let* e_kp = get_e st in
            let* rs_pk = get_rs st in
            do_dh st e_kp rs_pk
        | Responder ->
            let* s_kp = get_s st in
            let* re_pk = get_re st in
            do_dh st s_kp re_pk
      in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok { st with sym }
  | Pattern.SE ->
      let* shared =
        match st.role with
        | Initiator ->
            let* s_kp = get_s st in
            let* re_pk = get_re st in
            do_dh st s_kp re_pk
        | Responder ->
            let* e_kp = get_e st in
            let* rs_pk = get_rs st in
            do_dh st e_kp rs_pk
      in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok { st with sym }
  | Pattern.SS ->
      let* s_kp = get_s st in
      let* rs_pk = get_rs st in
      let* shared = do_dh st s_kp rs_pk in
      let sym = Symmetric_state.mix_key st.sym ~hash:h shared in
      Ok { st with sym }
  | Pattern.PSK -> (
      match st.psks with
      | [] -> Error (Error.Bad_psk "no pre-shared key left for PSK token")
      | psk :: psks' ->
          let sym = Symmetric_state.mix_key_and_hash st.sym ~hash:h psk in
          Ok { st with sym; psks = psks' })

let check_message_length bytes =
  let n = String.length bytes in
  if n > 65535 then Error (Error.Bad_message_length { max = 65535; got = n })
  else Ok ()

let write_message st payload =
  if st.finished then Error (Error.Invalid_state "handshake already finished")
  else if st.role <> st.turn then
    Error (Error.Invalid_state "out of turn: expected read_message")
  else
    match st.patterns with
    | [] -> Error (Error.Invalid_state "no more message patterns")
    | msg_pattern :: rest ->
        let buf = Buffer.create 256 in
        let* st', _ =
          List.fold_left
            (fun acc tok ->
              let* st_acc, buf_acc = acc in
              write_token st_acc buf_acc tok)
            (Ok (st, buf))
            msg_pattern
        in
        let c = cipher st' in
        let h = hash st' in
        let* ct_payload, sym' =
          Symmetric_state.encrypt_and_hash st'.sym ~hash:h ~cipher:c payload
        in
        Buffer.add_string buf ct_payload;
        let message = Buffer.contents buf in
        let* () = check_message_length message in
        (* The handshake hash is h after the final payload encryption: "h at the
       time Split() is called" (spec §11.2). *)
        let finished, h_final, patterns' =
          match rest with
          | [] -> (true, Some sym'.Symmetric_state.h, [])
          | _ -> (false, None, rest)
        in
        let st'' =
          {
            st' with
            sym = sym';
            patterns = patterns';
            finished;
            h_final;
            turn = peer st.turn;
          }
        in
        Ok (message, st'')

let read_message st message =
  if st.finished then Error (Error.Invalid_state "handshake already finished")
  else if st.role = st.turn then
    Error (Error.Invalid_state "out of turn: expected write_message")
  else begin
    let* () = check_message_length message in
    match st.patterns with
    | [] -> Error (Error.Invalid_state "no more message patterns")
    | msg_pattern :: rest ->
        let pos = ref 0 in
        let* st' =
          List.fold_left
            (fun acc tok ->
              let* st_acc = acc in
              read_token st_acc message pos tok)
            (Ok st) msg_pattern
        in
        let ct_payload =
          String.sub message !pos (String.length message - !pos)
        in
        let c = cipher st' in
        let h = hash st' in
        let* payload, sym' =
          Symmetric_state.decrypt_and_hash st'.sym ~hash:h ~cipher:c ct_payload
        in
        let finished, h_final, patterns' =
          match rest with
          | [] -> (true, Some sym'.Symmetric_state.h, [])
          | _ -> (false, None, rest)
        in
        let st'' =
          {
            st' with
            sym = sym';
            patterns = patterns';
            finished;
            h_final;
            turn = peer st.turn;
          }
        in
        Ok (payload, st'')
  end

let is_finished st = st.finished
let handshake_hash st = st.h_final

let split st =
  if not st.finished then
    Error (Error.Invalid_state "handshake not yet finished")
  else
    let cs1, cs2 = Symmetric_state.split st.sym ~hash:(hash st) in
    Ok (cs1, cs2)

let role_of st = st.role
let cipher_of st = st.suite.Suite.cipher
let one_way_of st = st.one_way