val map_ST: #a: Type -> #b: Type -> (a -> ST0 b) -> list a -> ST0 (list b)
val map_ST: #a: Type -> #b: Type -> (a -> ST0 b) -> list a -> ST0 (list b)
let rec map_ST: (#a:Type) -> (#b:Type) -> (a -> ST0 b) -> list a -> ST0 (list b) =
  fun #_ #_ f x -> match x with
  | [] -> []
  | a::tl -> f a :: map_ST f tl
module MiTLS.Old.KeySchedule
open MiTLS

open FStar.Heap
open FStar.HyperStack
open FStar.HyperStack.ST
open FStar.Seq
open FStar.Bytes
open FStar.Error
//open FStar.Integers 

open MiTLS.TLSError
open MiTLS.TLSConstants
open MiTLS.Extensions
open MiTLS.TLSInfo
open MiTLS.Range
open MiTLS.StatefulLHAE
open MiTLS.HKDF
open MiTLS.PSK

module MDM = FStar.Monotonic.DependentMap
module HS = FStar.HyperStack
module ST = FStar.HyperStack.ST
module H = MiTLS.Hashing.Spec

module HMAC_UFCMA = MiTLS.Old.HMAC.UFCMA

type finishedId = HMAC_UFCMA.finishedId

(* A flag for runtime debugging of computed keys.
   The F* normalizer will erase debug prints at extraction
   when this flag is set to false *)
let discard (b:bool): ST unit (requires (fun _ -> True))
 (ensures (fun h0 _ h1 -> h0 == h1)) = ()
let print s = discard (IO.debug_print_string ("KS | "^s^"\n"))
unfold let dbg : string -> ST unit (requires (fun _ -> True))
  (ensures (fun h0 _ h1 -> h0 == h1)) =
  if DebugFlags.debug_KS then print else (fun _ -> ())

#set-options "--admit_smt_queries true"

let print_share (#g:CommonDH.group) (s:CommonDH.share g) : ST unit
  (requires (fun h0 -> True))
  (ensures (fun h0 _ h1 -> modifies_none h0 h1))
  =
  let kb = CommonDH.serialize_raw #g s in
  let kh = FStar.Bytes.hex_of_bytes kb in
  dbg ("Share: "^kh)

(********************************************
*    Resumption PSK is disabled for now     *
*********************************************

abstract type res_psk (i:rmsId) =
  b:bytes{exists i.{:pattern index b i} index b i <> 0z}

abstract type res_context (i:rmsId) =
  b:bytes{length b = CoreCrypto.H.hash_len (rmsId_hash i)}

private type res_psk_entry (i:rmsId) =
  (res_psk i) * (res_context i) * ctx:psk_context * leaked:(rref tls_tables_region bool)

let res_psk_injective (m:MDM.map' rmsId res_psk_entry) =
  forall i1 i2.{:pattern (MDM.sel m i1); (MDM.sel m i2)}
       i1 = i2 <==> (match MDM.sel m i1, MDM.sel m i2 with
                  | Some (psk1, _, _, _), Some (psk2, _, _, _) -> b2t (equalBytes psk1 psk2)
                  | _ -> True)

let res_psk_table : MDM.t tls_tables_region rmsId res_psk_entry res_psk_injective =
  MDM.alloc #TLSConstants.tls_tables_region #rmsId #res_psk_entry #res_psk_injective

let registered_res_psk (i:rmsId) (h:HH.t) =
  b2t (Some? (MDM.sel (HS.sel h res_psk_table) i))

let res_psk_context (i:rmsId{registered_res_psk i}) =
  let (_, _, c, _) = Some.v (MDM.sel res_psk_table i) in c

private let res_psk_value (i:rmsId{registered_res_psk i}) =
  let (psk, _, _, _) = Some.v (MDM.sel res_psk_table i) in psk

**)

// PSK (internal/external multiplex, abstract)
// Note that application PSK is externally defined but should
// be idealized together with KS
val psk (i:esId) : Type0

let read_psk (i:PSK.pskid)
  : ST (esId * pskInfo * PSK.app_psk i)
  (requires fun h -> True)
  (ensures fun h0 _ h1 -> modifies_none h0 h1)
  =
  let c = PSK.psk_info i in
  let id =
    if Some? c.ticket_nonce then
      let (| li, rmsid |) = Ticket.dummy_rmsid c.early_ae c.early_hash in
      ResumptionPSK #li rmsid
    else
      ApplicationPSK #(c.early_hash) #(c.early_ae) i
    in
  (id, c, PSK.psk_value i)

type binderId = i:pre_binderId{valid (I_BINDER i)}
type hsId = i:TLSInfo.pre_hsId{valid (I_HS i)}
type asId = i:pre_asId{valid (I_AS i)}

// Resumption context
let rec esId_rc : (esId -> St bytes) =
  function
  | NoPSK h -> H.zeroHash h

and hsId_rc : (hsId -> St bytes) = function
  | HSID_DHE (Salt i) _ _ _ -> secretId_rc i
  | HSID_PSK (Salt i) -> secretId_rc i

and asId_rc : (asId -> St bytes) = function
  | ASID (Salt i) -> secretId_rc i

and secretId_rc : (secretId -> St bytes) = function
  | EarlySecretID i -> esId_rc i
  | HandshakeSecretID i -> hsId_rc i
  | ApplicationSecretID i -> asId_rc i

// miTLS 0.9:
// ==========
// PRF (type pms) -> TLSInfo (type id) -> KEF (extract pms)
//                     \ StatefulLHAE (coerce id) /
// TODO rework old 1.2 types
type ms = bytes
type pms = bytes

// Early secret (abstract)
val es (i:esId) : Type0

// Handshake secret (abstract)
val hs (i:hsId) : Type0
type fink (i:finishedId) = HMAC_UFCMA.key (HMAC_UFCMA.HMAC_Finished i) (fun _ -> True)
let trivial (_: bytes) = True
type binderKey (i:binderId) = HMAC_UFCMA.key (HMAC_UFCMA.HMAC_Binder i) trivial

// TLS 1.3 master secret (abstract)
val ams (i:asId) : Type0

type rekeyId (li:logInfo) = i:expandId li{
  (let ExpandedSecret _ t _ = i in
    ApplicationTrafficSecret? t \/
    ClientApplicationTrafficSecret? t \/
    ServerApplicationTrafficSecret? t)}

val rekey_secrets (#li:_) (i:expandId li) : Type0

type raw_rekey_secrets = {
  rekey_aead: aeadAlg;
  rekey_hash: hash_alg;
  rekey_client: bytes;
  rekey_server: bytes;
}

// Leaked to HS for tickets
(*abstract*) type rms #li (i:rmsId li) = H.tag (rmsId_hash i)

type ems #li (i:exportId li) = H.tag (exportId_hash i)

// TODO this is superseeded by StAE.state i
// but I'm waiting for it to be tested to switch over
// TODO use the newer index types
type recordInstance =
  | StAEInstance: 
    #id:TLSInfo.id ->
    r: StAE.reader (peerId id) ->
    w: StAE.writer id ->
    pn: option bytes * option bytes ->
    recordInstance

(* 2 choices - I prefer the second:
     (1) replace recordInstance in this module with Epochs.epoch, but that requires dependence on more than just $id
   (2) redefine recordInstance as follows, and then import epoch_region_inv over here from Epochs:
type recordInstance (rgn:rid) (n:TLSInfo.random) =
| RI: #id:StAE.id -> r:StAE.reader (peerId id) -> w:StAE.writer id{epoch_region_inv' rgn r w /\ I.nonce_of_id id = n} -> recordInstance rgn n

In (2) we would define Epochs.epoch as:
type epoch (hs_rgn:rgn) (n:TLSInfo.random) =
  | Epoch: h:handshake ->
           r:recordInstance hs_rgn n ->
           epich hs_rgn n
*)

type exportKey = (li:logInfo & i:exportId li & ems i)

// Note from old miTLS (in TLSInfo.fst)
// type id = {
//  msId   : msId;            // the index of the master secret used for key derivation
//  kdfAlg : kdfAlg_t;          // the KDF algorithm used for key derivation
//  pv     : protocolVersion; // should be part of aeAlg
//  aeAlg  : aeAlg;           // the authenticated-encryption algorithms
//  csrConn: csRands;         // the client-server random of the connection
//  ext: negotiatedExtensions;// the extensions negotiated for the current session
//  writer : role             // the role of the writer
//  }

type ks_alpha12 = pv:protocolVersion * cs:cipherSuite * ems:bool
type ks_alpha13 = ae:aeadAlg * h:hash_alg

type ks_client_state =
| C_Init: cr:random -> ks_client_state
| C_12_Resume_CH: cr:random -> si:sessionInfo -> msId:TLSInfo.msId -> ms:ms -> ks_client_state
| C_12_Full_CH: cr:random -> ks_client_state
| C_12_wait_MS: csr:csRands -> alpha:ks_alpha12 -> id:TLSInfo.pmsId -> pms:pms -> ks_client_state
| C_12_has_MS: csr:csRands -> alpha:ks_alpha12 -> id:TLSInfo.msId -> ms:ms -> ks_client_state
| C_13_wait_SH: cr:random -> esl: list (i:esId{~(NoPSK? i)} & es i) ->
                gs:list (g:CommonDH.group & CommonDH.ikeyshare g) -> ks_client_state
| C_13_wait_SF: alpha:ks_alpha13 -> (i:finishedId & cfk:fink i) -> (i:finishedId & sfk:fink i) ->
                (i:asId & ams:ams i) -> ks_client_state
| C_13_wait_CF: alpha:ks_alpha13 -> (i:finishedId & cfk:fink i) -> (i:asId & ams:ams i) ->
                (li:logInfo & i:rekeyId li & rekey_secrets i) -> ks_client_state
| C_13_postHS: alpha:ks_alpha13 -> (li:logInfo & i:rekeyId li & rekey_secrets i) ->
                (li:logInfo & i:rmsId li & rms i) -> ks_client_state
| C_Done

type ks_server_state =
| S_Init: sr:random -> ks_server_state
| S_12_wait_CKE_DH: csr:csRands -> alpha:ks_alpha12 -> our_share:(g:CommonDH.group & CommonDH.ikeyshare g) -> ks_server_state
| S_12_wait_CKE_RSA: csr: csRands -> alpha:ks_alpha12 -> ks_server_state
| S_12_has_MS: csr:csRands -> alpha:ks_alpha12 -> id:TLSInfo.msId -> ms:ms -> ks_server_state
| S_13_wait_SH: alpha:ks_alpha13 -> cr:random -> sr:random -> es:(i:esId & es i) ->
                hs:( i:hsId & hs i ) -> ks_server_state
| S_13_wait_SF: alpha:ks_alpha13 -> ( i:finishedId & cfk:fink i ) -> ( i:finishedId & sfk:fink i ) ->
                ( i:asId & ams:ams i ) -> ks_server_state
| S_13_wait_CF: alpha:ks_alpha13 -> ( i:finishedId & cfk:fink i ) -> ( i:asId & ams i ) ->
                ( li:logInfo & i:rekeyId li & rekey_secrets i ) ->  ks_server_state
| S_13_postHS: alpha:ks_alpha13 -> ( li:logInfo & i:rekeyId li & rekey_secrets i ) ->
	        ( li:logInfo & i:rmsId li & rms i) -> ks_server_state
| S_Done

// Reflecting state separation from HS
type ks_state =
| C: s:ks_client_state -> ks_state
| S: s:ks_server_state -> ks_state

// KeySchedule instances
(*
 * AR: changing state from rref to ref, with region captured in the refinement.
 *)
type ks =
| KS:
  #region:rid ->
  state:(ref ks_state){HS.frameOf state = region} ->
  is_quic: bool ->
  ks

//17-04-17 CF: expose it as a concrete ref?
//17-04-17 CF: no need to keep the region, already in the ref.

// Extract keys and IVs from a derived 1.3 secret
private let keygen_13 h secret ae is_quic : St (bytes * bytes * option bytes) =
  let kS = EverCrypt.aead_keyLen ae in
  let iS = 12ul in // IV length
  let lk, liv = if is_quic then "quic key", "quic iv" else "key", "iv" in
  let kb = HKDF.expand_label #h secret lk empty_bytes kS in
  let ib = HKDF.expand_label #h secret liv empty_bytes iS in
  let pn = if is_quic then
      Some (HKDF.expand_label #h secret "quic hp" empty_bytes kS)
    else None in
  (kb, ib, pn)

// Extract finished keys
private let finished_13 h secret : St (bytes) =
  HKDF.expand_label #h secret "finished" empty_bytes (Hacl.Hash.Definitions.hash_len h)

// Create a fresh key schedule instance
// We expect this to be called when the Handshake instance is created
let create: #rid:rid -> role -> is_quic:bool -> ST (ks * random)
  (requires fun h -> rid<>root)
  (ensures fun h0 (r,_) h1 ->
    let KS #ks_region state iq = r in
    HS.fresh_region ks_region h0 h1
    /\ extends ks_region rid
    /\ iq = is_quic
    /\ modifies (Set.singleton rid) h0 h1
    /\ HS.modifies_ref rid (Set.singleton (Heap.addr_of (as_ref state))) ( h0) ( h1)) =
  fun #rid r is_quic ->
  ST.recall_region rid;
  let ks_region = new_region rid in
  let nonce = Nonce.mkHelloRandom r ks_region in
  let istate = match r with
    | Client -> C (C_Init nonce)
    | Server -> S (S_Init nonce) in
  (KS #ks_region (ralloc ks_region istate) is_quic), nonce

private let group_of_valid_namedGroup
  (g:CommonDH.supportedNamedGroup)
  : CommonDH.group
  = Some?.v (CommonDH.group_of_namedGroup g)

effect ST0 (a:Type) = ST a (fun _ -> True) (fun h0 _ h1 -> modifies_none h0 h1)
/
prims.fst.checked MiTLS.TLSPRF.fst.checked MiTLS.TLSInfo.fst.checked MiTLS.TLSError.fst.checked MiTLS.TLSConstants.fsti.checked MiTLS.Ticket.fst.checked MiTLS.StatefulLHAE.fst.checked MiTLS.StAE.fst.checked MiTLS.RSAKey.fst.checked MiTLS.Range.fst.checked MiTLS.PSK.fst.checked MiTLS.Old.HMAC.UFCMA.fsti.checked MiTLS.Nonce.fsti.checked MiTLS.HKDF.fsti.checked MiTLS.Hashing.Spec.fst.checked MiTLS.Extensions.fsti.checked MiTLS.EverCrypt.fsti.checked MiTLS.DebugFlags.fst.checked MiTLS.CommonDH.fsti.checked Hacl.Hash.Definitions.fst.checked FStar.UInt32.fsti.checked FStar.Set.fsti.checked FStar.Seq.fst.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.Monotonic.DependentMap.fsti.checked FStar.List.Tot.fst.checked FStar.IO.fst.checked FStar.HyperStack.ST.fsti.checked FStar.HyperStack.fst.checked FStar.Heap.fst.checked FStar.Error.fst.checked FStar.Bytes.fsti.checked
MiTLS.Old.KeySchedule.fsti
f: (_: a -> MiTLS.Old.KeySchedule.ST0 b) -> x: Prims.list a
  -> MiTLS.Old.KeySchedule.ST0 (Prims.list b)
Prims.list Prims.Nil Prims.Cons MiTLS.Old.KeySchedule.map_ST
let rec map_ST: #a: Type -> #b: Type -> (a -> ST0 b) -> list a -> ST0 (list b) =
  
fun #_ #_ f x ->
  match x with
  | [] -> []
  | a :: tl -> f a :: map_ST f tl