val rsapss_sign_:
    a:Hash.hash_alg{hash_is_supported a}
  -> modBits:modBits_t
  -> skey:rsapss_skey modBits
  -> sLen:size_nat{
    sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a /\
    sLen + Hash.hash_length a + 2 <= blocks (modBits - 1) 8}
  -> salt:lbytes sLen
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen} ->
  tuple2 bool (lbytes (blocks modBits 8))
let rsapss_sign_ a modBits skey sLen salt msgLen msg =
  let pkey = Mk_rsapss_skey?.pkey skey in
  let n = Mk_rsapss_pkey?.n pkey in
  let e = Mk_rsapss_pkey?.e pkey in
  let d = Mk_rsapss_skey?.d skey in

  let k = blocks modBits 8 in
  FStar.Math.Lemmas.pow2_le_compat (8 * k) modBits;

  let emBits = modBits - 1 in
  let emLen = blocks emBits 8 in

  let em = pss_encode a sLen salt msgLen msg emBits in
  let m = os2ip #emLen em in
  os2ip_lemma emBits em;
  let s = pow_mod #n m d in
  let m' = pow_mod #n s e in
  let eq_m = m = m' in
  let s = if eq_m then s else 0 in
  (eq_m, i2osp k s)
val rsapss_sign_:
    a:Hash.hash_alg{hash_is_supported a}
  -> modBits:modBits_t
  -> skey:rsapss_skey modBits
  -> sLen:size_nat{
    sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a /\
    sLen + Hash.hash_length a + 2 <= blocks (modBits - 1) 8}
  -> salt:lbytes sLen
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen} ->
  tuple2 bool (lbytes (blocks modBits 8))
let rsapss_sign_ a modBits skey sLen salt msgLen msg =
let pkey = Mk_rsapss_skey?.pkey skey in
let n = Mk_rsapss_pkey?.n pkey in
let e = Mk_rsapss_pkey?.e pkey in
let d = Mk_rsapss_skey?.d skey in
let k = blocks modBits 8 in
FStar.Math.Lemmas.pow2_le_compat (8 * k) modBits;
let emBits = modBits - 1 in
let emLen = blocks emBits 8 in
let em = pss_encode a sLen salt msgLen msg emBits in
let m = os2ip #emLen em in
os2ip_lemma emBits em;
let s = pow_mod #n m d in
let m' = pow_mod #n s e in
let eq_m = m = m' in
let s = if eq_m then s else 0 in
(eq_m, i2osp k s)
Spec.RSAPSS.fst.checked
Spec.Agile.Hash.fsti.checked prims.fst.checked Lib.Sequence.fsti.checked Lib.RawIntTypes.fsti.checked Lib.NatMod.fsti.checked Lib.LoopCombinators.fsti.checked Lib.IntTypes.fsti.checked Lib.ByteSequence.fsti.checked FStar.UInt8.fsti.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.Mul.fst.checked FStar.Math.Lemmas.fst.checked FStar.Calc.fsti.checked
Spec.RSAPSS.fst
Spec.Hash.Definitions.hash_alg Prims.b2t Spec.RSAPSS.hash_is_supported Spec.RSAPSS.modBits_t Spec.RSAPSS.rsapss_skey Lib.IntTypes.size_nat Prims.l_and Prims.op_LessThanOrEqual Prims.op_Addition Spec.Hash.Definitions.hash_length Lib.IntTypes.max_size_t Spec.Hash.Definitions.less_than_max_input_length Spec.RSAPSS.blocks Prims.op_Subtraction Lib.ByteSequence.lbytes Prims.nat Lib.ByteSequence.bytes Prims.eq2 Lib.Sequence.length Lib.IntTypes.uint_t Lib.IntTypes.U8 Lib.IntTypes.SEC FStar.Pervasives.Native.Mktuple2 Prims.bool Spec.RSAPSS.i2osp Prims.op_LessThan Prims.pow2 Prims.op_Multiply FStar.Mul.op_Star Prims.op_Equality Prims.l_or Lib.NatMod.nat_mod Lib.NatMod.pow_mod Prims.unit Spec.RSAPSS.os2ip_lemma Spec.RSAPSS.os2ip Lib.Sequence.lseq Lib.IntTypes.int_t Spec.RSAPSS.pss_encode Prims.pos Prims.int FStar.Math.Lemmas.pow2_le_compat Spec.RSAPSS.__proj__Mk_rsapss_skey__item__d Spec.RSAPSS.__proj__Mk_rsapss_pkey__item__e Spec.RSAPSS.__proj__Mk_rsapss_pkey__item__n Spec.RSAPSS.rsapss_pkey Spec.RSAPSS.__proj__Mk_rsapss_skey__item__pkey FStar.Pervasives.Native.tuple2
module Spec.RSAPSS

