(* Model of WireGuard written by Benjamin Lipp and Bruno Blanchet, Inria Paris.

   This file generates several CryptoVerif scripts using the m4 preprocessor.
   It uses the following settings:
   m4_replay_prot    When this macro is defined, the first message is protected
   		     against replays using the timestamp.
   m4_zero_test      When this macro is defined, the process fails in case
   		     a Diffie-Hellman shared secret is zero.
*)


(* To speed up proofs of correspondences *)
set casesInCorresp = false.
(* To speed up the proof *)
set mergeBranches = false.
(* To save memory *)
set forgetOldGames = true.


(* The macro m4_out_game outputs the current game to a file.
   The file name contains a number that is increased at each output. *)
define(`m4_gamenumber', 0)
define(`m4_out_game',`define(`m4_gamenumber', incr(m4_gamenumber))out_game "m4_name.m4_gamenumber.out.cv" occ')



(* Proof indications *)

proof {
  insert after "in(ch1_rom3" "let rom3_input(x1_rom3, G8_to_G(x2_rom3), G8_to_G(x3_rom3), x4_rom3, G8_to_G(x5_rom3), G8_to_G(x6_rom3), v_psk) = x_rom3 in";
  crypto rom(rom3_intermediate);
  crypto splitter(concat_four_keys) *;
  crypto int_ctxt_corrupt(enc) *;
  success
}


param N_sessions.

(**** Gap Diffie-Hellman ****)
type G_t [bounded,large].  (* type of public keys *)
type G8_t [bounded,large]. (* type of { X^k, X \in F_p } *)
fun G_to_bitstring(G_t): bitstring [data].
type Z_t [bounded,large,nonuniform].   (* type of exponents (must be "bounded" and
                               "large", is the set of multiples of k prime to qq' modulo kqq'.) *)
const dummy_z: Z_t.      (* placeholder when we do not compromise the ephemerals *)
proba P_GDH.            (* probability of breaking the GDH assumption *)
(* Page 7 in the Noise paper, rev 33:
   The public_key either encodes some value in a large prime-order group
   (which may have multiple equivalent encodings), or is an invalid
   value. *)

expand DH_X25519(
  (* types *)
  G_t,  (* Public keys *)
  Z_t,  (* Exponents *)
  (* variables *)
  g,    (* base point *)
  exp,  (* exponentiation function *)
  mult, (* multiplication function for exponents *)
  G8_t,
  g8,
  exp_div8,
  exp_div8', (* a symbol that replaces exp_div8 after game transformation *)
  pow8,
  G8_to_G,
  zero,
  zero8
).

expand GDH(
  (* types *)
  G8_t,  (* Group elements *)
  Z_t,  (* Exponents *)
  (* variables *)
  g8,    (* a generator of the group *)
  exp_div8,  (* exponentiation function *)
  exp_div8', (* a symbol that replaces exp_div8 after game transformation *)
  mult, (* multiplication function for exponents *)
  (* probabilities *)
  P_GDH (* probability of breaking the GDH assumption *)
).
letfun DH(group_element: G_t, exponent: Z_t) =
  exp(group_element, exponent).

(**** Symmetric encryption ****)
type key_t [large,fixed].
fun key_to_bitstring(key_t): bitstring [data].
type psk_t [large,fixed].         (* 32 byte pre-shared symmetric key *)
const psk_0: psk_t.                (* pre-shared key with all zeros,  *)
                                   (* used in case the WireGuard user *)
                                   (* did not provide a psk.          *)
type nonce_t [large,fixed].        (* 12 byte counter nonce for AEAD. *)
const nonce_0: nonce_t.            (* const value for the zero nonce  *)
const empty_bitstring : bitstring. (* const value for the empty
                                     bitstring that will be encrypted *)
