Vote2.ml

(** Counting electronic votes, with ZK proofs of ballot well-formedness *)

(* Copyright Xavier Leroy.  License: GPL 2 or later *)

open Printf
open Cryptokit

open Algebra
open SigmaZKP

let rng = Random.(pseudo_rng (string secure_rng 32))

module EG = ElGamal.Make(P256)
module ZKP = Noninteractive(ElGamal01(P256))

(* Homomorphic addition of two EEG ciphertexts *)

let hom_add = EG.hmul

(* EEG encryption and production of a proof that the plaintext is 0 or 1.
   If the plaintext is neither 0 nor 1, an invalid proof is produced. *)

let encrypt_with_proof pk n r =
  let m = EG.encode_int n in
  let a = P256.(generator ** r) and b = P256.(pk ** r * m) in
  let p = ZKP.proof (a, b, pk) ((n = 1), r) in
  ((a, b), p)

(* Verification of a 0/1 proof *)

let check_proof pk (a, b) pf =
  ZKP.verify (a, b, pk) pf

(* Count the number of occurrences of character [c] in string [s]. *)

let count c s =
  String.fold_left
    (fun n c' -> if c' = c then n + 1 else n)
    0 s

(* Encode a vote as a pair of integers (a, b)
     a is 1 for an Alice vote, 0 otherwise
     b is 1 for a Bob vote, 0 otherwise
   Malicious votes are possible, e.g. "AAB" -> (2, 1)
*)

let encode_vote s =
  if s = "Alice" then (1, 0)
  else if s = "Bob" then (0, 1)
  else (count 'A' s, count 'B' s)

(* Each voter sends its encrypted, encoded vote to the operator,
   along with a ZK proof that the vote is well-formed:
   - proof that the number of A votes is 0 or 1
   - proof that the number of B votes is 0 or 1
   - proof that the number of A votes + the number of B votes is 0 or 1.
*)

let voter s =
  (* Get the public key from the authority *)
  let pk : EG.public_key = Multiparty.receive 0 in
  (* Encode and encrypt the vote; build the ZK proofs of well-formedness *)
  let (a, b) = encode_vote s in
  let ra = EG.rand_num ~rng and rb = EG.rand_num ~rng in
  let (a', pa) = encrypt_with_proof pk a ra
  and (b', pb) = encrypt_with_proof pk b rb
  and (_, pab) = encrypt_with_proof pk (a + b) (Z.(erem (ra + rb) P256.order)) in
  (* Send the encrypted vote and the proofs to the operator *)
  Multiparty.log "voting (%d, %d)\n" a b;
  Multiparty.send 1 (a', b', pa, pb, pab)

(* The vote operator collects the votes from the voters,
   checks the proofs of well-formedness, and sums the well-formed votes.
   Ill-formed votes are discarded.*)

let operator () =
  (* Get the public key from the authority *)
  let pk : EG.public_key = Multiparty.receive 0 in
  Multiparty.log "waiting for votes\n";
  (* Get the votes, check them, add them *)
  let rec total a b i =
    if i >= Mpi.(comm_size comm_world) then (a, b) else begin
      let (a1, b1, p1, p2, p3) :
          EG.ciphertext * EG.ciphertext * ZKP.proof * ZKP.proof * ZKP.proof =
        Multiparty.receive i in
      Multiparty.log "received encrypted vote from %d\n" i;
      (* Checking the ZK proofs *)
      let ok1 = check_proof pk a1 p1
      and ok2 = check_proof pk b1 p2
      and ok3 = check_proof pk (hom_add a1 b1) p3 in
      if not ok1 then Multiparty.log "A is not 0 or 1\n";
      if not ok2 then Multiparty.log "B is not 0 or 1\n";
      if not ok3 then Multiparty.log "A + B is not 0 or 1\n";
      if ok1 && ok2 && ok3 then begin
        Multiparty.log "valid vote from %d\n" i;
        total (hom_add a a1) (hom_add b b1) (i + 1)
      end else begin
        Multiparty.log "invalid vote from %d, ignored\n" i;
        total a b (i + 1)
      end
    end in
  let zero = EG.encrypt ~rng pk (EG.encode_int 0) in
  let (an, bn) = total zero zero 2 in
  (* Send the encrypted totals to the authority *)
  Multiparty.send 0 (an, bn)

(* The vote authority (participant 0) is the only one that has the private key.
   It performs the final decryption of totals. *)

let authority () =
  let (sk, pk) = EG.keygen ~rng () in
  (* Send the public key to all other participants *)
  for i = 1 to Multiparty.num_participants () - 1 do
    Multiparty.send i pk
  done;
  (* Receive the totals from the vote operator *)
  let (a, b) : EG.ciphertext * EG.ciphertext = Multiparty.receive 1 in
  Multiparty.log "received the encrypted totals\n";
  (* Decrypt the totals *)
  let a = EG.decode_int ~lo:0 ~hi:1000 (EG.decrypt sk a)
  and b = EG.decode_int ~lo:0 ~hi:1000 (EG.decrypt sk b) in
  (* Publish the results *)
  printf "Alice: %d\n" a;
  printf "Bob: %d\n%!" b

(* The protocol *)

let usage () =
  printf "Usage: mpirun -np 5 ./Vote2.exe <vote1> <vote2> <vote3>\n";
  printf "Each <vote> is one of\n\tAlice\n\tBob\n\tblank\n";
  printf "Malicious votes are possible, e.g. AAA (three Alice votes)\n"

let _ =
  if Multiparty.num_participants() < 5 then begin
    if Multiparty.self() = 0 then usage();
    exit 2
  end;
  begin match Multiparty.self() with
  | 0 ->
      authority()
  | 1 ->
      operator()
  | i ->
      let i = i - 1 in
      voter (if Array.length Sys.argv >= i + 1 then Sys.argv.(i) else "")
  end;
  Multiparty.barrier()