lib.rs 3.6 KB

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