diff --git a/Cargo.lock b/Cargo.lock index 98586d802..ca684ffb5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2200,6 +2200,8 @@ dependencies = [ name = "private_tx_linkedproof" version = "0.1.0" dependencies = [ + "env_logger", + "once_cell", "sunscreen", ] diff --git a/examples/private_tx_linkedproof/Cargo.toml b/examples/private_tx_linkedproof/Cargo.toml index d7734fb29..09b9faa4c 100644 --- a/examples/private_tx_linkedproof/Cargo.toml +++ b/examples/private_tx_linkedproof/Cargo.toml @@ -5,3 +5,5 @@ edition = "2021" [dependencies] sunscreen = { path = "../../sunscreen", features = ["linkedproofs"] } +env_logger = { workspace = true } +once_cell = { workspace = true } diff --git a/examples/private_tx_linkedproof/src/main.rs b/examples/private_tx_linkedproof/src/main.rs index 639c90084..9ce566911 100644 --- a/examples/private_tx_linkedproof/src/main.rs +++ b/examples/private_tx_linkedproof/src/main.rs @@ -1,69 +1,506 @@ +#![allow(unused)] + +//! In this example, we demonstrate how to use a [`LinkedProof`] to verify a private transaction. +//! We assume the private transactions are implemented on a transparent system like a blockchain, +//! where deterministic computation is performed on encrypted data. The linked proofs allow us to +//! validate the encrypted inputs. + +use std::collections::{hash_map::Entry, HashMap}; + +use once_cell::sync::Lazy; use sunscreen::{ bulletproofs::BulletproofsBackend, fhe_program, - linked::LogProofBuilder, + linked::{LinkedProof, LogProofBuilder, Sdlp}, types::{ bfv::Signed, zkp::{AsFieldElement, BfvSigned, BulletproofsField, ConstrainCmp, Field, FieldSpec}, Cipher, }, - zkp_program, zkp_var, Compiler, FheZkpRuntime, PlainModulusConstraint, Result, + zkp_program, zkp_var, Ciphertext, CompiledFheProgram, CompiledZkpProgram, Compiler, + FheZkpApplication, FheZkpRuntime, Params, PlainModulusConstraint, PrivateKey, PublicKey, + Result, ZkpProgramInput, }; +/// Subtract the transaction amount from the sender's balance. #[fhe_program(scheme = "bfv")] -fn update_balance(tx: Cipher, balance: Signed) -> Cipher { +fn update_balance_sender(balance: Cipher, tx: Cipher) -> Cipher { balance - tx } +/// Add the transaction amount to the receiver's balance. +#[fhe_program(scheme = "bfv")] +fn update_balance_receiver(balance: Cipher, tx: Cipher) -> Cipher { + balance + tx +} + +/// Validate a transfer transaction. #[zkp_program] -fn valid_transaction(#[linked] tx: BfvSigned, #[public] balance: Field) { +fn validate_transfer( + #[linked] tx: BfvSigned, + #[linked] sender_balance: BfvSigned, +) { let tx = tx.into_field_elem(); + let sender_balance = sender_balance.into_field_elem(); - // Constraint that tx is less than or equal to balance - tx.constrain_le_bounded(balance, 64); - - // Constraint that tx is greater than zero + // Transaction amount must be greater than 0. tx.constrain_gt_bounded(zkp_var!(0), 64); + // Transaction amount cannot exceed sender's balance. + tx.constrain_le_bounded(sender_balance, 64); +} + +/// Validate deposit/registration. The deposit amount is public, but we must prove that the +/// provided ciphertext encrypts the deposit amount. +#[zkp_program] +fn validate_deposit( + #[linked] encrypted_deposit: BfvSigned, + #[public] public_deposit: Field, +) { + let encrypted_deposit = encrypted_deposit.into_field_elem(); + encrypted_deposit.constrain_eq(public_deposit); +} + +/// A way to identify a user. +type Username = String; + +/// Perspective of a user. +pub struct User { + pub name: Username, + pub public_key: PublicKey, + private_key: PrivateKey, + // This app holds ZKP programs used to make proofs + app: App, + // The runtime is used for encryption/decryption and creating proofs + runtime: FheZkpRuntime, +} + +impl User { + pub fn new(name: &str) -> Result { + let app = App::new()?; + let runtime = FheZkpRuntime::new(app.params(), &BulletproofsBackend::new())?; + let (public_key, private_key) = runtime.generate_keys()?; + Ok(Self { + name: name.to_string(), + runtime, + public_key, + private_key, + app, + }) + } + + /// Create a private, validated transfer to send to another user. + pub fn create_transfer>( + &self, + chain: &Chain, + amount: i64, + receiver: U, + ) -> Result { + let receiver = receiver.into(); + let mut builder = LogProofBuilder::new(&self.runtime); + + // Encrypt tx amount under sender's public key. + println!(" {}: encrypting {} under own key", self.name, amount); + let (encrypted_amount_sender, amount_linked) = + builder.encrypt_returning_link(&Signed::from(amount), &self.public_key)?; + + // Encrypt tx amount under receiver's public key, implicitly proving that the two + // ciphertexts encrypt the same value. + println!( + " {}: encrypting {} under receiver key", + self.name, amount + ); + let recv_pk = chain.keys.get(&receiver).unwrap(); + let encrypted_amount_receiver = builder.encrypt_msg(&amount_linked, recv_pk)?; + + // Decrypt current balance, needed to prove tx validity + let balance_enc = chain.balances.get(&self.name).unwrap(); + let (balance, balance_linked) = + builder.decrypt_returning_link::(balance_enc, &self.private_key)?; + + // Create transfer proof + println!( + " {}: creating transfer linkedproof, proving {} <= {}", + self.name, amount, balance + ); + let proof = builder + .zkp_program(self.app.get_transfer_zkp())? + .linked_input(amount_linked) + .linked_input(balance_linked) + .build_linkedproof()?; + + Ok(Transfer { + proof, + sender: self.name.clone(), + receiver, + encrypted_amount_sender, + encrypted_amount_receiver, + }) + } + + /// Create a public deposit to a private balance. + pub fn create_deposit(&self, amount: i64) -> Result { + let mut builder = LogProofBuilder::new(&self.runtime); + + // Encrypt deposit amount + println!(" {}: encrypting and linking {}", self.name, amount); + let (amount_enc, amount_linked) = + builder.encrypt_returning_link(&Signed::from(amount), &self.public_key)?; + + // Create deposit proof + println!(" {}: creating deposit linkedproof", self.name); + let proof = builder + .zkp_program(self.app.get_deposit_zkp())? + .linked_input(amount_linked) + .public_input(BulletproofsField::from(amount)) + .build_linkedproof()?; + + Ok(Deposit { + proof, + encrypted_amount: amount_enc, + public_amount: amount, + name: self.name.clone(), + }) + } + + /// Create a refresh balance transaction. + pub fn create_refresh_balance(&self, chain: &Chain) -> Result { + let mut builder = LogProofBuilder::new(&self.runtime); + + // Decrypt current balance, returning a reference to the underlying message + let balance_encrypted = chain.balances.get(&self.name).unwrap(); + let (balance, msg) = + builder.decrypt_returning_msg::(balance_encrypted, &self.private_key)?; + + // Re-encrypt the current balance + println!(" {}: re-encrypting balance of {}", self.name, balance); + let fresh_balance = builder.encrypt_msg(&msg, &self.public_key)?; + + // Generate proof that the ciphertexts encrypt the same underlying message + println!(" {}: creating refresh balance logproof", self.name); + let proof = builder.build_logproof()?; + + Ok(RefreshBalance { + proof, + fresh_balance, + name: self.name.clone(), + }) + } + + pub fn create_register(&self, initial_deposit: i64) -> Result { + let deposit = self.create_deposit(initial_deposit)?; + Ok(Register { + public_key: self.public_key.clone(), + deposit, + }) + } +} + +/// A register transaction is an initial deposit plus an identifying public key. +#[derive(Clone)] +pub struct Register { + public_key: PublicKey, + deposit: Deposit, +} + +/// A deposit transaction. +/// +/// The SDLP in the linked proof proves that the ciphertext is a valid, fresh encryption. The R1CS +/// ZKP in the linked proof proves that the amount encrypted matches the public amount deposited. +#[derive(Clone)] +pub struct Deposit { + proof: LinkedProof, + encrypted_amount: Ciphertext, + public_amount: i64, + name: Username, +} + +/// A private transfer transaction. +/// +/// The SDLP in the linked proof proves that the ciphertexts are valid, fresh encryptions of the +/// same value. The R1CS ZKP in the linked proof proves that the amount encrypted does not exceed +/// the sender's current balance. +#[derive(Clone)] +pub struct Transfer { + proof: LinkedProof, + // Transfer amount encrypted under sender's key + encrypted_amount_sender: Ciphertext, + // Transfer amount encrypted under receiver's key + encrypted_amount_receiver: Ciphertext, + sender: Username, + receiver: Username, +} + +/// A refresh private balance transaction. +/// +/// The SDLP proves that the fresh balance is a valid, fresh encryption and also that the +/// underlying value matches the current on chain balance. +/// +/// Note that this proof is not a linked proof, as these validations can be proven by an +/// [`Sdlp`] alone. +// +// TODO do we need a BfvSigned::constrain_is_fresh_encryption that proves each coefficient's +// absolute value <= 1 and degree is less than 64? +#[derive(Clone)] +pub struct RefreshBalance { + proof: Sdlp, + fresh_balance: Ciphertext, + name: Username, +} + +/// A chain transaction. +pub enum Transaction { + Register(Register), + Deposit(Deposit), + Transfer(Transfer), + RefreshBalance(RefreshBalance), +} + +/// Perspective of the blockchain, basically just a place where user's encrypted balances are +/// stored and transparent FHE computations take place in the form of atomic transactions. +/// +/// In this simple example, assume read-only references `&Chain` provide "call" functionalities, +/// i.e. non-mutating methods for reading chain data, and mutable references `&mut Chain` provide +/// "send" functionalities, i.e. sending transactions that can mutate chain data. +pub struct Chain { + /// The current balances + balances: HashMap, + /// The user's public keys + keys: HashMap, + /// Ledger of transactions + // TODO log transactions + ledger: Vec, + /// App holding FHE and ZKP programs + app: App, + /// Runtime to run FHE programs and verify proofs + runtime: FheZkpRuntime, +} + +impl Chain { + pub fn new() -> Result { + let app = App::new()?; + let runtime = FheZkpRuntime::new(app.params(), &BulletproofsBackend::new())?; + Ok(Self { + balances: HashMap::new(), + keys: HashMap::new(), + ledger: Vec::new(), + runtime, + app, + }) + } + + pub fn register(&mut self, register: Register) -> Result<()> { + self.ledger.push(Transaction::Register(register.clone())); + let Register { + public_key, + deposit, + } = register; + + self.keys.insert(deposit.name.clone(), public_key); + self.deposit(deposit) + } + + pub fn deposit(&mut self, deposit: Deposit) -> Result<()> { + self.ledger.push(Transaction::Deposit(deposit.clone())); + let Deposit { + proof, + encrypted_amount, + public_amount, + name, + } = deposit; + + // First, verify that the encrypted amount matches the public amount + proof.verify( + self.app.get_deposit_zkp(), + vec![BulletproofsField::from(public_amount)], + vec![], + )?; + + // Then deposit into the user's balance + let pk = self.keys.get(&name).unwrap(); + match self.balances.entry(name) { + // Update existing balance + Entry::Occupied(mut entry) => { + let curr = entry.get().clone(); + let updated = self + .runtime + .run( + self.app.get_update_balance_receiver_fhe(), + vec![curr, encrypted_amount], + pk, + )? + .remove(0); + entry.insert(updated); + } + // Or insert the amount if this is an initial deposit + Entry::Vacant(entry) => { + entry.insert(encrypted_amount); + } + } + Ok(()) + } + + pub fn transfer(&mut self, transfer: Transfer) -> Result<()> { + self.ledger.push(Transaction::Transfer(transfer.clone())); + let Transfer { + proof, + encrypted_amount_sender, + encrypted_amount_receiver, + sender, + receiver, + } = transfer; + + // First verify the transfer is valid + proof.verify::(self.app.get_transfer_zkp(), vec![], vec![])?; + + // Update the sender's balance: + let sender_pk = self.keys.get(&sender).unwrap(); + let sender_balance = self.balances.get_mut(&sender).unwrap(); + let new_balance = self + .runtime + .run( + self.app.get_update_balance_sender_fhe(), + vec![sender_balance.clone(), encrypted_amount_sender], + sender_pk, + )? + .remove(0); + *sender_balance = new_balance; + + // Update receiver's balance + let receiver_pk = self.keys.get(&receiver).unwrap(); + let receiver_balance = self.balances.get_mut(&receiver).unwrap(); + let new_balance = self + .runtime + .run( + self.app.get_update_balance_receiver_fhe(), + vec![receiver_balance.clone(), encrypted_amount_receiver], + receiver_pk, + )? + .remove(0); + *receiver_balance = new_balance; + Ok(()) + } + + pub fn refresh_balance(&mut self, refresh_balance: RefreshBalance) -> Result<()> { + self.ledger + .push(Transaction::RefreshBalance(refresh_balance.clone())); + let RefreshBalance { + proof, + fresh_balance, + name, + } = refresh_balance; + + // Verify the balance refresh is valid + // TODO fix logproofs for computed ciphertexts + // proof.verify()?; + + // Use the freshly encrypted balance + self.balances + .insert(name, fresh_balance) + .expect("User should be registered"); + Ok(()) + } + + pub fn print_ledger(&self) { + for (i, tx) in self.ledger.iter().enumerate() { + match tx { + Transaction::Register(r) => println!( + "{i}. User {} registered with initial deposit {}", + r.deposit.name, r.deposit.public_amount + ), + Transaction::Deposit(d) => { + println!("{i}. User {} deposited {}", d.name, d.public_amount) + } + Transaction::Transfer(t) => println!( + "{i}. User {} transferred to {}", + t.sender, t.receiver + ), + Transaction::RefreshBalance(b) => { + println!("{i}. User {} refreshed their balance", b.name) + } + } + } + } +} + +pub struct App(&'static Lazy); + +impl App { + pub fn new() -> Result { + static APP: Lazy = Lazy::new(|| { + Compiler::new() + .fhe_program(update_balance_sender) + .fhe_program(update_balance_receiver) + .zkp_backend::() + .zkp_program(validate_transfer) + .zkp_program(validate_deposit) + .compile() + .unwrap() + }); + Ok(Self(&APP)) + } + + pub fn get_transfer_zkp(&self) -> &CompiledZkpProgram { + self.0.get_zkp_program(validate_transfer).unwrap() + } + + pub fn get_deposit_zkp(&self) -> &CompiledZkpProgram { + self.0.get_zkp_program(validate_deposit).unwrap() + } + + pub fn get_update_balance_sender_fhe(&self) -> &CompiledFheProgram { + self.0.get_fhe_program(update_balance_sender).unwrap() + } + + pub fn get_update_balance_receiver_fhe(&self) -> &CompiledFheProgram { + self.0.get_fhe_program(update_balance_receiver).unwrap() + } + + pub fn params(&self) -> &Params { + self.0.params() + } } fn main() -> Result<()> { - println!("Compiling FHE and ZKP programs..."); - let app = Compiler::new() - .fhe_program(update_balance) - // this is not strictly necessary, but will help performance - .plain_modulus_constraint(PlainModulusConstraint::Raw(1024)) - .zkp_backend::() - .zkp_program(valid_transaction) - .compile()?; - let rt = FheZkpRuntime::new(app.params(), &BulletproofsBackend::new())?; + env_logger::init(); - let valid_tx_zkp = app.get_zkp_program(valid_transaction).unwrap(); + println!("Starting a new chain..."); + let mut chain = Chain::new()?; - println!("Generating FHE keys..."); - let (public_key, _secret_key) = rt.generate_keys()?; + println!(); - let balance = 2002_i64; - let tx = 5_i64; - println!("Using balance {balance} and tx {tx}..."); - let balance = BulletproofsField::from(balance); - let tx = Signed::from(tx); + println!("Running Alice's transactions..."); + let alice = User::new("Alice")?; + let deposit = 100; + println!("Registering with a deposit of {deposit}"); + chain.register(alice.create_register(deposit)?)?; + let deposit = 50; + println!("Depositing an extra {deposit}"); + chain.deposit(alice.create_deposit(deposit)?)?; - let mut proof_builder = LogProofBuilder::new(&rt); + println!(); - println!("Encrypting and sharing transaction..."); - let (_ct, tx_msg) = proof_builder.encrypt_returning_link(&tx, &public_key)?; + println!("Running Bob's transactions..."); + let bob = User::new("Bob")?; + let deposit = 100; + println!("Registering with a deposit of {deposit}"); + chain.register(bob.create_register(deposit)?)?; + let tx = 50; + println!("Transfering {tx} to Alice"); + chain.transfer(bob.create_transfer(&chain, tx, "Alice")?)?; - println!("Building linkedproof..."); - let lp = proof_builder - .zkp_program(valid_tx_zkp)? - .linked_input(tx_msg) - .public_input(balance) - .build_linkedproof()?; + println!(); - println!("Verifying linkedproof..."); - lp.verify(valid_tx_zkp, vec![BulletproofsField::from(balance)], vec![])?; + println!("Refreshing Alice's balance..."); + let refresh_balance = alice.create_refresh_balance(&chain)?; + chain.refresh_balance(refresh_balance)?; + + println!("Done!"); + + println!(); + println!("========================== Ledger =========================="); + println!(); + chain.print_ledger(); - println!("Success! Transaction is valid."); Ok(()) }