let max (a b:nat) =
//  let open FStar.UInt64 in
  if a > b then a else b
let max (a b: nat) =
if a > b then a else b
Model.QUIC.fsti.checked
Spec.Agile.Cipher.fsti.checked Spec.Agile.AEAD.fsti.checked QUIC.UInt62.fst.checked QUIC.Spec.PacketNumber.Base.fst.checked QUIC.Spec.Header.Parse.fsti.checked QUIC.Spec.Header.Base.fst.checked QUIC.Spec.Header.fsti.checked QUIC.Spec.fsti.checked QUIC.Secret.Int.fsti.checked prims.fst.checked Model.PNE.fsti.checked Model.Indexing.fsti.checked Model.Helpers.fsti.checked Model.AEAD.fsti.checked Mem.fst.checked LowStar.Modifies.fst.checked LowParse.BitFields.fsti.checked Lib.RawIntTypes.fsti.checked FStar.UInt8.fsti.checked FStar.UInt64.fsti.checked FStar.UInt32.fsti.checked FStar.UInt128.fsti.checked FStar.Seq.fst.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.HyperStack.fst.checked
Model.QUIC.fsti
Prims.nat Prims.op_GreaterThan Prims.bool
module Model.QUIC

module HS = FStar.HyperStack
module I = Model.Indexing
module U32 = FStar.UInt32
module U64 = FStar.UInt64
module U128 = FStar.UInt128
module M = LowStar.Modifies

module Spec = QUIC.Spec
module QH = QUIC.Spec.Header
module AE = Model.AEAD
module SAE = Spec.Agile.AEAD
module PNE = Model.PNE
module SPNE = Spec.Agile.Cipher
module BF = LowParse.BitFields
module U62 = QUIC.UInt62
module Secret = QUIC.Secret.Int

open FStar.UInt32
open Mem

type id = i:AE.id & j:PNE.id{PNE.is_safe j <==> AEAD.is_safe i}

let safe (i:id) = AEAD.is_safe (dfst i)
let safe_id = i:id{safe i}
let unsafe (i:id) = not (safe i)
let unsafe_id = i:id{unsafe i}

/// Package for QUIC plaintexts
/// (i.e. the payload of QUIC packets)
/// This would be implemented as a serializable list of frames
/// (using EverParse refined lists)

type qplain_len = n:AEAD.plain_length_at_least 3{n < Spec.max_plain_length}
let pnl = PNE.pne_plain_length

val pne_pkg: PNE.pne_plain_pkg

noeq type info = {
  region: r:subq{r `HS.disjoint` q_ae_region};
}

let max_ctr = pow2 62 - 1
type epn (nl:pnl) = Spec.lbytes nl
type pn = n:nat{n <= max_ctr}
type rpn = n:U64.t{U64.v n < max_ctr}
let rpn_of_nat (j:nat{j < max_ctr}) : rpn = U64.uint_to_t j

val stream_writer: (k:id) -> Type u#1
val stream_reader: #k:id -> w:stream_writer k -> Type u#1
val writer_info: #k:id -> w:stream_writer k -> info
val reader_info: #k:id -> #w:stream_writer k -> r:stream_reader w -> i:info{i == writer_info w}

val writer_ae_info: #k:id -> w:stream_writer k -> a:AEAD.info (dfst k){a.AEAD.min_len == 3}
val reader_ae_info: #k:id -> #w:stream_writer k -> r:stream_reader w -> a:AEAD.info (dfst k){a.AEAD.min_len == 3}
val writer_pne_info: #k:id -> w:stream_writer k -> a:PNE.info (dsnd k){a.PNE.calg == Spec.Agile.AEAD.cipher_alg_of_supported_alg (writer_ae_info w).AEAD.alg /\ a.PNE.halg == (writer_ae_info w).AEAD.halg}
val reader_pne_info: #k:id -> #w:stream_writer k -> r:stream_reader w -> a:PNE.info (dsnd k){a.PNE.calg == Spec.Agile.AEAD.cipher_alg_of_supported_alg (reader_ae_info r).AEAD.alg /\ a.PNE.halg == (reader_ae_info r).AEAD.halg}

