#define CHAINING_KEY_EVENTS 1
#define MESSAGE_TRANSMISSION_EVENTS 0
#define SESSION_START_EVENTS 0
#define RANDOMIZED_CALL_IDS 0
#define COOKIE_EVENTS 1
#define KEM_EVENTS 1

#include "config.mpv"
#include "prelude/basic.mpv"
#include "crypto/key.mpv"
#include "crypto/kem.mpv"
#include "rosenpass/handshake_state.mpv"

/* The cookie data structure is implemented based on the WireGuard protocol.
 * The ip and port is based purely on the public key and the implementation of the private cookie key is intended to mirror the biscuit key.
 * The code tests the response to a possible DOS attack by setting up alternative branches for the protocol
 * processes: Oinit_conf, Oinit_hello and resp_hello to simulate what happens when the responder or initiator is overloaded.
 * When under heavy load a valid cookie is required. When such a cookie is not present a cookie message is sent as a response.
 * Queries then test to make sure that expensive KEM operations are only conducted after a cookie has been successfully validated.
 */

type CookieMsg_t.
fun CookieMsg(
  SessionId,  // sender
  bits,     // nonce
  bits      // cookie
) : CookieMsg_t [data].

#define COOKIE_EVENTS(eventLbl) \
  COOKIE_EV(event MCAT(eventLbl, _UnderLoadEV) (SessionId, SessionId, Atom).)   \
  COOKIE_EV(event MCAT(eventLbl, _CookieValidated) (SessionId, SessionId, Atom).)   \
  COOKIE_EV(event MCAT(eventLbl, _CookieSent) (SessionId, SessionId, Atom, CookieMsg_t).)

fun cookie_key(kem_sk) : key [private].
fun ip_and_port(kem_pk):bits.
letfun create_mac2_key(sskm:kem_sk, spkt:kem_pk) = prf(cookie_key(sskm), ip_and_port(spkt)).
letfun create_cookie(sskm:kem_sk, spkm:kem_pk, spkt:kem_pk, nonce:bits, msg:bits) = xaead_enc(lprf2(COOKIE, kem_pk2b(spkm), nonce),
    k2b(create_mac2_key(sskm, spkm)), msg).

#define COOKIE_PROCESS(eventLbl, innerFunc) \
  new nonce:bits; \
  in(C, Ccookie(mac1, mac2)); \
  COOKIE_EV(event MCAT(eventLbl, _UnderLoadEV) (sidi, sidr, call);)  \
  msgB <- Envelope(mac1, msg); \
  mac2_key <- create_mac2_key(sskm, spkt); \
  if k2b(create_mac2(mac2_key, msgB)) = mac2 then  \
    COOKIE_EV(event MCAT(eventLbl, _CookieValidated) (sidi, sidr, call);)  \
    innerFunc \
  else \
    cookie <- create_cookie(sskm, spkm, spkt, nonce, msg); \
    cookie_msg <- CookieMsg(sidi, nonce, cookie); \
    COOKIE_EV(event MCAT(eventLbl, _CookieSent) (sidi, sidr, call, cookie_msg);)  \
    out(C, cookie_msg). \

#include "rosenpass/oracles.mpv"

#include "rosenpass/responder.macro"
COOKIE_EVENTS(Oinit_conf)
let Oinit_conf_underLoad() =
  in(C, Cinit_conf(Ssskm, Spsk, Sspkt, ic));
  in(C, last_cookie:bits);

  msg <- IC2b(ic);
  let InitConf(sidi, sidr, biscuit, auth) = ic in 

  new call:Atom;

  SETUP_HANDSHAKE_STATE()

  COOKIE_PROCESS(Oinit_conf, Oinit_conf_inner(Ssskm, Spsk, Sspkt, ic, call))

