lib.rs 3.5 KB

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