val writer_aead_state : (#k:id) -> (w:stream_writer k) ->
  aw:AEAD.aead_writer (dfst k)
val reader_aead_state : #k:id -> #w:stream_writer k -> r:stream_reader w ->
  ar:AEAD.aead_reader (writer_aead_state w)
val writer_pne_state : #k:id -> w:stream_writer k -> PNE.pne_state (writer_pne_info w)
val reader_pne_state : #k:id -> #w:stream_writer k -> r:stream_reader w -> PNE.pne_state (reader_pne_info r)

type key_t #i (w:stream_writer i) =
  Spec.lbytes (SAE.key_length (writer_ae_info w).AEAD.alg) &
  Spec.lbytes (PNE.key_len (writer_pne_info w))

let writer_leak (#k:unsafe_id) (w:stream_writer k) : key_t w =
  Model.Helpers.reveal #(SAE.key_length (writer_ae_info w).AEAD.alg) (AEAD.wkey (writer_aead_state w)),
  Model.Helpers.reveal #(PNE.key_len (writer_pne_info w)) (PNE.key (writer_pne_state w))

let reader_leak (#k:unsafe_id) (#w:stream_writer k) (r:stream_reader w) : key_t w = writer_leak w

val invariant: #k:id -> w:stream_writer k -> h:mem ->
  t:Type0{t ==> AEAD.winvariant (writer_aead_state w) h}
val rinvariant: #k:id -> #w:stream_writer k -> r:stream_reader w -> h:mem ->
  t:Type0{t ==> invariant w h}

val writer_offset: #k:id -> w:stream_writer k -> pn
val reader_offset: #k:id -> #w:stream_writer k -> stream_reader w -> pn

val wctrT: #k:id -> w:stream_writer k -> h:mem{invariant w h} -> 
  GTot (n:nat{n >= writer_offset w /\ n <= max_ctr})
  
val wctr: #k:id -> w:stream_writer k -> ST (n:nat{n >= writer_offset w /\ n <= max_ctr})
  (requires fun h0 -> invariant w h0)
  (ensures fun h0 c h1 -> h0 == h1 /\ c = wctrT w h1)

type qiv (k:id) = Spec.lbytes 12 // SAE.iv (I.ae_id_ginfo (fst k))
val writer_static_iv: #k:id -> w:stream_writer k -> qiv k
val reader_static_iv: #k:id -> #w:stream_writer k -> r:stream_reader w ->
  iv:qiv k{iv == writer_static_iv w}

val expected_pnT: #k:id -> #w:stream_writer k -> r:stream_reader w -> h:mem{rinvariant r h} ->
  GTot (n:nat{n >= writer_offset w /\ n < max_ctr})
  
val expected_pn: #k:id -> #w:stream_writer k -> r:stream_reader w ->
  ST (n:nat{n >= writer_offset w /\ n <= max_ctr})
  (requires fun h0 -> rinvariant r h0)
  (ensures fun h0 c h1 -> h0 == h1 /\
    (c == expected_pnT #k #w r h0))

let wincrementable (#k:id) (w:stream_writer k) (h:mem{invariant w h}) =
  wctrT w h < max_ctr

val footprint: #k:id -> w:stream_writer k -> GTot M.loc
val rfootprint: #k:id -> #w:stream_writer k -> r:stream_reader w -> GTot M.loc

val frame_invariant: #k:id -> w:stream_writer k -> h0:mem -> ri:M.loc -> h1:mem ->
  Lemma
  (requires
    (invariant w h0 /\
    M.modifies ri h0 h1 /\
    M.loc_disjoint ri (footprint w)))
  (ensures invariant w h1 /\
    wctrT w h0 == wctrT w h1)
    //AEAD.wlog (writer_aead_state w) h1 == AEAD.wlog (writer_aead_state w) h0 /\
    //PNE.table (writer_pne_state w) h1 == PNE.table (writer_pne_state w) h0)
  [ SMTPat (M.modifies ri h0 h1); SMTPat (invariant w h1) ]
  (*[ SMTPatOr [
      [ SMTPat (M.modifies ri h0 h1); SMTPat (invariant w h1) ];
      [ SMTPat (M.modifies ri h0 h1); SMTPat (AEAD.wlog (writer_aead_state w) h1) ];
      [ SMTPat (M.modifies ri h0 h1); SMTPat (PNE.table (writer_pne_state w) h1) ]
  ]]*)

val rframe_invariant: #k:id -> #w:stream_writer k -> r:stream_reader w ->
  h0:mem -> ri:M.loc -> h1:mem ->
  Lemma
  (requires
    (rinvariant r h0 /\
    M.modifies ri h0 h1 /\
    M.loc_disjoint ri (rfootprint r)))
  (ensures rinvariant r h1 /\
    expected_pnT r h0 == expected_pnT r h1)
  [ SMTPat (M.modifies ri h0 h1); SMTPat (rinvariant r h1) ]

val wframe_log: #k:id{AEAD.is_safe (dfst k)} -> w:stream_writer k -> l:Seq.seq (AEAD.entry (dfst k) (AEAD.wgetinfo (writer_aead_state w))) ->
  h0:mem -> ri:M.loc -> h1:mem ->
  Lemma
  (requires
    invariant w h0 /\
    AEAD.wlog (writer_aead_state w) h0 == l /\
    M.modifies ri h0 h1 /\
    M.loc_disjoint ri (footprint w))
  (ensures invariant w h1 ==> AEAD.wlog (writer_aead_state w) h1 == l)

val rframe_log: #k:id{AEAD.is_safe (dfst k)} -> #w:stream_writer k -> r:stream_reader w -> l:Seq.seq (AEAD.entry (dfst k) (AEAD.rgetinfo (reader_aead_state r))) ->
  h0:mem -> ri:M.loc -> h1:mem ->
  Lemma
  (requires
    invariant w h0 /\
    AEAD.rlog (reader_aead_state r) h0 == l /\
    M.modifies ri h0 h1 /\
    M.loc_disjoint ri (rfootprint r))
  (ensures invariant w h1 ==> AEAD.rlog (reader_aead_state r) h1 == l)

val wframe_pnlog: #k:id{PNE.is_safe (dsnd k)} -> w:stream_writer k -> l:Seq.seq (PNE.entry (writer_pne_info w)) ->
  h0:mem -> ri:M.loc -> h1:mem ->
  Lemma
  (requires
    invariant w h0 /\
    PNE.table (writer_pne_state w) h0 == l /\
    M.modifies ri h0 h1 /\
    M.loc_disjoint ri (footprint w))
  (ensures PNE.table (writer_pne_state w) h1 == l)

val rframe_pnlog: #k:id{PNE.is_safe (dsnd k)} ->  #w:stream_writer k -> r:stream_reader w -> l:Seq.seq (PNE.entry (reader_pne_info r)) ->
  h0:mem -> ri:M.loc -> h1:mem ->
  Lemma
  (requires
    rinvariant r h0 /\
    PNE.table (reader_pne_state r) h0 == l /\
    M.modifies ri h0 h1 /\
    M.loc_disjoint ri (rfootprint r))
  (ensures PNE.table (reader_pne_state r) h1 == l)

val create: k:id -> u:info ->
  u1:AEAD.info (dfst k) -> u2:PNE.info (dsnd k) -> init: pn ->
  ST (stream_writer k)
  (requires fun h0 -> u2.PNE.calg ==
    Spec.Agile.AEAD.cipher_alg_of_supported_alg u1.AEAD.alg /\
    u2.PNE.halg == u1.AEAD.halg /\
    u2.PNE.plain == pne_pkg /\
    u1.AEAD.min_len == 3)
  (ensures fun h0 w h1 ->
    invariant w h1 /\
    M.modifies (footprint w) h0 h1 /\
    writer_offset w == init /\
    writer_ae_info w == u1 /\
    writer_pne_info w == u2 /\
    writer_info w == u /\
    wctrT w h1 == writer_offset w /\
    (safe k ==>
      (AEAD.wlog (writer_aead_state w) h1 == Seq.empty /\
      PNE.table (writer_pne_state w) h1 == Seq.empty
    ))
  )

#push-options "--z3rlimit 15"

val coerce: k:unsafe_id -> u:info ->
  u1:AEAD.info (dfst k) -> u2:PNE.info (dsnd k) ->  
  init: pn -> ts:AEAD.traffic_secret u1.AEAD.halg ->
  ST (stream_writer k)
  (requires fun h0 -> u2.PNE.calg ==
    Spec.Agile.AEAD.cipher_alg_of_supported_alg u1.AEAD.alg /\
    u2.PNE.halg == u1.AEAD.halg /\
    u2.PNE.plain == pne_pkg /\
    u1.AEAD.min_len == 3)
  (ensures fun h0 w h1 ->
    let (k1, k2) = writer_leak w in
    invariant w h1 /\
    M.modifies (footprint w) h0 h1 /\
    writer_offset w == init /\
    wctrT w h1 == writer_offset w /\
    writer_ae_info w == u1 /\
    writer_pne_info w == u2 /\
    writer_info w == u /\
    Model.Helpers.hide (writer_static_iv w) ==
      Spec.derive_secret u1.AEAD.halg ts
        Spec.label_iv 12 /\
    Model.Helpers.hide k1 == Spec.derive_secret u1.AEAD.halg ts
        Spec.label_key (SAE.key_length u1.AEAD.alg) /\
    Model.Helpers.hide k2 == QUIC.Spec.derive_secret u2.PNE.halg ts
        QUIC.Spec.label_hp (PNE.key_len u2)
  )
#pop-options

val createReader: parent:rgn -> #k:id -> w:stream_writer k ->
  ST (stream_reader w)
  (requires fun h0 -> invariant w h0 /\
  writer_offset w < max_ctr)
  (ensures fun h0 r h1 ->
    invariant w h1 /\ rinvariant r h1 /\
    M.modifies (rfootprint r) h0 h1 /\
    expected_pnT r h1 == writer_offset w)

#reset-options "--z3rlimit 50 --fuel 1"

let _ = assert_norm(Spec.max_plain_length < pow2 32)
let _ = assert_norm(pow2 32 < pow2 64)

(*
      let (ne,c) = split #k #(nl+l) nec nl in
      let rpn = rpn_of_nat (wctrT w h0) in
      let npn = npn_encode j rpn nl in
      let alg = ((AEAD.wgetinfo aw).AEAD.alg) in
      let nce:AEAD.nonce (AEAD.wgetinfo aw) = create_nonce #k #alg (writer_iv w) rpn in
      let ad:AEAD.adata = Bytes.append (bytes_of_quic_header hd) npn in
      let s:PNE.sample = sample_quic_protect nec in
      let nn = pne_plain_of_header_pn hd npn in
      let cc = pne_cipher_of_pheader_epn ph ne in
      AEAD.wlog aw h1 ==
        Seq.snoc
          (AEAD.wlog aw h0)
          (AEAD.Entry #i #(AEAD.wgetinfo aw) nce ad #l p c) /\
           PNE.table ps h1 ==
              Seq.snoc
                (PNE.table ps h0)
                (PNE.Entry #j #pne_plain_pkg s #(nl+1) nn cc))
*)

module PN = QUIC.Spec.PacketNumber.Base
module U62 = QUIC.UInt62
module Secret = QUIC.Secret.Int

let set_pn_long (l:Spec.long_header_specifics{not (Spec.MRetry? l)}) (pn:PN.packet_number_t) =
  let open Spec in
  match l with
  | MInitial r t p pnl _ -> MInitial r t p pnl pn
  | MZeroRTT r p pnl _ -> MZeroRTT r p pnl pn
  | MHandshake r p pnl _ -> MHandshake r p pnl pn

#push-options "--z3rlimit 200"
let set_pn (h:Spec.header{not (Spec.is_retry h)}) (pn:nat{pn <= max_ctr}) =
  let pn : PN.packet_number_t = Secret.hide (U62.uint_to_t pn) in
  match h with 
  | Spec.MLong b d s l -> Spec.MLong b d s (set_pn_long l pn)
  | Spec.MShort r s k d pnl _ -> Spec.MShort r s k d pnl pn
#pop-options

val encrypt
  (#k:id)
  (w:stream_writer k)
  (h:Spec.header)
  (#l:qplain_len)
  (p:(writer_ae_info w).AEAD.plain_pkg.AEAD.plain (dfst k) l)
  : ST Spec.packet
  (requires fun h0 ->
    invariant w h0 /\
    wincrementable w h0 /\
    (if Spec.is_retry h then l = 0
    else (
      (Lib.RawIntTypes.u64_from_UInt64 (UInt64.uint_to_t (wctrT w h0 + 1))) == QUIC.Spec.Header.Base.packet_number h /\ (
      Spec.has_payload_length h ==>
        Secret.v (Spec.payload_length h) == l
	  + Spec.Agile.AEAD.tag_length (writer_ae_info w).AEAD.alg))
  ))
  (ensures fun h0 c h1 ->
    let (|i,j|) = k in
    let aw = writer_aead_state w in
    let ps = writer_pne_state w in
    M.modifies (footprint w) h0 h1 /\
    invariant w h1 /\
    wctrT w h1 == wctrT w h0 + 1 /\
    (safe k ==> True) /\
    (unsafe k ==>
      (let ea = (writer_ae_info w).AE.alg in
      let k1, k2 = writer_leak w in
      let plain_pkg = (writer_ae_info w).AEAD.plain_pkg in
      let plain = plain_pkg.AEAD.repr i l p in
      c == Spec.encrypt ea (Helpers.hide k1) (Helpers.hide (writer_static_iv w)) (Helpers.hide k2) h (Helpers.reveal #l plain))
    ))

#push-options "--fuel 1 --z3rlimit 30"
noeq type model_result (#k:id) (#w:stream_writer k) (r:stream_reader w) =
| M_Success:
  h: Spec.header{not (Spec.is_retry h)} ->
  l: qplain_len ->
  p: (reader_ae_info r).AEAD.plain_pkg.AEAD.plain (dfst k) l ->
  remainder: Spec.bytes ->
  model_result r
| M_Failure

let max (a b:nat) =
val max : a: Prims.nat -> b: Prims.nat -> Prims.nat
a: Prims.nat -> b: Prims.nat -> Prims.nat