val gen: parent:rgn -> i:stae_id -> ST (writer i)
  (requires (fun h0 -> True))
  (ensures (genPost parent))
val gen: parent:rgn -> i:stae_id -> ST (writer i)
  (requires (fun h0 -> True))
  (ensures (genPost parent))
let gen parent i =
  if is_stream i then
    Stream () (Stream.gen parent i)
  else
    StLHAE () (StLHAE.gen parent i)
(**
Authenticated encryptions of streams of TLS fragments (from Content)
multiplexing StatefulLHAE and StreamAE with (some) length hiding
(for now, under-specifying ciphertexts lengths and values)
*)
module MiTLS.StAE
open MiTLS
module HST = FStar.HyperStack.ST //Added automatically


open FStar.HyperStack
open FStar.Bytes

open MiTLS.Mem
open MiTLS.TLSConstants
open MiTLS.TLSInfo


module HS   = FStar.HyperStack

module MS   = FStar.Monotonic.Seq
module C    = MiTLS.Content

module Stream = MiTLS.StreamAE
module StLHAE = MiTLS.StatefulLHAE
module Range = MiTLS.Range
#set-options "--initial_fuel 0 --max_fuel 0 --initial_ifuel 1 --max_ifuel 1"

////////////////////////////////////////////////////////////////////////////////
//Distinguishing the two multiplexing choices of StAE based on the ids
////////////////////////////////////////////////////////////////////////////////
let is_stream i = ID13? i

let is_stlhae (i:id) =
  if ID12? i && AEAD? (aeAlg_of_id i) then
    let is_aes128 = AEAD?._0 (aeAlg_of_id i) = EverCrypt.AES128_GCM in
    let is_aes256 = AEAD?._0 (aeAlg_of_id i) = EverCrypt.AES256_GCM in
    is_aes128 || is_aes256
  else
    false

// type id = i:id {is_stream i \/ is_stlhae i}

// PLAINTEXTS are defined in Content.fragment i
//16-06-08 see also StreamPlain and StatefulPlain.

type stae_id = i:id {is_stream i \/ is_stlhae i}

////////////////////////////////////////////////////////////////////////////////
//Various utilities related to lengths of ciphers and fragments
////////////////////////////////////////////////////////////////////////////////

