val encrypt_spec
      (a: Spec.cipher_alg)
      (l: pne_plain_length)
      (pn: lbytes l)
      (b: length_bits l)
      (s: sample)
      (k: Spec.key a)
    : pne_cipher l
let encrypt_spec (a: Spec.cipher_alg)
  (l: pne_plain_length)
  (pn: lbytes l)
  (b: length_bits l)
  (s: sample)
  (k: Spec.key a):
  pne_cipher l
=
  let open QUIC.Spec.Lemmas in
  // We need the packet number length in order to know where to find the mask in the cipher block.
  let mask = Model.Helpers.reveal #16 (QUIC.TotSpec.block_of_sample a k (Model.Helpers.hide s)) in
  let pnmask = and_inplace (Seq.slice mask 1 (l + 1)) (QUIC.Spec.Header.pn_sizemask (l - 1)) 0 in
  // Classify, because HACL* specs require secret integers.
  let pnmask = Seq.init (Seq.length pnmask) (fun i -> Lib.IntTypes.u8 (UInt8.v (Seq.index pnmask i))) in
  assert (Seq.length pnmask == l);
  // We now have enough to build the encrypted pn. Note that because our
  // input is a sliced packet number, we don't mask at an offset like
  // header_encrypt does.
  let encrypted_pn = pointwise_op (Lib.IntTypes.(logxor #U8 #SEC)) pn pnmask 0 in
  // Now on to bit protection. Since we receive as input only the protected
  // bits, there is a proof obligation for a caller, to show that get_bf
  // (header.[0] `xor` mask.[0]) == get_bf header.[0] `xor` get_bf mask.[0]
  let mask_bits: bits = LowParse.BitFields.get_bitfield (UInt8.v (Seq.index mask 0)) 0 5 in
  let protected_bits = mask_bits `FStar.UInt.logxor` b in
  QUIC.UInt.lemma_logxor_lt #8 mask_bits b 5;
  encrypted_pn, protected_bits
val encrypt_spec
      (a: Spec.cipher_alg)
      (l: pne_plain_length)
      (pn: lbytes l)
      (b: length_bits l)
      (s: sample)
      (k: Spec.key a)
    : pne_cipher l
let encrypt_spec
      (a: Spec.cipher_alg)
      (l: pne_plain_length)
      (pn: lbytes l)
      (b: length_bits l)
      (s: sample)
      (k: Spec.key a)
    : pne_cipher l =
let open QUIC.Spec.Lemmas in
let mask = Model.Helpers.reveal #16 (QUIC.TotSpec.block_of_sample a k (Model.Helpers.hide s)) in
let pnmask = and_inplace (Seq.slice mask 1 (l + 1)) (QUIC.Spec.Header.pn_sizemask (l - 1)) 0 in
let pnmask =
  Seq.init (Seq.length pnmask) (fun i -> Lib.IntTypes.u8 (UInt8.v (Seq.index pnmask i)))
in
assert (Seq.length pnmask == l);
let encrypted_pn = pointwise_op (let open Lib.IntTypes in logxor #U8 #SEC) pn pnmask 0 in
let mask_bits:bits = LowParse.BitFields.get_bitfield (UInt8.v (Seq.index mask 0)) 0 5 in
let protected_bits = mask_bits `FStar.UInt.logxor` b in
QUIC.UInt.lemma_logxor_lt #8 mask_bits b 5;
encrypted_pn, protected_bits
Model.PNE.fsti.checked
Spec.Hash.Definitions.fst.checked Spec.Agile.Cipher.fsti.checked QUIC.UInt.fst.checked QUIC.TotSpec.fsti.checked QUIC.Spec.Lemmas.fst.checked QUIC.Spec.Header.fsti.checked QUIC.Spec.fsti.checked prims.fst.checked Model.Indexing.fsti.checked Model.Helpers.fsti.checked Mem.fst.checked LowStar.Buffer.fst.checked LowParse.BitFields.fsti.checked Lib.IntTypes.fsti.checked FStar.UInt8.fsti.checked FStar.UInt32.fsti.checked FStar.UInt128.fsti.checked FStar.UInt.fsti.checked FStar.Seq.fst.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.HyperStack.fst.checked
Model.PNE.fsti
Spec.Agile.Cipher.cipher_alg Model.PNE.pne_plain_length Model.Helpers.lbytes Model.PNE.length_bits Model.PNE.sample Spec.Agile.Cipher.key FStar.Pervasives.Native.Mktuple2 Model.PNE.bits Prims.unit QUIC.UInt.lemma_logxor_lt FStar.UInt.uint_t FStar.UInt.logxor LowParse.BitFields.get_bitfield FStar.UInt8.n FStar.UInt8.v FStar.Seq.Base.index QUIC.Spec.Base.byte FStar.Seq.Base.seq Lib.IntTypes.int_t Lib.IntTypes.U8 Lib.IntTypes.SEC QUIC.Spec.Lemmas.pointwise_op Lib.IntTypes.uint8 Lib.IntTypes.logxor Prims._assert Prims.eq2 Prims.nat FStar.Seq.Base.length FStar.Seq.Base.init Prims.b2t Prims.op_LessThan Lib.IntTypes.u8 QUIC.Spec.Base.bytes QUIC.Spec.Lemmas.and_inplace FStar.Seq.Base.slice Prims.op_Addition QUIC.Spec.Header.pn_sizemask Prims.op_Subtraction QUIC.Spec.Base.lbytes Model.Helpers.reveal QUIC.TotSpec.block_of_sample Model.Helpers.hide Model.PNE.pne_cipher
module Model.PNE

module B = LowStar.Buffer
module HS = FStar.HyperStack
module I = Model.Indexing
module U32 = FStar.UInt32
module U128 = FStar.UInt128
module Spec = Spec.Agile.Cipher


#set-options "--fuel 0 --ifuel 0"

open FStar.HyperStack
open Mem
open Model.Helpers

/// Abbreviations for idealization
/// ------------------------------

type id = I.pne_id
type alg = I.ca

let is_safe (i:id) =
  I.is_pne_honest i && I.ideal_PNE

type safe_id =
  i:id{is_safe i}

type unsafe_id =
  i:id{not (is_safe i)}

/// QUIC payload sampling
/// --------------------

// For simplicity, we do not distinguish between long and short headers; caller
// of this module will just 0-left-pad the protected bits in the case of a long
// header.
let bits: Type0 = LowParse.BitFields.ubitfield 8 5

// Note: the sample is PUBLIC so is using QUIC.Spec.lbytes which do not operate over secret integers.
let sample_length = 16
let sample = QUIC.Spec.lbytes sample_length
let pne_plain_length = l:nat {l >= 1 /\ l <= 4}

let length_bits (l: pne_plain_length) =
  b:bits { LowParse.BitFields.get_bitfield b 0 2 + 1 == l }

let pne_cipher (l:pne_plain_length) = lbytes l & bits
let pne_cipherpad = lbytes 4 & bits

/// Restrict a generated cipherpad to the length of the encoded packet number.
val clip_cipherpad : (cp:pne_cipherpad) -> (l:pne_plain_length) -> pne_cipher l

// We define a separate package type here (different from the one in AEAD). We
// idealize them separately, and there's a specific part here that's about the
// protected bits of the header.
noeq type pne_plain_pkg =
  | PNEPlainPkg:
    pne_plain: (j:id -> l:pne_plain_length -> t:eqtype) ->
    as_bytes: (j:id -> l:pne_plain_length -> pne_plain j l -> GTot (lbytes l & length_bits l)) ->
    repr: (j:unsafe_id -> l:pne_plain_length -> n:pne_plain j l ->
      Tot (b:(lbytes l & length_bits l) {b == as_bytes j l n})) ->
    mk: (j:id -> l:pne_plain_length -> n:lbytes l -> b:length_bits l -> p:pne_plain j l { as_bytes j l p == (n, b) }) ->
    xor: (j:id -> l:pne_plain_length -> p:pne_plain j l -> cp:pne_cipherpad ->
      (l': pne_plain_length & pne_plain j l')) ->
    lemma_xor: (j:id -> l:pne_plain_length -> p:pne_plain j l -> cp:pne_cipherpad -> Lemma
        (requires True)
        (ensures (
          let (| l', p' |) = xor j l p cp in
          let _, bits = as_bytes j l' p' in
          let l: pne_plain_length = LowParse.BitFields.get_bitfield bits 0 2 + 1 in
          let l': pne_plain_length = l' in
          l == l'
        ))) ->
    pne_plain_pkg

noeq type info' = {
  calg: alg;
  halg: I.ha;
  plain: pne_plain_pkg;
}

let info (j:id) =
  info:info'{
    I.pne_id_ginfo j == info.calg /\
    I.pne_id_ghash j == info.halg
  }

let pne_plain (#j:id) (u:info j) (l:pne_plain_length) =
  PNEPlainPkg?.pne_plain u.plain j l

noeq
type entry (#j:id) (u:info j) =
  | Entry :
    s:sample ->
    #l:pne_plain_length ->
    n:pne_plain u l ->
    // We do not perform truncation at encryption-time, but rather at decryption-time
    c:pne_cipherpad ->
    entry u

val pne_state : (#j:id) -> (u:info j) -> Type u#1

val table : (#j:safe_id) -> (#u:info j) -> (st:pne_state u) -> (h:mem) -> GTot (Seq.seq (entry u))

// Header protection key
let key_len (#j:id) (u:info j) = Spec.Agile.Cipher.key_length u.calg
val key: #j:unsafe_id -> #u:info j -> st:pne_state u -> lbytes (key_len u)

val footprint : #j:id -> #u:info j -> st:pne_state u -> GTot B.loc

val invariant: #j:id -> #u:info j -> st:pne_state u -> mem -> Type0

val frame_invariant: #j:id -> #u:info j -> st:pne_state u -> l:B.loc -> h0:mem -> h1:mem -> Lemma
  (requires (
    invariant st h0 /\
    B.modifies l h0 h1 /\
    B.loc_disjoint l (footprint st)))
  (ensures (
    invariant st h1))

val frame_table: #j:safe_id -> #u:info j -> st:pne_state u ->
  r:B.loc -> h0:mem -> h1:mem ->
  Lemma
    (requires
      invariant st h0 /\
      B.modifies r h0 h1 /\
      r `B.loc_disjoint` (footprint st))
    (ensures table st h1 == table st h0)

// Equality for the pair of (pne, protected_bits) we manipulate in this module.
let pne_bits_eq #l (x y: pne_cipher l) =
  fst x `lbytes_eq` fst y && snd x = snd y

// Finding the entry for a given sample.
let sample_filter (#j:id) (u:info j) (s:sample) (e:entry u) : bool =
  Entry?.s e = s

let entry_for_sample (#j:safe_id) (#u:info j) (s:sample) (st:pne_state u) (h:mem) :
  GTot (option (entry u)) =
  Seq.find_l (sample_filter u s) (table st h)

/// XXX unused
let contains_sample (#j:safe_id) (#u:info j) (s:sample) (st:pne_state u) (h:mem) :
  GTot bool =
  Some? (entry_for_sample s st h)

let fresh_sample (#j:safe_id) (#u:info j) (s:sample) (st:pne_state u) (h:mem) :
  GTot bool =
  None? (entry_for_sample s st h)


/// Stateful API
/// ------------

val create (j:id) (u:info j) : ST (pne_state u)
  (requires fun _ -> True)
  (ensures fun h0 st h1 ->
    invariant st h1 /\
    B.modifies B.loc_none h0 h1 /\
    B.fresh_loc (footprint st) h0 h1 /\
    B.(loc_includes (loc_pne_region ()) (footprint st)) /\

    (is_safe j ==> table st h1 == Seq.empty))

let lemma_max_hash_len ha
  : Lemma (Spec.Hash.Definitions.hash_length ha <= 64 /\
  (if Some? (Spec.Hash.Definitions.max_input_length ha) then
    Some?.v (Spec.Hash.Definitions.max_input_length ha) >= pow2 61 - 1
  else
    True) /\
  pow2 61 - 1 > 64)
  [SMTPat (Spec.Hash.Definitions.hash_length ha)]
  =
  assert_norm (pow2 61 < pow2 125);
  assert_norm (pow2 61 - 1 > 64);
  assert_norm (pow2 64 > pow2 61);
  assert_norm (pow2 128 > pow2 64)

let traffic_secret ha =
  lbytes (Spec.Hash.Definitions.hash_length ha)

val coerce (j:unsafe_id) (u:info j)
  (k:lbytes (key_len u))
  : ST (pne_state #j u)
  (requires fun _ -> True)
  (ensures fun h0 st h1 ->
    invariant st h1 /\
    B.modifies B.loc_none h0 h1 /\
    B.fresh_loc (footprint st) h0 h1 /\
    B.(loc_includes (loc_pne_region ()) (footprint st)) /\
    key st == k)

val quic_coerce (j:unsafe_id) (u:info j)
  (ts:traffic_secret u.halg)
  : ST (pne_state #j u)
  (requires fun _ -> True)
  (ensures fun h0 st h1 ->
    invariant st h1 /\
    B.modifies B.loc_none h0 h1 /\
    B.fresh_loc (footprint st) h0 h1 /\
    B.(loc_includes (loc_pne_region ()) (footprint st)) /\
    key st ==
      QUIC.Spec.derive_secret u.halg ts
        QUIC.Spec.label_hp (key_len u))

let encrypt_spec (a: Spec.cipher_alg)
  (l: pne_plain_length)
  (pn: lbytes l)
  (b: length_bits l)
  (s: sample)
  (k: Spec.key a):
  pne_cipher l
val encrypt_spec
      (a: Spec.cipher_alg)
      (l: pne_plain_length)
      (pn: lbytes l)
      (b: length_bits l)
      (s: sample)
      (k: Spec.key a)
    : pne_cipher l

    a: Spec.Agile.Cipher.cipher_alg ->
    l: Model.PNE.pne_plain_length ->
    pn: Model.Helpers.lbytes l ->
    b: Model.PNE.length_bits l ->
    s: Model.PNE.sample ->
    k: Spec.Agile.Cipher.key a
  -> Model.PNE.pne_cipher l