open FStar.Mul
open Lib.IntTypes

open Lib.NatMod
open Lib.Sequence
open Lib.ByteSequence
open Lib.LoopCombinators

module Hash = Spec.Agile.Hash

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

(**
 RFC spec of RSA-PSS: https://tools.ietf.org/html/rfc8017#section-8.1
 List of supported hash functions: https://tools.ietf.org/html/rfc8446#appendix-B.3.1.3
 Mask generation function: https://tools.ietf.org/html/rfc8017#page-67
*)

///
/// Auxillary functions
///

val blocks: x:size_pos -> m:size_pos -> Tot (r:size_pos{x <= m * r})
let blocks x m = (x - 1) / m + 1

val xor_bytes: #len:size_pos -> b1:lbytes len -> b2:lbytes len -> Tot (lbytes len)
let xor_bytes #len b1 b2 = map2 (fun x y -> x ^. y) b1 b2

let hash_is_supported (a:Hash.hash_alg) : Tot bool =
  match a with
  | Hash.SHA2_256 -> true
  | Hash.SHA2_384 -> true
  | Hash.SHA2_512 -> true
  | _ -> false

(* Mask Generation Function *)
val mgf_hash_f:
    a:Hash.hash_alg{hash_is_supported a}
  -> len:size_nat{len + 4 <= max_size_t /\ (len + 4) `Hash.less_than_max_input_length` a}
  -> i:size_nat
  -> mgfseed_counter:lbytes (len + 4) ->
  lbytes (len + 4) & lbytes (Hash.hash_length a)

let mgf_hash_f a len i mgfseed_counter =
  let counter = nat_to_intseq_be 4 i in
  let mgfseed_counter = update_sub mgfseed_counter len 4 counter in
  let block = Hash.hash a mgfseed_counter in
  mgfseed_counter, block

let mgf_hash_a (len:size_nat{len + 4 <= max_size_t}) (n:pos) (i:nat{i <= n}) = lbytes (len + 4)

val mgf_hash:
    a:Hash.hash_alg{hash_is_supported a}
  -> len:size_nat{len + 4 <= max_size_t /\ (len + 4) `Hash.less_than_max_input_length` a}
  -> mgfseed:lbytes len
  -> maskLen:size_pos{(blocks maskLen (Hash.hash_length a)) * Hash.hash_length a < pow2 32} ->
  Tot (lbytes maskLen)

let mgf_hash a len mgfseed maskLen =
  let mgfseed_counter = create (len + 4) (u8 0) in
  let mgfseed_counter = update_sub mgfseed_counter 0 len mgfseed in

  let hLen = Hash.hash_length a in
  let n = blocks maskLen hLen in
  let _, acc = generate_blocks #uint8 hLen n n (mgf_hash_a len n) (mgf_hash_f a len) mgfseed_counter in
  sub #uint8 #(n * hLen) acc 0 maskLen


(* Bignum convert functions *)
val os2ip: #len:size_nat -> b:lbytes len -> Tot (res:nat{res < pow2 (8 * len)})
let os2ip #len b = nat_from_bytes_be b

val i2osp: len:size_nat -> n:nat{n < pow2 (8 * len)} -> Tot (lbytes len)
let i2osp len n = nat_to_intseq_be len n

///
///  RSA
///

type modBits_t = modBits:size_nat{1 < modBits}

noeq type rsapss_pkey (modBits:modBits_t) =
  | Mk_rsapss_pkey: n:pos{pow2 (modBits - 1) < n /\ n < pow2 modBits} -> e:pos -> rsapss_pkey modBits

noeq type rsapss_skey (modBits:modBits_t) =
  | Mk_rsapss_skey: pkey:rsapss_pkey modBits -> d:pos -> rsapss_skey modBits


val db_zero: #len:size_pos -> db:lbytes len -> emBits:size_nat ->
  Pure (lbytes len)
  (requires True)
  (ensures fun res ->
    (v res.[0] == (if emBits % 8 > 0 then v db.[0] % pow2 (emBits % 8) else v db.[0])))

