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
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
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
QUIC.Spec.Header.Base.long_header_specifics Prims.b2t Prims.op_Negation QUIC.Spec.Header.Base.uu___is_MRetry QUIC.Spec.PacketNumber.Base.packet_number_t QUIC.Spec.Base.bitfield QUIC.Spec.Base.vlbytes QUIC.Spec.Base.token_max_len QUIC.Spec.Base.payload_and_pn_length_t QUIC.Spec.PacketNumber.Base.packet_number_length_t QUIC.Spec.Header.Base.MInitial QUIC.Spec.Header.Base.MZeroRTT QUIC.Spec.Header.Base.MHandshake
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
val set_pn_long : l: QUIC.Spec.Header.Base.long_header_specifics{Prims.op_Negation (MRetry? l)} ->
    pn: QUIC.Spec.PacketNumber.Base.packet_number_t
  -> QUIC.Spec.Header.Base.long_header_specifics

    l: QUIC.Spec.Header.Base.long_header_specifics{Prims.op_Negation (MRetry? l)} ->
    pn: QUIC.Spec.PacketNumber.Base.packet_number_t
  -> QUIC.Spec.Header.Base.long_header_specifics