lib.rs 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. #[macro_use]
  2. extern crate clap;
  3. use bellman::groth16;
  4. use bls12_381::{Bls12, Scalar};
  5. use std::collections::{HashMap, HashSet};
  6. pub mod async_serial;
  7. pub mod blockchain;
  8. pub mod bls_extensions;
  9. pub mod circuit;
  10. pub mod cli;
  11. pub mod client;
  12. pub mod crypto;
  13. pub mod endian;
  14. pub mod error;
  15. pub mod net;
  16. pub mod rpc;
  17. pub mod serial;
  18. pub mod service;
  19. pub mod state;
  20. pub mod system;
  21. pub mod tx;
  22. pub mod util;
  23. pub mod vm;
  24. pub mod vm_serial;
  25. pub mod wallet;
  26. pub use crate::bls_extensions::BlsStringConversion;
  27. pub use crate::error::{Error, Result};
  28. pub use crate::net::p2p::P2p;
  29. pub use crate::serial::{Decodable, Encodable};
  30. pub use crate::vm::{
  31. AllocType, ConstraintInstruction, CryptoOperation, VariableIndex, VariableRef,
  32. ZkVirtualMachine, ZkVmCircuit,
  33. };
  34. pub type Bytes = Vec<u8>;
  35. pub struct ZkContract {
  36. pub name: String,
  37. pub vm: ZkVirtualMachine,
  38. params_map: HashMap<String, VariableIndex>,
  39. pub params: HashMap<VariableIndex, Scalar>,
  40. public_map: bimap::BiMap<String, VariableIndex>,
  41. }
  42. pub struct ZkProof {
  43. pub public: HashMap<String, Scalar>,
  44. pub proof: groth16::Proof<Bls12>,
  45. }
  46. impl ZkContract {
  47. // Just have a load() and save()
  48. // Load the contract, do the setup, save it...
  49. pub fn setup(&mut self, filename: &str) -> Result<()> {
  50. self.vm.setup()?;
  51. let buffer = std::fs::File::create(filename)?;
  52. self.vm.params.as_ref().unwrap().write(buffer)?;
  53. Ok(())
  54. }
  55. pub fn load_setup(&mut self, filename: &str) -> Result<()> {
  56. let buffer = std::fs::File::open(filename)?;
  57. let setup = groth16::Parameters::<Bls12>::read(buffer, false)?;
  58. let vk = groth16::prepare_verifying_key(&setup.vk);
  59. self.vm.params = Some(setup);
  60. self.vm.verifying_key = Some(vk);
  61. Ok(())
  62. }
  63. pub fn param_names(&self) -> Vec<String> {
  64. self.params_map.keys().cloned().collect()
  65. }
  66. pub fn set_param(&mut self, name: &str, value: Scalar) -> Result<()> {
  67. match self.params_map.get(name) {
  68. Some(index) => {
  69. self.params.insert(*index, value);
  70. Ok(())
  71. }
  72. None => Err(Error::InvalidParamName),
  73. }
  74. }
  75. pub fn prove(&mut self) -> Result<ZkProof> {
  76. // Error if params not all set
  77. let user_params: HashSet<_> = self.params.keys().collect();
  78. let req_params: HashSet<_> = self.params_map.values().collect();
  79. if user_params != req_params {
  80. return Err(Error::MissingParams);
  81. }
  82. // execute
  83. let params = std::mem::replace(&mut self.params, HashMap::default());
  84. self.vm.initialize(&params.into_iter().collect())?;
  85. // prove
  86. let proof = self.vm.prove();
  87. let mut public = HashMap::new();
  88. for (index, value) in self.vm.public() {
  89. match self.public_map.get_by_right(&index) {
  90. Some(name) => {
  91. public.insert(name.clone(), value);
  92. }
  93. None => return Err(Error::BadContract),
  94. }
  95. }
  96. // return proof and public values (Hashmap string -> scalars)
  97. Ok(ZkProof { public, proof })
  98. }
  99. pub fn verify(&self, proof: &ZkProof) -> bool {
  100. let mut public = vec![];
  101. for (name, value) in &proof.public {
  102. match self.public_map.get_by_left(name) {
  103. Some(index) => {
  104. public.push((index, value.clone()));
  105. }
  106. None => return false,
  107. }
  108. }
  109. public.sort_by(|a, b| a.0.partial_cmp(b.0).unwrap());
  110. let (_, public): (Vec<VariableIndex>, Vec<Scalar>) = public.into_iter().unzip();
  111. // Takes proof and public values
  112. self.vm.verify(&proof.proof, &public)
  113. }
  114. }