let db_zero #len db emBits =
  let msBits = emBits % 8 in
  if msBits > 0 then begin
    let r = db.[0] <- db.[0] &. (u8 0xff >>. size (8 - msBits)) in
    Math.Lemmas.pow2_plus msBits (8 - msBits);
    logand_mask db.[0] (u8 0xff >>. size (8 - msBits)) msBits;
    r end
  else db


val pss_encode:
    a:Hash.hash_alg{hash_is_supported a}
  -> sLen:size_nat{sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a}
  -> salt:lbytes sLen
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen}
  -> emBits:size_pos{Hash.hash_length a + sLen + 2 <= blocks emBits 8} ->
  Pure (lbytes (blocks emBits 8))
  (requires True)
  (ensures fun em -> if emBits % 8 > 0 then v em.[0] < pow2 (emBits % 8) else v em.[0] < pow2 8)

let pss_encode a sLen salt msgLen msg emBits =
  let mHash = Hash.hash a msg in
  let hLen = Hash.hash_length a in

  //m1 = [8 * 0x00; mHash; salt]
  let m1Len = 8 + hLen + sLen in
  let m1 = create m1Len (u8 0) in
  let m1 = update_sub m1 8 hLen mHash in
  let m1 = update_sub m1 (8 + hLen) sLen salt in
  let m1Hash = Hash.hash a m1 in

  //db = [0x00;..; 0x00; 0x01; salt]
  let emLen = blocks emBits 8 in
  let dbLen = emLen - hLen - 1 in
  let db = create dbLen (u8 0) in
  let last_before_salt = dbLen - sLen - 1 in
  let db = db.[last_before_salt] <- u8 1 in
  let db = update_sub db (last_before_salt + 1) sLen salt in

  let dbMask = mgf_hash a hLen m1Hash dbLen in
  let maskedDB = xor_bytes db dbMask in
  let maskedDB = db_zero maskedDB emBits in

  //em = [maskedDB; m1Hash; 0xbc]
  let em = create emLen (u8 0) in
  let em = update_sub em 0 dbLen maskedDB in
  let em = update_sub em dbLen hLen m1Hash in
  assert (v em.[0] == v maskedDB.[0]);
  em.[emLen - 1] <- u8 0xbc


val pss_verify_:
    a:Hash.hash_alg{hash_is_supported a}
  -> sLen:size_nat{sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a}
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen}
  -> emBits:size_pos{sLen + Hash.hash_length a + 2 <= blocks emBits 8}
  -> em:lbytes (blocks emBits 8) ->
  Tot bool

let pss_verify_ a sLen msgLen msg emBits em =
  let hLen = Hash.hash_length a in
  let emLen = blocks emBits 8 in
  let dbLen = emLen - hLen - 1 in
  let maskedDB = sub em 0 dbLen in
  let m1Hash = sub em dbLen hLen in

  let dbMask = mgf_hash a hLen m1Hash dbLen in
  let db = xor_bytes dbMask maskedDB in
  let db = db_zero db emBits in

  let padLen = emLen - sLen - hLen - 1 in
  let pad2 = create padLen (u8 0) in
  let pad2 = pad2.[padLen - 1] <- u8 0x01 in

  let pad  = sub db 0 padLen in
  let salt = sub db padLen sLen in

  if not (lbytes_eq pad pad2) then false
  else begin
    let mHash = Hash.hash a msg in
    let m1Len = 8 + hLen + sLen in
    let m1 = create m1Len (u8 0) in
    let m1 = update_sub m1 8 hLen mHash in
    let m1 = update_sub m1 (8 + hLen) sLen salt in
    let m1Hash0 = Hash.hash a m1 in
    lbytes_eq m1Hash0 m1Hash
  end


val pss_verify:
    a:Hash.hash_alg{hash_is_supported a}
  -> sLen:size_nat{sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a}
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen}
  -> emBits:size_pos
  -> em:lbytes (blocks emBits 8) ->
  Tot bool

let pss_verify a sLen msgLen msg emBits em =
  let emLen = blocks emBits 8 in
  let msBits = emBits % 8 in

  let em_0 = if msBits > 0 then em.[0] &. (u8 0xff <<. size msBits) else u8 0 in
  let em_last = em.[emLen - 1] in

  if (emLen < sLen + Hash.hash_length a + 2) then false
  else begin
    if not (FStar.UInt8.(Lib.RawIntTypes.u8_to_UInt8 em_last =^ 0xbcuy) &&
       FStar.UInt8.(Lib.RawIntTypes.u8_to_UInt8 em_0 =^ 0uy))
    then false
    else pss_verify_ a sLen msgLen msg emBits em end


