clutch.approxis.examples.kem_elgamal

From Ltac2 Require Import Ltac2.
Set Default Proof Mode "Classic".
From clutch.prob_lang Require Import advantage typing.tychk.
From clutch.approxis Require Import map reltac2 approxis option.
From clutch.clutch.examples.crypto Require ElGamal_bijection.
From clutch.approxis.examples Require Import
  valgroup diffie_hellman prf_local_state security_aux symmetric_init option xor
  ElGamal_defs bounded_oracle pubkey advantage_laws iterable_expression KEM.
From mathcomp Require Import ssrbool.
From mathcomp Require fingroup.fingroup.
Set Default Proof Using "All".
Import ElGamal_bijection.bij_nat.
Import valgroup_notation.
(* Import fingroup_Notations. *)
Import map.

Section ElGamal_KEM.

  Import KEM.aux.

  Context {vg : val_group}. (* A group on a subset of values. *)
  Context {cg : clutch_group_struct}. (* Implementations of the vg group operations *)
  Context {vgg : @val_group_generator vg}. (* G is generated by g. *)
  Context {cgg : @clutch_group_generator vg cg vgg}. (* g is well-typed *)

  Let N := S n''.
  Let SymKey := N.
  Let Input := N.
  Variable SymInput : nat.
  Let SymOutput := N.

  Local Instance sym_params : SYM_init_params := {|
      card_key := SymKey
    ; card_message := SymInput
    ; card_cipher := SymOutput
  |}.

  Context {symmetric_scheme : SYM_init}.
  Variable key_carrier : (Σ : gFunctors), approxisRGS Σ lrel Σ.

  Axiom sym_keygen_sem_typed : (Σ : gFunctors) (A : approxisRGS Σ),
     REL symmetric_init.keygen << symmetric_init.keygen : () key_carrier Σ A.

  Variable vg_of_symkey : val.
  Variable symkey_of_vg : val.
  Axiom vg_of_symkey_sem_typed : (Σ : gFunctors) (A : approxisRGS Σ),
     REL vg_of_symkey << vg_of_symkey : key_carrier Σ A lrel_G.
  Axiom symkey_of_vg_sem_typed : (Σ : gFunctors) (A : approxisRGS Σ),
     REL symkey_of_vg << symkey_of_vg : lrel_G key_carrier Σ A.

  Definition keygen_kem : val := keygen.

  Definition encaps : val :=
    λ: "pk",
      let: "k" := symmetric_init.get_keygen sym_scheme #() in
      let: "kg" := vg_of_symkey "k" in
      let: "c_kem" := enc "pk" "kg" in
      ("c_kem", "k").

  Definition decaps : val :=
    λ: "sk" "c",
      symkey_of_vg (dec "sk" "c").

  Local Instance elgamal_kem_struct : KEM_STRUCT := {|
      KeyGen := keygen_kem
    ; Encaps := encaps
    ; Decaps := decaps
  |}.

  Definition elgamal_sharedsecret_sem_type (Σ : gFunctors) (A : approxisRGS Σ) : lrel Σ :=
    @lrel_int_bounded Σ 0 SymKey.
  Definition elgamal_cipher_sem_type (Σ : gFunctors) (A : approxisRGS Σ) : lrel Σ :=
    lrel_G * lrel_G.
  Definition elgamal_publickey_sem_type (Σ : gFunctors) (A : approxisRGS Σ) : lrel Σ :=
    lrel_G.
  Definition elgamal_secretkey_sem_type (Σ : gFunctors) (A : approxisRGS Σ) : lrel Σ :=
    @lrel_int_bounded Σ 0 N.

  Definition rand_bound (n : nat) : val :=
    λ: <>, rand #n.
  (* Definition rand_group_elt : val :=
    rec: "rejection_sampler" <> :=
      let: "r" := rand N in match: vg_of_int "r" with | SOME "g" => "g" | NONE => "rejection_sampler" () 
      end. *)

    Definition rand_group_elt : val :=
      λ: <>, let: "a" := rand #N in vgval g ^ "a".
    Definition rand_cipher_elgamal : val :=
      λ: <>, let: "a" := rand #N in
        let: "b" := rand #N in (vgval g ^ "a", vgval g ^ "b").

  Lemma rand_bounded_sem_typed (n : nat) (Σ : gFunctors) (A : approxisRGS Σ) :
     REL rand_bound n << rand_bound n : () lrel_int_bounded 0 n.
  Proof with (rel_pures_l; rel_pures_r).
    rewrite /rand_bound. rel_arrow_val.
    iIntros (v1 v2 [eq1 eq2]); subst...
    rel_apply refines_couple_UU; first done.
    iIntros (r Hrbound). iModIntro.
    rel_vals.
  Qed.

  Lemma rand_G_sem_typed (Σ : gFunctors) (A : approxisRGS Σ)
    {G : clutch_group (vg := vg) (cg := cg)} :
     REL rand_group_elt << rand_group_elt : () lrel_G.
  Proof with (rel_pures_l; rel_pures_r).
    rewrite /rand_group_elt. rel_arrow_val.
    iIntros (v1 v2 [eq1 eq2]); subst...
    rel_apply refines_couple_UU; first done.
    iIntros (r Hrbound). iModIntro...
    rel_apply refines_exp_l;
    rel_apply refines_exp_r.
    rel_vals.
  Qed.

  Lemma rand_cipher_elgamal_sem_typed (Σ : gFunctors) (A : approxisRGS Σ)
    {G : clutch_group (vg := vg) (cg := cg)} :
     REL rand_cipher_elgamal << rand_cipher_elgamal : () lrel_G * lrel_G.
  Proof with (rel_pures_l; rel_pures_r).
    rewrite /rand_cipher_elgamal. rel_arrow_val.
    iIntros (v1 v2 [eq1 eq2]); subst...
    rel_apply refines_couple_UU; first done.
    iIntros (a Habound). iModIntro...
    rel_apply refines_couple_UU; first done.
    iIntros (b Hbbound). iModIntro...
    rel_apply refines_exp_l;
    rel_apply refines_exp_r;
    rel_apply refines_exp_l;
    rel_apply refines_exp_r...
    rel_vals; iExists _; done.
  Qed.

  Context {GS : (Σ : gFunctors) (A : approxisRGS Σ),
    @clutch_group Σ A vg cg}. (* reasonable assumption *)

  Definition elgamal_sharedsecret : typ := {|
      carrier := key_carrier
    ; rnd := symmetric_init.keygen
    ; rnd_typed := sym_keygen_sem_typed
  |}.
  Definition elgamal_cipher : typ := {|
      carrier := elgamal_cipher_sem_type
    ; rnd := rand_cipher_elgamal
    ; rnd_typed := λ (Σ : gFunctors) (A : approxisRGS Σ),
      @rand_cipher_elgamal_sem_typed Σ A (GS Σ A)
  |}.
  Definition elgamal_publickey : typ := {|
      carrier := elgamal_publickey_sem_type
    ; rnd := rand_group_elt
    ; rnd_typed := λ (Σ : gFunctors) (A : approxisRGS Σ),
      @rand_G_sem_typed Σ A (GS Σ A)
  |}.
  Definition elgamal_privatekey : typ := {|
      carrier := elgamal_secretkey_sem_type
    ; rnd := rand_bound SymKey
    ; rnd_typed := rand_bounded_sem_typed SymKey
  |}.

  Local Instance elgamal_kem_params : KEM_PARAMS := {|
      SharedSecret := elgamal_sharedsecret
    ; Ciphertext := elgamal_cipher
    ; PublicKey := elgamal_publickey
    ; SecretKey := elgamal_privatekey
  |}.

  Lemma elgamal_keygen_sem_typed (Σ : gFunctors) (A : approxisRGS Σ) :
     REL KeyGen << KeyGen : lrel_unit SecretKey * PublicKey.
  Proof with (rel_pures_l; rel_pures_r).
    rewrite /KeyGen. simpl.
    rewrite /keygen_kem/keygen.
    rel_arrow_val.
    iIntros (v1 v2 [eq1 eq2]); subst...
    rel_apply refines_couple_UU; first done.
    iIntros (n Hnbound). iModIntro...
    rel_apply refines_exp_l.
    rel_apply refines_exp_r... rel_vals.
    - rewrite /elgamal_sharedsecret_sem_type/SymKey/N.
      iExists n.
      iPureIntro; repeat split; lia.
    - rewrite /elgamal_publickey_sem_type.
      iExists _. done.
  Qed.

  Lemma elgamal_encaps_sem_typed (Σ : gFunctors) (A : approxisRGS Σ) :
     REL Encaps << Encaps : PublicKey Ciphertext * SharedSecret.
  Proof with (rel_pures_l; rel_pures_r).
    rewrite /Encaps. simpl.
    rewrite /encaps. rel_arrow_val.
    rewrite /elgamal_publickey_sem_type.
    iIntros (v1 v2 [x [eq1 eq2]]); subst...
    rewrite /get_keygen...
    rel_bind_l (symmetric_init.keygen _).
    rel_bind_r (symmetric_init.keygen _).
    rel_apply refines_bind.
    {
      rel_apply refines_app.
      - rel_apply sym_keygen_sem_typed.
      - rel_vals.
    }
    iIntros (k k') "#Hrelk"...
    rel_bind_l (vg_of_symkey _).
    rel_bind_r (vg_of_symkey _).
    rel_apply (refines_bind).
    {
      rel_apply refines_app.
      - rel_apply vg_of_symkey_sem_typed.
      - rel_vals.
    }
    iIntros (v1 v2 [kg [eq1 eq2]]); subst...
    rewrite /enc...
    rel_apply refines_couple_UU; first done.
    iIntros (b Hbbound); iModIntro...
    rel_apply refines_exp_l.
    rel_apply refines_exp_r...
    rel_apply refines_exp_l.
    rel_apply refines_exp_r...
    rel_apply refines_mult_l.
    rel_apply refines_mult_r...
    rel_vals.
    - iExists _. done.
    - iExists _. done.
    - done.
  Qed.

  Lemma elgamal_decaps_sem_typed (Σ : gFunctors) (A : approxisRGS Σ) :
     REL Decaps << Decaps : SecretKey Ciphertext SharedSecret.
  Proof with (rel_pures_l; rel_pures_r).
    rewrite /Decaps. simpl.
    rewrite /decaps.
    rewrite /elgamal_sharedsecret_sem_type/elgamal_cipher_sem_type
      /elgamal_secretkey_sem_type.
    rel_arrow_val.
    iIntros (v1 v2 [x [eq1 [eq2 Hbound]]]); subst...
    rel_arrow_val.
    iIntros (kg1 kg2 [v1 [v2 [v3 [v4 [eqpair1 [eqpair2
      [[g1 [eq1 eq2]] [g2 [eq1' eq2']]]
]]]]]]); subst...
    rel_apply refines_app.
    - rel_apply symkey_of_vg_sem_typed.
    - rewrite /dec...
      rewrite -(Z2Nat.id x); last lia.
      rel_apply_l refines_exp_l.
      rel_apply_r refines_exp_r.
      rel_apply_l refines_inv_l.
      rel_apply_r refines_inv_r.
      rel_apply_l refines_mult_l.
      rel_apply_r refines_mult_r.
      rel_vals.
  Qed.

  Local Instance elgamal_kem : KEM := {|
      KeyGen_typed := elgamal_keygen_sem_typed
    ; Encaps_typed := elgamal_encaps_sem_typed
    ; Decaps_typed := elgamal_decaps_sem_typed
  |}.

  Class KEM {KEM_struct : KEM_STRUCT} {KEM_params : KEM_PARAMS} :=
    { KeyGen_typed `{!approxisRGS Σ} : REL KeyGen << KeyGen : lrel_unit SecretKey * PublicKey
    ; Encaps_typed `{!approxisRGS Σ} : REL Encaps << Encaps : PublicKey Ciphertext * SharedSecret
    ; Decaps_typed `{!approxisRGS Σ} : REL Decaps << Decaps : SecretKey Ciphertext SharedSecret
    }.

End ElGamal_KEM.