proba P_enc.
proba P_encctxt.
expand AEAD_nonce(
  (* types *)
  key_t,     (* keys *)
  bitstring, (* plaintext *)
  bitstring, (* ciphertext *)
  bitstring, (* additional data *)
  nonce_t,   (* nonces *)
  (* functions *)
  enc,     (* encryption:
           (* enc(plaintext, additional data, key, nonce): ciphertext *)
  dec,  (* decryption:
        (* dec(ciphertext, additional data, key, nonce): bitstringbot *)
  injbot,    (* injection from plaintext to bitstringbot:
             (* injbot(plaintext): bitstringbot *)
  Zero,      (* returns a plaintext of same length, consisting of zeros:
             (* Zero(plaintext): plaintext *)
  (* probabilities *)
  P_enc,     (* breaking IND-CPA *)
  P_encctxt  (* breaking INT-CTXT *)
).

(**** Hash and HKDF ****)
type hashkey_t [fixed].
type hashoutput_t [large,fixed].
fun hashoutput_to_bitstring(hashoutput_t): bitstring [data].

(* This models the derivation of a first intermediate symmetric key. *)
type rom1_input_t.
fun rom1_input(G_t, G_t): rom1_input_t [data].

expand ROM_hash_refactored(
  (* types *)
  hashkey_t,   (* key of the hash function, models the choice of the
               (* hash function *)
  rom1_input_t, (* input type *)
  key_t,       (* output type *)
  (* functions *)
  rom1_intermediate,  (* name of the random oracle hash function:
               (* rom1(hashkey_t, rom2_input_t): key_t *)
  (* processes *)
  rom1_oracle, (* name of the oracle that will be available to the
               (* attacker *)
  (* variables *)
  r_rom1,
  x_rom1,
  (* channels for random oracle *)
  ch1_rom1,
  ch2_rom1,
  (* parameters *)
  N_qH1        (* number of queries to the oracle by the attacker *)
).

letfun rom1(hk: hashkey_t, x1_rom1: G_t, x2_rom1: G_t) =
   rom1_intermediate(hk, rom1_input(x1_rom1, x2_rom1)).

(* This models the derivation of a second intermediate symmetric key. *)
type rom2_input_t.
fun rom2_input(G_t, G_t, G_t): rom2_input_t [data].

expand ROM_hash_refactored(
  (* types *)
  hashkey_t,   (* key of the hash function, models the choice of the
               (* hash function *)
  rom2_input_t, (* input type *)
  key_t,       (* output type *)
  (* functions *)
  rom2_intermediate,        (* Name of the random oracle hash function:
               (* rom2(hashkey_t, rom2_input_t): key_t *)
  (* processes *)
  rom2_oracle, (* Name of the oracle that will be available to the
               (* attacker. *)
  (* variables *)
  r_rom2,
  x_rom2,
  (* channels for random oracle *)
  ch1_rom2,
  ch2_rom2,
  (* parameters *)
  N_qH2        (* Number of queries to the oracle by the attacker. *)
).

letfun rom2(hk: hashkey_t, x1_rom2: G_t, x2_rom2: G_t, x3_rom2: G_t) =
   rom2_intermediate(hk, rom2_input(x1_rom2, x2_rom2, x3_rom2)).

(* This models the derivation of
   * a third intermediate symmetric key along with
   * a token that has the same length as a key and
   * the two
   * transport data keys.
   This makes it 4 variables we derive based on 7 input variables.

   A random oracle is used to derive an intermediate type. We define
   functions to retrieve the four individual parts and equivalence that
   shows that all the four parts are indepently random. *)
type rom3_input_t.
fun rom3_input(G_t, G_t, G_t, G_t, G_t, G_t, psk_t): rom3_input_t [data].
type four_keys_t [large,fixed].

expand ROM_hash_refactored(
  (* types *)
  hashkey_t, (* Key of the hash function, models the choice of the hash
             (* function. *)
  rom3_input_t, (* input type *)
  four_keys_t,  (* intermediary output type of the actual random      *)
                (* oracle *)
  (* functions *)
  rom3_intermediate, (* name of the random oracle hash function,      *)
                     (* returning the intermediate type:              *)
   (* rom3_intermediate(hashkey_t, rom3_input_t): four_keys_t             *)
  (* processes *)
  rom3_oracle,        (* name of the oracle available to the attacker *)
  (* variables *)
  r_rom3,
  x_rom3,
  (* channels for random oracle *)
  ch1_rom3,
  ch2_rom3,
  (* parameters *)
  N_qH3            (* number of queries to the oracle by the attacker *)
).

   (* wrapper around the random oracle hash function,      *)
   (* returning the concatenation of 4 keys:               *)
   (* rom3(hashkey_t, G_t, G_t, G_t, G_t, G_t, G_t, psk_t): four_keys_t *)

letfun rom3(hk: hashkey_t, e_i: G_t, es: G_t, ss: G_t, e_r: G_t, ee: G_t, se: G_t, v_psk: psk_t) =
  rom3_intermediate(hk, rom3_input(e_i, es, ss, e_r, ee, se, v_psk)).

expand split_4(
  (* types *)
  four_keys_t,  (* intermediary output type of the actual random      *)
                (* oracle *)
  key_t,        (* All the 4 parts to extract have type key_t *)
  key_t,
  key_t,
  key_t,
  (* functions *)
  concat_four_keys, (* function that concatenates the four parts. *)
  (* variables *)
  tau, (* Names of the variables used when applying the equivalence   *)
  k,   (* This makes the games in the proof much more readable.       *)
  T_i_send,
  T_i_recv
).

  


(* A collision resistant hash function is used for chaining hashes
   with parts of the protocol transcript. To the previous hash output,
   the next part of the transcript is appended. We model this as being
   a hash function that has two arguments, the first being of type
   hashoutput and the second of bitstring. *)
proba P_hash.    (* probability of breaking collision resistance *)
expand CollisionResistant_hash_2(
  (* types *)
  hashkey_t,    (* key of the hash function, models the choice of *)
                (* the hash function *)
  hashoutput_t, (* first argument that gets hashed. See the comment *)
                (* just above this macro for an explanation. *)
  bitstring,    (* second argument that gets hashed. *)
  hashoutput_t, (* output type of the hash function *)
  (* functions *)
  hash,     (* name of the hash function: *)
            (* hash(hashkey_t, hashoutput_t, bitstring): hashoutput_t *)
  (* processes *)
  hash_oracle,  (* name of the oracle that will make available the *)
                (* hash key to the attacker *)
  (* parameters *)
  P_hash        (* probability of breaking collision resistance *)
).
(* constants used in the transcript hashing *)
const hash_construction_identifier : hashoutput_t.
  (* This is hash( hash("Noise_IK…") || "WireGuard v1 …" ), and it's *)
  (* the same for all parties, so no need to calculate it with hash() *)

(* channel names *)
channel c_start.
channel c_config_initiator, c_config_responder, c_final_leak.
channel c_init2resp_send, c_resp2init_recv.
channel c_init2resp_recv, c_resp2init_send.
channel c_keyconfirm_send, c_keyconfirm_recv, c_wait_before_2nd_part.
channel c_N_init_send_config, c_N_resp_send_config.
channel c_N_init_send, c_N_resp_send.
channel c_N_resp_recv, c_N_init_recv.
channel c_N_resp_recv_res, c_N_init_recv_res.
channel c_publickeys.
channel c_corrupt_S_i, c_corrupt_S_r, c_corrupt_psk.


(* WireGuard specific types *)
type G_set_t. (* set of public keys *)
fun is_in(G_t, G_set_t): bool.

type counter_t [fixed].     (* 8 byte counter in the data message *)
const counter_0: counter_t. (* constant for counter with value 0  *)
fun counter_to_nonce(counter_t) : nonce_t [data].
  (* This is [data] because it is an injection,
     from 8 byte counters to 12 bytes nonces. *)

type msg_type_t [fixed]. (* 1 byte msg type field *)
const msg_type_init2resp:    msg_type_t.
const msg_type_resp2init:    msg_type_t.
const msg_type_data:         msg_type_t.
const msg_type_cookie_reply: msg_type_t.

type reserved_t [fixed]. (* 3 byte reserved field *)
const reserved: reserved_t.

type session_index_t [fixed]. (* 4 byte session identifier field *)
type timestamp_t [fixed].     (* 12 byte timestamps *)
fun timestamp_to_bitstring(timestamp_t): bitstring [data].

table rcvd_timestamps(G_t, G_t, timestamp_t).
(* In the security game we want to be able to abort if the attacker
   wants us to encrypt with an already used nonce. *)
type side.
const is_initiator: side.
const is_responder: side.
const is_initiator_sw: side.
const is_responder_sw: side.
(* the bitstring is used as tuple (side, replication_index) *)
table sent_counters(bitstring, counter_t).
table rcvd_counters(bitstring, counter_t).

(* Convenience wrappers around hash that take care of
   type conversion. *)
letfun mix_hash_G(key_hash: hashkey_t, prev_hash: hashoutput_t, value: G_t) =
  hash(key_hash, prev_hash, G_to_bitstring(value)).

letfun mix_hash_bitstring(key_hash: hashkey_t, prev_hash: hashoutput_t, value: bitstring) =
  hash(key_hash, prev_hash, value).

letfun mix_hash_key(key_hash: hashkey_t, prev_hash: hashoutput_t, value: key_t) =
  hash(key_hash, prev_hash, key_to_bitstring(value)).

(* Convenience wrappers for enc and dec that take care
   of type conversions. *)
letfun enc_G(group_element: G_t, current_hash: hashoutput_t, k: key_t, n: nonce_t) =
  enc(G_to_bitstring(group_element), hashoutput_to_bitstring(current_hash), k, n).

letfun dec_ad_hash(ciphertext: bitstring, current_hash: hashoutput_t, k: key_t, n: nonce_t) =
  dec(ciphertext, hashoutput_to_bitstring(current_hash), k, n).

letfun enc_timestamp(timestamp: timestamp_t, current_hash: hashoutput_t, k: key_t, n: nonce_t) =
  enc(timestamp_to_bitstring(timestamp), hashoutput_to_bitstring(current_hash), k, n).

letfun enc_bitstring(plaintext: bitstring, current_hash: hashoutput_t, k: key_t, n: nonce_t) =
  enc(plaintext, hashoutput_to_bitstring(current_hash), k, n).

(* Define a function for choosing from two attacker-provided plaintexts based
   on a bit. Also, defines some equations on it so CryptoVerif is able
   to reason about it. *)
expand boolean_choice_for_encryption(
  (* types *)
  bitstring,  (* type of the values *)
  (* functions *)
  Zero,       (* the Zero function provided by the encryption scheme. *)
              (* Needed for some equations about the function. *)
  test	      (* Name of the choice function: *)
              (* test(bool, bitstring, bitstring): bitstring *)
).

event initiator_corr(
  (* values sent in the first protocol message *)
  G_t,                (* initiator's ephemeral public key *)
  bitstring,          (* msg.static *)
  bitstring,          (* msg.timestamp *)
  (* values sent in the second protocol message, that
     have not yet been sent in the first *)
  G_t,                (* responder's ephemeral public key *)
  bitstring,          (* encrypted empty bitstring *)
  (* values for which we want to prove correctness *)
  key_t,               (* T_i_send = T_r_recv *)
  key_t,               (* T_i_recv = T_r_send *)
  hashoutput_t         (* H7 *)
).
event responder_corr(
  (* values sent in the first protocol message *)
  G_t,                (* initiator's ephemeral public key *)
  bitstring,          (* msg.static *)
  bitstring,          (* msg.timestamp *)
  (* values sent in the second protocol message, that
     have not yet been sent in the first *)
  G_t,                (* responder's ephemeral public key *)
  bitstring,          (* encrypted empty bitstring *)
  (* values for which we want to prove correctness *)
  key_t,               (* T_i_send = T_r_recv *)
  key_t,               (* T_i_recv = T_r_send *)
  hashoutput_t         (* H7 *)
).

(* Correctness can be proved in a model with all keys (psk, long-term,
   and ephemeral) compromised. *)

(* Correctness: T_i_send = T_r_recv && T_i_recv = T_r_send. *)
query
  (* Values sent in the first protocol message. *)
  E_i_pub: G_t,
  static_i_enc: bitstring,
  timestamp_i_enc: bitstring,
  (* Values sent in the second protocol message, that
     have not yet been sent in the first. *)
  E_r_pub: G_t,
  empty_bitstring_enc: bitstring,
  (* Values for which we want to prove correctness. *)
  T_i_send: key_t,
  T_i_recv: key_t,
  T_r_recv: key_t,
  T_r_send: key_t,
  Hi7: hashoutput_t,        (* H7 *)
  Hr7: hashoutput_t;
  event(responder_corr(
	    E_i_pub, static_i_enc, timestamp_i_enc, 
	    E_r_pub, empty_bitstring_enc, 
	    T_r_recv, T_r_send, Hr7))
  &&
  event(initiator_corr(
	    E_i_pub, static_i_enc, timestamp_i_enc, 
	    E_r_pub, empty_bitstring_enc, 
	    T_i_send, T_i_recv, Hi7))
  ==> T_i_send = T_r_recv && T_i_recv = T_r_send && Hi7 = Hr7.


(* Type definitions for return values of functions that
   prepare or process messages 
   
   prepare1res_t is an "option" type:
   prepare1succ(...) is the success case
   prepare1fail is the failure case
   The disequation prepare1succ(x1,x2,x3,x4,x5,x6,x7,x8) <> prepare1fail 
   guarantees that there is no confusion between these two cases. *)
type prepare1res_t.
fun prepare1succ(session_index_t, Z_t, G_t, bitstring, bitstring, G_t, G_t, hashoutput_t): prepare1res_t [data].
const prepare1fail: prepare1res_t.
equation forall x1: session_index_t, x2: Z_t, x3: G_t, x4: bitstring, x5: bitstring, x6: G_t, x7: G_t, x8: hashoutput_t;
   prepare1succ(x1,x2,x3,x4,x5,x6,x7,x8) <> prepare1fail.

type process1res_t.
fun process1succ(G_t, G_t, G_t, hashoutput_t, timestamp_t):process1res_t [data].
const process1fail: process1res_t.
equation forall x1: G_t, x2: G_t, x3: G_t, x4: hashoutput_t, x5: timestamp_t;
   process1succ(x1,x2,x3,x4,x5) <> process1fail.

type prepare2res_t.
fun prepare2succ(session_index_t, Z_t, G_t, key_t, key_t, hashoutput_t, bitstring):prepare2res_t [data].
const prepare2fail: prepare2res_t.
equation forall x1: session_index_t, x2: Z_t, x3: G_t, x4: key_t, x5: key_t, x6: hashoutput_t, x7: bitstring;
   prepare2succ(x1,x2,x3,x4,x5,x6,x7) <> prepare2fail.
   
type process2res_t.
fun process2succ(key_t, key_t, hashoutput_t):process2res_t [data].
const process2fail: process2res_t.
equation forall x1: key_t, x2: key_t, x3: hashoutput_t;
   process2succ(x1,x2,x3) <> process2fail.

type preparemsgres_t.
fun preparemsgsucc(bitstring, bitstring):preparemsgres_t [data].
const preparemsgfail: preparemsgres_t.
equation forall x1: bitstring, x2: bitstring;
   preparemsgsucc(x1,x2) <> preparemsgfail.

type processmsgres_t.
fun processmsgsucc(bitstring):processmsgres_t [data].
const processmsgfail: processmsgres_t.
equation forall x: bitstring;
   processmsgsucc(x) <> processmsgfail.

(* Prepare the first message e, es, s, ss, {t} *)

letfun prepare1(
          key_hash: hashkey_t,
          key_rom1: hashkey_t,
          key_rom2: hashkey_t,
          S_X_pub_rcvd: G_t,
          S_i_priv: Z_t,
          S_i_pub: G_t,
          timestamp_i: timestamp_t) =

  new I_i: session_index_t;
  new E_i_priv: Z_t;
  (* TODO Why are we not using DH here? *)
  let E_i_pub: G_t = exp(g, E_i_priv) in

  let H_i1: hashoutput_t = mix_hash_G(key_hash, hash_construction_identifier, S_X_pub_rcvd) in
  let H_i2: hashoutput_t = mix_hash_G(key_hash, H_i1, E_i_pub) in

  let es_i: G_t = DH(S_X_pub_rcvd, E_i_priv) in
  ifdef(`m4_zero_test',`
  if es_i = zero then prepare1fail else
  ')
  let k_i2: key_t = rom1(key_rom1, E_i_pub, es_i) in

  let static_i_enc: bitstring = enc_G(S_i_pub, H_i2, k_i2, nonce_0) in
  let H_i3: hashoutput_t = mix_hash_bitstring(key_hash, H_i2, static_i_enc) in

  let ss_i: G_t = DH(S_X_pub_rcvd, S_i_priv) in
  ifdef(`m4_zero_test',`
  if ss_i = zero then prepare1fail else
  ')
  let k_i3: key_t = rom2(key_rom2, E_i_pub, es_i, ss_i) in

  let timestamp_i_enc: bitstring = enc_timestamp(timestamp_i, H_i3, k_i3, nonce_0) in

  let H_i4: hashoutput_t = mix_hash_bitstring(key_hash, H_i3, timestamp_i_enc) in

  prepare1succ(I_i, E_i_priv, E_i_pub, static_i_enc, timestamp_i_enc, es_i, ss_i, H_i4).

(* Process a received first message *)

letfun process1(
          key_hash: hashkey_t,
          key_rom1: hashkey_t,
          key_rom2: hashkey_t,
	  allowed_peers: G_set_t,
          S_r_priv: Z_t,
          S_r_pub: G_t,
          E_i_pub_rcvd: G_t,
          static_i_enc_rcvd: bitstring,
          timestamp_i_enc_rcvd: bitstring
        ) =


    let H_r1: hashoutput_t = mix_hash_G(key_hash, hash_construction_identifier, S_r_pub) in
    let H_r2: hashoutput_t = mix_hash_G(key_hash, H_r1, E_i_pub_rcvd) in

    let es_r: G_t = DH(E_i_pub_rcvd, S_r_priv) in
    ifdef(`m4_zero_test',`
    if es_r = zero then process1fail else
    ')
    let k_r2: key_t = rom1(key_rom1, E_i_pub_rcvd, es_r) in

    let injbot(G_to_bitstring(S_i_pub_rcvd: G_t)) = dec_ad_hash(static_i_enc_rcvd, H_r2, k_r2, nonce_0) in
    (
      if is_in(S_i_pub_rcvd, allowed_peers) then
      (
        let H_r3: hashoutput_t = mix_hash_bitstring(key_hash, H_r2, static_i_enc_rcvd) in

        let ss_r: G_t = DH(S_i_pub_rcvd, S_r_priv) in
        ifdef(`m4_zero_test',`
        if ss_r = zero then process1fail else
        ')
        let k_r3: key_t = rom2(key_rom2, E_i_pub_rcvd, es_r, ss_r) in

        let injbot(timestamp_to_bitstring(timestamp_i_rcvd: timestamp_t)) = dec_ad_hash(timestamp_i_enc_rcvd, H_r3, k_r3, nonce_0) in
        (
          let H_r4: hashoutput_t = mix_hash_bitstring(key_hash, H_r3, timestamp_i_enc_rcvd) in
          process1succ(es_r, ss_r, S_i_pub_rcvd, H_r4, timestamp_i_rcvd)
        ) else (
          (* timestamp did not decrypt *)
          process1fail
        )
      ) else (
        (* peer not allowed *)
        process1fail
      )      
    ) else (
      (* static did not decrypt *)
      process1fail
    ).

(* Prepare the second message e, ee, se, psk, {} *)

letfun prepare2(
    key_hash: hashkey_t,
    key_rom3: hashkey_t,
    S_i_pub_rcvd: G_t, E_i_pub_rcvd: G_t,
    H_r4: hashoutput_t, es_r: G_t, ss_r: G_t, Q: psk_t) =

  new I_r: session_index_t;
  new E_r_priv: Z_t;
  let E_r_pub: G_t = exp(g, E_r_priv) in

  let ee_r: G_t = DH(E_i_pub_rcvd, E_r_priv) in
  ifdef(`m4_zero_test',`
  if ee_r = zero then prepare2fail else
  ')
  let se_r: G_t = DH(S_i_pub_rcvd, E_r_priv) in
  ifdef(`m4_zero_test',`
  if se_r = zero then prepare2fail else
  ')

  let H_r5: hashoutput_t = mix_hash_G(key_hash, H_r4, E_r_pub) in

  let concat_four_keys(tau_r4: key_t, k_r4: key_t, T_r_recv: key_t, T_r_send: key_t) = rom3(key_rom3, E_i_pub_rcvd, es_r, ss_r, E_r_pub, ee_r, se_r, Q) in
  (
    let H_r6: hashoutput_t = mix_hash_key(key_hash, H_r5, tau_r4) in

    let empty_bitstring_r_enc: bitstring = enc_bitstring(empty_bitstring, H_r6, k_r4, nonce_0) in
    let H_r7: hashoutput_t = mix_hash_bitstring(key_hash, H_r6, empty_bitstring_r_enc) in

    prepare2succ(I_r, E_r_priv, E_r_pub, T_r_recv, T_r_send, H_r7, empty_bitstring_r_enc)
  ) else (
    prepare2fail
  ).

(* Process a received second message *)

letfun process2(
          key_hash: hashkey_t,
          key_rom3: hashkey_t,
          E_i_priv: Z_t, E_i_pub: G_t, S_i_priv: Z_t, S_i_pub: G_t,
          E_r_pub_rcvd: G_t, empty_bitstring_r_enc_rcvd: bitstring, 
          H_i4: hashoutput_t,
          es_i: G_t, ss_i: G_t, Q: psk_t) =
  let ee_i: G_t = DH(E_r_pub_rcvd, E_i_priv) in
  ifdef(`m4_zero_test',`
  if ee_i = zero then process2fail else
  ')
  let se_i: G_t = DH(E_r_pub_rcvd, S_i_priv) in
  ifdef(`m4_zero_test',`
  if se_i = zero then process2fail else
  ')

  let H_i5: hashoutput_t = mix_hash_G(key_hash, H_i4, E_r_pub_rcvd) in

  let concat_four_keys(tau_i4: key_t, k_i4: key_t, T_i_send: key_t, T_i_recv: key_t) = rom3(key_rom3, E_i_pub, es_i, ss_i, E_r_pub_rcvd, ee_i, se_i, Q) in
  (
    let H_i6: hashoutput_t = mix_hash_key(key_hash, H_i5, tau_i4) in

    let injbot(=empty_bitstring) = dec_ad_hash(empty_bitstring_r_enc_rcvd, H_i6, k_i4, nonce_0) in
    (
      let H_i7: hashoutput_t = mix_hash_bitstring(key_hash, H_i6, empty_bitstring_r_enc_rcvd) in
      process2succ(T_i_send, T_i_recv, H_i7)
    ) else (
      (* empty_bitstring_r_enc_rcvd did not decrypt *)
      process2fail
    )
  ) else (
    (* weird case where the rom3 pattern matching did not work - actually never happens *)
    process2fail
  ).


(* Prepare the payload messages *)

letfun prepare_msg(
    side_index: bitstring, secret_bit_I: bool,
    plaintext_0: bitstring, plaintext_1: bitstring, counter: counter_t,
    honest: bool,
    T_i_send: key_t) =

  if Zero(plaintext_0) = Zero(plaintext_1) && (honest || (plaintext_0 = plaintext_1)) then
  (
    get sent_counters(=side_index, =counter) in preparemsgfail else
    insert sent_counters(side_index, counter);

    let plaintext = test(secret_bit_I, plaintext_0, plaintext_1) in
    let ciphertext = enc(plaintext, empty_bitstring, T_i_send, counter_to_nonce(counter)) in
    preparemsgsucc(ciphertext, plaintext)
  ) else (
    preparemsgfail
  ).

(* Process the payload messages *)

letfun process_msg(
    side_index: bitstring, counter_rcvd: counter_t,
    ciphertext_rcvd: bitstring, T_i_recv: key_t) =

  get rcvd_counters(=side_index, =counter_rcvd) in processmsgfail else
  insert rcvd_counters(side_index, counter_rcvd);
  let injbot(plaintext) = dec(ciphertext_rcvd, empty_bitstring, T_i_recv, counter_to_nonce(counter_rcvd)) in
  (
    processmsgsucc(plaintext)
  ) else (
    (* decryption failed *)
    processmsgfail
  ).

(* The initiator *)

let initiator_responder(key_hash: hashkey_t,
              key_rom1: hashkey_t, key_rom2: hashkey_t, key_rom3: hashkey_t,
              S_i_priv: Z_t, S_i_pub: G_t, S_r_priv: Z_t, S_r_pub: G_t) =

  ! i <= N_sessions
  (* Initiator. Receive the public key of the responder we should communicate *)
  (* with, and the timestamp the initiator should use. *)
  (* The adversary also chooses the pre-shared key Q_i *)
  in(c_config_initiator, (S_X_pub: G_t, timestamp_i: timestamp_t , Q_i: psk_t));
  

    let prepare1succ(I_i: session_index_t,
         E_i_priv: Z_t, E_i_pub_sent: G_t, static_i_enc: bitstring,
         timestamp_i_enc: bitstring, 
         es_i: G_t, ss_i: G_t, H_i4: hashoutput_t) =
      prepare1(key_hash, key_rom1, key_rom2, S_X_pub, S_i_priv, S_i_pub, timestamp_i) in


    out(c_init2resp_send, (E_i_priv (* Leak ephemeral *),
                           (msg_type_init2resp, reserved, I_i, E_i_pub_sent, static_i_enc, timestamp_i_enc) (* First message: initiator -> responder *)));

    (* Now the responder. Configuration: the adversary chooses the pre-shared key (Q_r)
       and the set of initiators it accepts connections from (allowed_peers);
       We allow the adversary to alter I_i, that does not change anything. *)
    in(c_config_responder, (Q_r: psk_t, allowed_peers: G_set_t, I_i_rcvd: session_index_t));
    (* The responder receives the message sent by the initiator *)
    let (=msg_type_init2resp, =reserved, E_i_pub_rcvd: G_t, static_i_enc_rcvd: bitstring, timestamp_i_enc_rcvd: bitstring) =
      (msg_type_init2resp, reserved, E_i_pub_sent, static_i_enc, timestamp_i_enc) in
    
    let process1succ(es_r: G_t, ss_r: G_t, S_i_pub_rcvd: G_t, H_r4: hashoutput_t, timestamp_i_rcvd: timestamp_t) =
      process1(
        key_hash, key_rom1, key_rom2, allowed_peers, S_r_priv, S_r_pub,
        E_i_pub_rcvd, static_i_enc_rcvd, timestamp_i_enc_rcvd) in
    
ifdef(`m4_replay_prot',`
    get rcvd_timestamps(=S_r_pub, =S_i_pub_rcvd, =timestamp_i_rcvd) in yield else
    insert rcvd_timestamps(S_r_pub, S_i_pub_rcvd, timestamp_i_rcvd);
')

    let prepare2succ(I_r: session_index_t, E_r_priv: Z_t, E_r_pub_sent: G_t, T_r_recv: key_t, T_r_send: key_t, H_r7: hashoutput_t, empty_bitstring_enc: bitstring) =
      prepare2(key_hash, key_rom3,
        S_i_pub_rcvd, E_i_pub_rcvd, H_r4, es_r, ss_r, Q_r) in
    event responder_corr(E_i_pub_rcvd, static_i_enc_rcvd, timestamp_i_enc_rcvd, E_r_pub_sent, empty_bitstring_enc, T_r_recv, T_r_send, H_r7);

    (* Back in the initiator: the initiator receives the message sent by the responder *)
    let (=msg_type_resp2init, =reserved, E_r_pub_rcvd: G_t, empty_bitstring_r_enc_rcvd: bitstring) =
      (msg_type_resp2init, reserved, E_r_pub_sent, empty_bitstring_enc)
    in

    let process2succ(T_i_send: key_t, T_i_recv: key_t, H_i7: hashoutput_t) =
      process2(key_hash, key_rom3, E_i_priv, E_i_pub_sent, S_i_priv, S_i_pub, E_r_pub_rcvd, empty_bitstring_r_enc_rcvd, H_i4, es_i, ss_i, Q_i) in
    event initiator_corr(E_i_pub_sent, static_i_enc, timestamp_i_enc, E_r_pub_rcvd, empty_bitstring_r_enc_rcvd, T_i_send, T_i_recv, H_i7);

    out(c_final_leak, (E_r_priv (* Leak ephemeral *),
    		       (msg_type_resp2init, reserved, I_r, I_i_rcvd, E_r_pub_sent, empty_bitstring_enc) (* Message 2 responder -> initiator *),
		       (T_i_send, T_i_recv), (T_r_recv, T_r_send) (* Keys so that the adversary can continue the protocol *))).



process
  in(c_start, ());
  new key_hash: hashkey_t;
  new key_rom1: hashkey_t;
  new key_rom2: hashkey_t;
  new key_rom3: hashkey_t;

  new S_A_priv: Z_t;
  let S_A_pub = exp(g, S_A_priv) in
  new S_B_priv: Z_t;
  let S_B_pub = exp(g, S_B_priv) in

  (* hand over control to the attacker *)
  out(c_publickeys, (S_A_priv, S_B_priv, S_A_pub, S_B_pub));

  (initiator_responder(key_hash, key_rom1, key_rom2, key_rom3, S_A_priv, S_A_pub, S_B_priv, S_B_pub) |
   rom1_oracle(key_rom1) | rom2_oracle(key_rom2) | rom3_oracle(key_rom3) |
   hash_oracle(key_hash)
  )

(* The adversary has all keys (long-term, ephemeral, psk), so it can run other sessions,
   including sessions with swapped initiator and responder, by itself. *)

ifdef(`m4_replay_prot',`
(* EXPECTED
All queries proved.
3.012s (user 2.984s + system 0.028s), max rss 29208K
END *)
',`
(* EXPECTED
All queries proved.
2.484s (user 2.456s + system 0.028s), max rss 29328K
END *)
')