val os2ip_lemma: emBits:size_pos -> em:lbytes (blocks emBits 8) -> Lemma
  (requires (if emBits % 8 > 0 then v em.[0] < pow2 (emBits % 8) else v em.[0] < pow2 8))
  (ensures  os2ip #(blocks emBits 8) em < pow2 emBits)

let os2ip_lemma emBits em =
  let emLen = blocks emBits 8 in

  let aux (emBits:size_pos{emBits % 8 > 0}) : Lemma (8 * ((emBits - 1) / 8) + emBits % 8 <= emBits) =
    assert ((emBits - 1) / 8 <= emBits / 8) in

  if emBits % 8 > 0 then begin
    nat_from_intseq_be_slice_lemma em 1;
    nat_from_intseq_be_lemma0 (slice em 0 1);
    assert (nat_from_bytes_be em == nat_from_bytes_be (slice em 1 emLen) + pow2 ((emLen - 1) * 8) * v em.[0]);
    assert (nat_from_bytes_be em < pow2 ((emLen - 1) * 8) + pow2 ((emLen - 1) * 8) * v em.[0]);

    calc (<=) {
      pow2 ((emLen - 1) * 8) + pow2 ((emLen - 1) * 8) * v em.[0];
      (==) { Math.Lemmas.distributivity_add_right (pow2 (8 * (emLen - 1))) 1 (v em.[0]) }
      pow2 (8 * (emLen - 1)) * (1 + v em.[0]);
      (<=) { Math.Lemmas.lemma_mult_le_left (pow2 (8 * emLen - 1)) (v em.[0] + 1) (pow2 (emBits % 8)) }
      pow2 (8 * (emLen - 1)) * pow2 (emBits % 8);
      (==) { Math.Lemmas.pow2_plus (8 * (emLen - 1)) (emBits % 8) }
      pow2 (8 * ((emBits - 1) / 8) + emBits % 8);
      (<=) { aux emBits; Math.Lemmas.pow2_le_compat emBits ((emBits - 1) - (emBits - 1) % 8 + emBits % 8) }
      pow2 emBits;
    };
   assert (nat_from_bytes_be em < pow2 emBits) end
  else begin
    assert (nat_from_bytes_be em < pow2 (emLen * 8));
    assert (emBits == 8 * emLen) end


val rsapss_sign_:
    a:Hash.hash_alg{hash_is_supported a}
  -> modBits:modBits_t
  -> skey:rsapss_skey modBits
  -> sLen:size_nat{
    sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a /\
    sLen + Hash.hash_length a + 2 <= blocks (modBits - 1) 8}
  -> salt:lbytes sLen
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen} ->
  tuple2 bool (lbytes (blocks modBits 8))
val rsapss_sign_:
    a:Hash.hash_alg{hash_is_supported a}
  -> modBits:modBits_t
  -> skey:rsapss_skey modBits
  -> sLen:size_nat{
    sLen + Hash.hash_length a + 8 <= max_size_t /\
    (sLen + Hash.hash_length a + 8) `Hash.less_than_max_input_length` a /\
    sLen + Hash.hash_length a + 2 <= blocks (modBits - 1) 8}
  -> salt:lbytes sLen
  -> msgLen:nat{msgLen `Hash.less_than_max_input_length` a}
  -> msg:bytes{length msg == msgLen} ->
  tuple2 bool (lbytes (blocks modBits 8))

    a: Spec.Hash.Definitions.hash_alg{Spec.RSAPSS.hash_is_supported a} ->
    modBits: Spec.RSAPSS.modBits_t ->
    skey: Spec.RSAPSS.rsapss_skey modBits ->
    sLen:
      Lib.IntTypes.size_nat
        { sLen + Spec.Hash.Definitions.hash_length a + 8 <= Lib.IntTypes.max_size_t /\
          Spec.Hash.Definitions.less_than_max_input_length (sLen +
              Spec.Hash.Definitions.hash_length a +
              8)
            a /\
          sLen + Spec.Hash.Definitions.hash_length a + 2 <= Spec.RSAPSS.blocks (modBits - 1) 8 } ->
    salt: Lib.ByteSequence.lbytes sLen ->
    msgLen: Prims.nat{Spec.Hash.Definitions.less_than_max_input_length msgLen a} ->
    msg: Lib.ByteSequence.bytes{Lib.Sequence.length msg == msgLen}
  -> Prims.bool * Lib.ByteSequence.lbytes (Spec.RSAPSS.blocks modBits 8)