let frag_plain_len (#i:id{is_stream i}) (f:C.fragment i): StreamPlain.plainLen =
  let r = snd (C.rg i f) + 1 in
  assume (StreamPlain.plainLength r);
  r

let frag_cipher_len (#i:id{is_stream i}) (f:C.fragment i) =
  frag_plain_len f + Stream.ltag i

// CONCRETE KEY MATERIALS, for leaking & coercing.
// (each implementation splits it into encryption keys, IVs, MAC keys, etc)
// ADL: this can now be factored going through the common AEADProvider interface
let aeKeySize (i:stae_id) =
  if is_stream i
  then
    UInt32.v (EverCrypt.aead_keyLen (Stream.alg i)) +
    AEADProvider.iv_length i
  else
    UInt32.v (EverCrypt.aead_keyLen (AEAD?._0 (aeAlg_of_id i))) +
    AEADProvider.salt_length i

type keyBytes (i:stae_id) = lbytes (aeKeySize i)


////////////////////////////////////////////////////////////////////////////////
//`state i rw`, a sum to cover StreamAE (1.3) and StatefulLHAE (1.2)
////////////////////////////////////////////////////////////////////////////////
#set-options "--admit_smt_queries true"
noeq type state (i:id) (rw:rw) =
  | Stream: u:unit{is_stream i} -> Stream.state i rw -> state i rw
  | StLHAE: u:unit{is_stlhae i} -> StLHAE.state i rw -> state i rw

let stream_state (#i:id{is_stream i}) (#rw:rw) (s:state i rw)
  : Tot (Stream.state i rw)
  = let Stream _ s = s in s

let stlhae_state (#i:id{is_stlhae i}) (#rw:rw) (s:state i rw)
  : Tot (StatefulLHAE.state i rw)
  = let StLHAE _ s = s in s

val region: #i:id -> #rw:rw -> state i rw -> Tot rgn
let region (#i:id) (#rw:rw) (s:state i rw): Tot rgn =
  match s with
  | Stream _ x -> Stream.State?.region x
  | StLHAE _ x -> StLHAE.region x

val log_region: #i:id -> #rw:rw -> state i rw -> Tot rgn
let log_region (#i:id) (#rw:rw) (s:state i rw): Tot rgn =
  match s with
  | Stream _ s -> Stream.State?.log_region s
  | StLHAE _ s -> AEAD_GCM.State?.log_region s

type reader i = state i Reader
type writer i = state i Writer

// For 0-RTT: we ignore a failed decryption for the
// handshake traffic key reader if the counter is 0
let tolerate_decrypt_failure (#i:id) (r:reader i)
  : ST bool
  (requires fun h0 -> True)
  (ensures fun h0 _ h1 -> modifies_none h0 h1)
  =
  match r with
  | StLHAE _ _ -> false
  | Stream _ st ->
    let ctr = HST.op_Bang (StreamAE.ctr st.StreamAE.counter) in
    let ID13 (KeyID #li (ExpandedSecret _ t _)) = i in
    0 = ctr && ClientHandshakeTrafficSecret? t

let tolerate_ccs (#i:id) (r:reader i)
  : ST bool
  (requires fun h0 -> True)
  (ensures fun h0 _ h1 -> modifies_none h0 h1)
  =
  match r with
  | StLHAE _ _ -> false
  | Stream _ st ->
    let ctr = ! (StreamAE.ctr st.StreamAE.counter) in
    let ID13 (KeyID #li (ExpandedSecret _ t _)) = i in
    0 = ctr && (ServerHandshakeTrafficSecret? t ||
      ClientHandshakeTrafficSecret? t || ClientEarlyTrafficSecret? t)

// our view to AE's ideal log (when idealized, ignoring ciphers) and counter
// TODO: write down their joint monotonic specification: both are monotonic,
// and seqn = length log when ideal

////////////////////////////////////////////////////////////////////////////////
//Logs of fragments, defined as projections on the underlying entry logs
////////////////////////////////////////////////////////////////////////////////
// TODO: consider adding constraint on terminator fragments
type frags (i:id{~ (PlaintextID? i)}) = Seq.seq (C.fragment i)

let ideal_log (r:rgn) (i:id) =
  if is_stream i then
    Stream.ideal_log r i
  else if is_stlhae i then
    AEAD_GCM.ideal_log r i
  else False

let ilog (#i:id) (#rw:rw) (s:state i rw{authId i}): Tot (ideal_log (log_region s) i) =
  match s with
  | Stream u s -> Stream.ilog (Stream.State?.log s)
  | StLHAE u s -> AEAD_GCM.ilog (AEAD_GCM.State?.log s)

let entry (i:id) =
   if is_stream i then
     Stream.entry i
   else if is_stlhae i then
     AEAD_GCM.entry i
   else False

private
let ptext (#i:id) (ent:entry i): Tot (C.fragment i) =
  if is_stream i then
    Stream.Entry?.p #i ent
  else
    AEAD_GCM.Entry?.p #i ent

//A projection of fragments from Stream.entries
let fragments (#i:id) (#rw:rw) (s:state i rw{authId i}) (h:mem): GTot (frags i) =
  let entries = HS.sel #_ #MS.grows h (ilog s) in
  MS.map ptext entries

val lemma_fragments_snoc_commutes: #i:id -> w:writer i{authId i}
  -> h0:mem -> h1:mem -> e:entry i
  -> Lemma (let log = ilog w in
           HS.sel #_ #MS.grows h1 log ==
	   Seq.snoc (HS.sel #_ #MS.grows h0 log) e ==>
	   fragments w h1 == Seq.snoc (fragments w h0) (ptext e))
let lemma_fragments_snoc_commutes #i w h0 h1 e =
  let log = ilog w in
  MS.map_snoc ptext (HS.sel #_ #MS.grows h0 log) e

//A predicate stating that the fragments have fs as a prefix
let fragments_prefix (#i:id) (#rw:rw) (w:state i rw{authId i}) (fs:frags i) (h:mem) : GTot Type0 =
  MS.map_prefix #_ #_ #(log_region w) (ilog w) ptext fs h

//In order to witness fragments_prefix s fs, we need to prove that it is stable
val fragments_prefix_stable: #i:id -> #rw:rw
  -> w:state i rw{authId i} -> h:mem
  -> Lemma (let fs = fragments w h in
	   MS.grows fs fs
	   /\ HST.stable_on_t #(log_region w) #_ #(MS.grows #(entry i)) (ilog w)
	     (fragments_prefix w fs))
let fragments_prefix_stable #i #rw w h =
  let fs = fragments w h in
  let log = ilog w in
  // MS.seq_extension_reflexive fs; //NS: seems no longer necessary
  MS.map_prefix_stable #_ #_ #(log_region w) log ptext fs


////////////////////////////////////////////////////////////////////////////////
//Projecting sequence numbers
////////////////////////////////////////////////////////////////////////////////

let seqnT (#i:id) (#rw:rw) (s:state i rw) h : GTot nat =
  match s with
  | Stream _ s -> HS.sel h (Stream.ctr (Stream.State?.counter s))
  | StLHAE _ s -> HS.sel h (AEAD_GCM.ctr (StLHAE.counter s))

//it's incrementable if it doesn't overflow
let incrementable (#i:id) (#rw:rw) (s:state i rw) (h:mem) =
  seqnT s h < 18446744073709551615

// Some invariants:
// - the writer counter is the length of the log; the reader counter is lower or equal
// - gen is called at most once for each (i:id), generating distinct refs for each (i:id)
// - the log is monotonic

// We generate first the writer, then the reader (possibly several of them)


////////////////////////////////////////////////////////////////////////////////
//Framing
////////////////////////////////////////////////////////////////////////////////
val frame_fragments : #i:id -> #rw:rw -> st:state i rw -> h0:mem -> h1:mem -> s:Set.set rid
  -> Lemma
    (requires modifies s h0 h1
	      /\ Map.contains (HS.get_hmap h0) (log_region st)
	      /\ not (Set.mem (log_region st) s))
    (ensures authId i ==> fragments st h0 == fragments st h1)
let frame_fragments #i #rw st h0 h1 s = ()

#set-options "--z3rlimit 100 --max_ifuel 1 --initial_ifuel 3 --max_fuel 3 --initial_fuel 3"
val frame_seqnT : #i:id -> #rw:rw -> st:state i rw -> h0:mem -> h1:mem -> s:Set.set rid
	       -> Lemma
    (requires modifies s h0 h1
    	  /\ Map.contains (HS.get_hmap h0) (region st)
	      /\ not (Set.mem (region st) s))
    (ensures seqnT st h0 = seqnT st h1)
let frame_seqnT #i #rw st h0 h1 s = ()

let trigger_frame (h:mem) = True

let frame_f (#a:Type) (f:mem -> GTot a) (h0:mem) (s:Set.set rid) =
  forall h1.{:pattern trigger_frame h1}
        trigger_frame h1
        /\ (HS.equal_on s (HS.get_hmap h0) (HS.get_hmap h1) ==> f h0 == f h1)

val frame_seqT_auto: i:id -> rw:rw -> s:state i rw -> h0:mem -> h1:mem ->
  Lemma (requires   HS.equal_on (Set.singleton (region s)) (HS.get_hmap h0)
  (HS.get_hmap h1)
		  /\ Map.contains (HS.get_hmap h0) (region s))
        (ensures seqnT s h0 = seqnT s h1)
	[SMTPat (seqnT s h0);
	 SMTPat (seqnT s h1)]
//	 SMTPatT (trigger_frame h1)]
let frame_seqT_auto i rw s h0 h1 = ()

val frame_fragments_auto: i:id{authId i} -> rw:rw -> s:state i rw -> h0:mem -> h1:mem ->
  Lemma (requires    HS.equal_on (Set.singleton (log_region s)) (HS.get_hmap h0)
  (HS.get_hmap h1)
		  /\ Map.contains (HS.get_hmap h0) (log_region s))
        (ensures fragments s h0 == fragments s h1)
	[SMTPat (fragments s h0);
	 SMTPat (fragments s h1)]
	 (* SMTPatT (trigger_frame h1)] *)
let frame_fragments_auto i rw s h0 h1 = ()


////////////////////////////////////////////////////////////////////////////////
//Experimenting with reads clauses: probably unnecessary
////////////////////////////////////////////////////////////////////////////////
let reads (s:Set.set rid) (a:Type) =
    f: (h:mem -> GTot a){forall h1 h2. (HS.equal_on s (HS.get_hmap h1)
    (HS.get_hmap h2) /\ Set.subset s (Map.domain (HS.get_hmap h1)))
				  ==> f h1 == f h2}

(*
val fragments' : #i:id -> #rw:rw -> s:state i rw{ authId i } -> Tot (reads (Set.singleton (log_region s)) (frags i))
let fragments' #i #rw s = fun h -> fragments #i #rw s h
*)

(*------------------------------------------------------------------*)
let genPost (#i:id) parent h0 (w:writer i) h1 =
  let r = region #i #Writer w in
  HS.modifies_transitively Set.empty h0 h1 /\
  HS.extends r parent /\
  HS.fresh_region r h0 h1 /\
  color r = color parent /\
  seqnT #i #Writer w h1 = 0 /\
  (authId i ==> fragments #i #Writer w h1 == Seq.empty) // we need to re-apply #i knowning authId

// Generate a fresh instance with index i in a fresh sub-region
val gen: parent:rgn -> i:stae_id -> ST (writer i)
  (requires (fun h0 -> True))
  (ensures (genPost parent))
/
prims.fst.checked MiTLS.TLSInfo.fst.checked MiTLS.TLSConstants.fsti.checked MiTLS.StreamPlain.fst.checked MiTLS.StreamAE.fst.checked MiTLS.StatefulPlain.fst.checked MiTLS.StatefulLHAE.fst.checked MiTLS.Range.fst.checked MiTLS.Mem.fst.checked MiTLS.LHAEPlain.fst.checked MiTLS.EverCrypt.fsti.checked MiTLS.Content.fst.checked MiTLS.AEADProvider.fsti.checked MiTLS.AEAD_GCM.fst.checked FStar.UInt32.fsti.checked FStar.Set.fsti.checked FStar.Seq.fst.checked FStar.Pervasives.fsti.checked FStar.Monotonic.Seq.fst.checked FStar.Map.fsti.checked FStar.HyperStack.ST.fsti.checked FStar.HyperStack.fst.checked FStar.Bytes.fsti.checked
MiTLS.StAE.fst
parent: MiTLS.Mem.rgn -> i: MiTLS.StAE.stae_id -> FStar.HyperStack.ST.ST (MiTLS.StAE.writer i)
MiTLS.Mem.rgn MiTLS.StAE.stae_id MiTLS.StAE.is_stream MiTLS.StAE.Stream MiTLS.TLSConstants.Writer MiTLS.StAE.writer MiTLS.StreamAE.state MiTLS.StreamAE.gen MiTLS.StreamAE.writer Prims.bool MiTLS.StAE.StLHAE MiTLS.StatefulLHAE.state MiTLS.StatefulLHAE.gen MiTLS.StatefulLHAE.writer
let gen parent i =
  
if is_stream i then Stream () (Stream.gen parent i) else StLHAE () (StLHAE.gen parent i)