#include "rosenpass/responder.macro"
COOKIE_EVENTS(Oinit_hello)
let Oinit_hello_underLoad() =

  in(C, Cinit_hello(sidr, biscuit_no, Ssskm, Spsk, Sspkt, Septi, Sspti, ih));
  in(C, Oinit_hello_last_cookie:key);
  new call:Atom;

  msg <- IH2b(ih);
  let InitHello(sidi, epki, sctr, pidic, auth) = ih in
  SETUP_HANDSHAKE_STATE()

  COOKIE_PROCESS(Oinit_hello, Oinit_hello_inner(sidr, biscuit_no, Ssskm, Spsk, Sspkt, Septi, Sspti, ih, Oinit_hello_last_cookie, C, call))

let rosenpass_dos_main() = 0
  | !Oreveal_kem_pk
  | REP(INITIATOR_BOUND, Oinitiator)
  | REP(RESPONDER_BOUND, Oinit_hello)
  | REP(RESPONDER_BOUND, Oinit_conf)
  | REP(RESPONDER_BOUND, Oinit_hello_underLoad)
  | REP(RESPONDER_BOUND, Oinit_conf_underLoad).

let main = rosenpass_dos_main.

select cookie:CookieMsg_t;         attacker(cookie)/6220[hypothesis].
nounif v:key;         attacker(prepare_key( v ))/6217[hypothesis].
nounif v:seed;        attacker(prepare_seed( v ))/6216[hypothesis].
nounif v:seed;        attacker(prepare_seed( v ))/6216[hypothesis].
nounif v:seed;        attacker(rng_kem_sk( v ))/6215[hypothesis].
nounif v:seed;        attacker(rng_key( v ))/6214[hypothesis].
nounif v:kem_sk;      attacker(prepare_kem_sk( v ))/6210[hypothesis].

// nounif Spk:kem_sk_tmpl;
//   attacker(Creveal_kem_pk(Spk))/6110[conclusion].
// nounif sid:SessionId, Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, Seski:seed_tmpl, Ssptr:seed_tmpl;
//   attacker(Cinitiator( *sid, *Ssskm, *Spsk, *Sspkt, *Seski, *Ssptr ))/6109[conclusion].
// nounif sid:SessionId, biscuit_no:Atom, Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, Septi:seed_tmpl, Sspti:seed_tmpl, ih:InitHello_t;
//   attacker(Cinit_hello( *sid, *biscuit_no, *Ssskm, *Spsk, *Sspkt, *Septi, *Sspti, *ih ))/6108[conclusion].
nounif rh:RespHello_t;
  attacker(Cresp_hello( *rh ))/6107[conclusion].
nounif Ssskm:kem_sk_tmpl, Spsk:key_tmpl, Sspkt:kem_sk_tmpl, ic:InitConf_t;
  attacker(Cinit_conf( *Ssskm, *Spsk, *Sspkt, *ic ))/6106[conclusion].

@reachable "DOS protection: cookie sent"
query sidi:SessionId, sidr:SessionId, call:Atom, cookieMsg:CookieMsg_t;
  event (Oinit_hello_CookieSent(sidi, sidr, call, cookieMsg)).

@lemma "DOS protection: Oinit_hello kem use when under load implies validated cookie"
lemma sidi:SessionId, sidr:SessionId, call:Atom;
event(Oinit_hello_UnderLoadEV(sidi, sidr, call))
  && event(Oinit_hello_KemUse(sidi, sidr, call))
  ==> event(Oinit_hello_CookieValidated(sidi, sidr, call)).

@lemma "DOS protection: Oinit_conf kem use when under load implies validated cookie"
lemma sidi:SessionId, sidr:SessionId, call:Atom;
event(Oinit_conf_UnderLoadEV(sidi, sidr, call))
  && event(Oinit_conf_KemUse(sidi, sidr, call))
  ==> event(Oinit_conf_CookieValidated(sidi, sidr, call)).

@lemma "DOS protection: Oresp_hello kem use when under load implies validated cookie"
lemma sidi:SessionId, sidr:SessionId, call:Atom;
event(Oresp_hello_UnderLoadEV(sidi, sidr, call))
  && event(Oresp_hello_KemUse(sidi, sidr, call))
  ==> event(Oresp_hello_CookieValidated(sidi, sidr, call)).