Rename crypto crate to sdk

zk-token-sdk/src/instruction/close_account.rs

@@ -0,0 +1,163 @@
+use {
+    crate::pod::*,
+    bytemuck::{Pod, Zeroable},
+};
+#[cfg(not(target_arch = "bpf"))]
+use {
+    crate::{
+        encryption::elgamal::{ElGamalCT, ElGamalSK},
+        errors::ProofError,
+        instruction::Verifiable,
+        transcript::TranscriptProtocol,
+    },
+    curve25519_dalek::{
+        ristretto::RistrettoPoint,
+        scalar::Scalar,
+        traits::{IsIdentity, MultiscalarMul},
+    },
+    merlin::Transcript,
+    rand::rngs::OsRng,
+    std::convert::TryInto,
+};
+
+/// This struct includes the cryptographic proof *and* the account data information needed to verify
+/// the proof
+///
+/// - The pre-instruction should call CloseAccountData::verify_proof(&self)
+/// - The actual program should check that `balance` is consistent with what is
+///   currently stored in the confidential token account
+///
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct CloseAccountData {
+    /// The source account available balance in encrypted form
+    pub balance: PodElGamalCT, // 64 bytes
+
+    /// Proof that the source account available balance is zero
+    pub proof: CloseAccountProof, // 64 bytes
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl CloseAccountData {
+    pub fn new(source_sk: &ElGamalSK, balance: ElGamalCT) -> Self {
+        let proof = CloseAccountProof::new(source_sk, &balance);
+
+        CloseAccountData {
+            balance: balance.into(),
+            proof,
+        }
+    }
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl Verifiable for CloseAccountData {
+    fn verify(&self) -> Result<(), ProofError> {
+        let balance = self.balance.try_into()?;
+        self.proof.verify(&balance)
+    }
+}
+
+/// This struct represents the cryptographic proof component that certifies that the encrypted
+/// balance is zero
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+#[allow(non_snake_case)]
+pub struct CloseAccountProof {
+    pub R: PodCompressedRistretto, // 32 bytes
+    pub z: PodScalar,              // 32 bytes
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_arch = "bpf"))]
+impl CloseAccountProof {
+    fn transcript_new() -> Transcript {
+        Transcript::new(b"CloseAccountProof")
+    }
+
+    pub fn new(source_sk: &ElGamalSK, balance: &ElGamalCT) -> Self {
+        let mut transcript = Self::transcript_new();
+
+        // add a domain separator to record the start of the protocol
+        transcript.close_account_proof_domain_sep();
+
+        // extract the relevant scalar and Ristretto points from the input
+        let s = source_sk.get_scalar();
+        let C = balance.decrypt_handle.get_point();
+
+        // generate a random masking factor that also serves as a nonce
+        let r = Scalar::random(&mut OsRng); // using OsRng for now
+        let R = (r * C).compress();
+
+        // record R on transcript and receive a challenge scalar
+        transcript.append_point(b"R", &R);
+        let c = transcript.challenge_scalar(b"c");
+
+        // compute the masked secret key
+        let z = c * s + r;
+
+        CloseAccountProof {
+            R: R.into(),
+            z: z.into(),
+        }
+    }
+
+    pub fn verify(&self, balance: &ElGamalCT) -> Result<(), ProofError> {
+        let mut transcript = Self::transcript_new();
+
+        // add a domain separator to record the start of the protocol
+        transcript.close_account_proof_domain_sep();
+
+        // extract the relevant scalar and Ristretto points from the input
+        let C = balance.message_comm.get_point();
+        let D = balance.decrypt_handle.get_point();
+
+        let R = self.R.into();
+        let z = self.z.into();
+
+        // generate a challenge scalar
+        //
+        // use `append_point` as opposed to `validate_and_append_point` as the ciphertext is
+        // already guaranteed to be well-formed
+        transcript.append_point(b"R", &R);
+        let c = transcript.challenge_scalar(b"c");
+
+        // decompress R or return verification error
+        let R = R.decompress().ok_or(ProofError::VerificationError)?;
+
+        // check the required algebraic relation
+        let check = RistrettoPoint::multiscalar_mul(vec![z, -c, -Scalar::one()], vec![D, C, R]);
+
+        if check.is_identity() {
+            Ok(())
+        } else {
+            Err(ProofError::VerificationError)
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use {super::*, crate::encryption::elgamal::ElGamal};
+
+    #[test]
+    fn test_close_account_correctness() {
+        let (source_pk, source_sk) = ElGamal::keygen();
+
+        // If account balance is 0, then the proof should succeed
+        let balance = source_pk.encrypt(0_u64);
+
+        let proof = CloseAccountProof::new(&source_sk, &balance);
+        assert!(proof.verify(&balance).is_ok());
+
+        // If account balance is not zero, then the proof verification should fail
+        let balance = source_pk.encrypt(55_u64);
+
+        let proof = CloseAccountProof::new(&source_sk, &balance);
+        assert!(proof.verify(&balance).is_err());
+
+        // A zeroed cyphertext should be considered as an account balance of 0
+        let zeroed_ct: ElGamalCT = PodElGamalCT::zeroed().try_into().unwrap();
+        let proof = CloseAccountProof::new(&source_sk, &zeroed_ct);
+        assert!(proof.verify(&zeroed_ct).is_ok());
+    }
+}

zk-token-sdk/src/instruction/mod.rs

@@ -0,0 +1,21 @@
+mod close_account;
+pub mod transfer;
+mod update_account_pk;
+mod withdraw;
+
+#[cfg(not(target_arch = "bpf"))]
+use crate::errors::ProofError;
+pub use {
+    close_account::CloseAccountData,
+    transfer::{
+        TransferComms, TransferData, TransferEphemeralState, TransferPubKeys,
+        TransferRangeProofData, TransferValidityProofData,
+    },
+    update_account_pk::UpdateAccountPkData,
+    withdraw::WithdrawData,
+};
+
+#[cfg(not(target_arch = "bpf"))]
+pub trait Verifiable {
+    fn verify(&self) -> Result<(), ProofError>;
+}

zk-token-sdk/src/instruction/transfer.rs

@@ -0,0 +1,546 @@
+use {
+    crate::pod::*,
+    bytemuck::{Pod, Zeroable},
+};
+#[cfg(not(target_arch = "bpf"))]
+use {
+    crate::{
+        encryption::{
+            elgamal::{ElGamalCT, ElGamalPK, ElGamalSK},
+            pedersen::{Pedersen, PedersenBase, PedersenComm, PedersenDecHandle, PedersenOpen},
+        },
+        errors::ProofError,
+        instruction::Verifiable,
+        range_proof::RangeProof,
+        transcript::TranscriptProtocol,
+    },
+    curve25519_dalek::{
+        ristretto::{CompressedRistretto, RistrettoPoint},
+        scalar::Scalar,
+        traits::{IsIdentity, MultiscalarMul, VartimeMultiscalarMul},
+    },
+    merlin::Transcript,
+    rand::rngs::OsRng,
+    std::convert::TryInto,
+};
+
+/// Just a grouping struct for the data required for the two transfer instructions. It is
+/// convenient to generate the two components jointly as they share common components.
+pub struct TransferData {
+    pub range_proof: TransferRangeProofData,
+    pub validity_proof: TransferValidityProofData,
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl TransferData {
+    pub fn new(
+        transfer_amount: u64,
+        spendable_balance: u64,
+        spendable_ct: ElGamalCT,
+        source_pk: ElGamalPK,
+        source_sk: &ElGamalSK,
+        dest_pk: ElGamalPK,
+        auditor_pk: ElGamalPK,
+    ) -> Self {
+        // split and encrypt transfer amount
+        //
+        // encryption is a bit more involved since we are generating each components of an ElGamal
+        // ciphertext separately.
+        let (amount_lo, amount_hi) = split_u64_into_u32(transfer_amount);
+
+        let (comm_lo, open_lo) = Pedersen::commit(amount_lo);
+        let (comm_hi, open_hi) = Pedersen::commit(amount_hi);
+
+        let handle_source_lo = source_pk.gen_decrypt_handle(&open_lo);
+        let handle_dest_lo = dest_pk.gen_decrypt_handle(&open_lo);
+        let handle_auditor_lo = auditor_pk.gen_decrypt_handle(&open_lo);
+
+        let handle_source_hi = source_pk.gen_decrypt_handle(&open_hi);
+        let handle_dest_hi = dest_pk.gen_decrypt_handle(&open_hi);
+        let handle_auditor_hi = auditor_pk.gen_decrypt_handle(&open_hi);
+
+        // message encoding as Pedersen commitments, which will be included in range proof data
+        let amount_comms = TransferComms {
+            lo: comm_lo.into(),
+            hi: comm_hi.into(),
+        };
+
+        // decryption handles, which will be included in the validity proof data
+        let decryption_handles_lo = TransferHandles {
+            source: handle_source_lo.into(),
+            dest: handle_dest_lo.into(),
+            auditor: handle_auditor_lo.into(),
+        };
+
+        let decryption_handles_hi = TransferHandles {
+            source: handle_source_hi.into(),
+            dest: handle_dest_hi.into(),
+            auditor: handle_auditor_hi.into(),
+        };
+
+        // grouping of the public keys for the transfer
+        let transfer_public_keys = TransferPubKeys {
+            source_pk: source_pk.into(),
+            dest_pk: dest_pk.into(),
+            auditor_pk: auditor_pk.into(),
+        };
+
+        // subtract transfer amount from the spendable ciphertext
+        let spendable_comm = spendable_ct.message_comm;
+        let spendable_handle = spendable_ct.decrypt_handle;
+
+        let new_spendable_balance = spendable_balance - transfer_amount;
+        let new_spendable_comm = spendable_comm - combine_u32_comms(comm_lo, comm_hi);
+        let new_spendable_handle =
+            spendable_handle - combine_u32_handles(handle_source_lo, handle_source_hi);
+
+        let new_spendable_ct = ElGamalCT {
+            message_comm: new_spendable_comm,
+            decrypt_handle: new_spendable_handle,
+        };
+
+        // range_proof and validity_proof should be generated together
+        let (transfer_proofs, ephemeral_state) = TransferProofs::new(
+            source_sk,
+            &source_pk,
+            &dest_pk,
+            &auditor_pk,
+            (amount_lo as u64, amount_hi as u64),
+            &open_lo,
+            &open_hi,
+            new_spendable_balance,
+            &new_spendable_ct,
+        );
+
+        // generate data components
+        let range_proof = TransferRangeProofData {
+            amount_comms,
+            proof: transfer_proofs.range_proof,
+            ephemeral_state,
+        };
+
+        let validity_proof = TransferValidityProofData {
+            decryption_handles_lo,
+            decryption_handles_hi,
+            transfer_public_keys,
+            new_spendable_ct: new_spendable_ct.into(),
+            proof: transfer_proofs.validity_proof,
+            ephemeral_state,
+        };
+
+        TransferData {
+            range_proof,
+            validity_proof,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct TransferRangeProofData {
+    /// The transfer amount encoded as Pedersen commitments
+    pub amount_comms: TransferComms, // 64 bytes
+
+    /// Proof that certifies:
+    ///   1. the source account has enough funds for the transfer (i.e. the final balance is a
+    ///      64-bit positive number)
+    ///   2. the transfer amount is a 64-bit positive number
+    pub proof: PodRangeProof128, // 736 bytes
+
+    /// Ephemeral state between the two transfer instruction data
+    pub ephemeral_state: TransferEphemeralState, // 128 bytes
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl Verifiable for TransferRangeProofData {
+    fn verify(&self) -> Result<(), ProofError> {
+        let mut transcript = Transcript::new(b"TransferRangeProof");
+
+        // standard range proof verification
+        let proof: RangeProof = self.proof.try_into()?;
+        proof.verify_with(
+            vec![
+                &self.ephemeral_state.spendable_comm_verification.into(),
+                &self.amount_comms.lo.into(),
+                &self.amount_comms.hi.into(),
+            ],
+            vec![64_usize, 32_usize, 32_usize],
+            Some(self.ephemeral_state.x.into()),
+            Some(self.ephemeral_state.z.into()),
+            &mut transcript,
+        )
+    }
+}
+
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct TransferValidityProofData {
+    /// The decryption handles that allow decryption of the lo-bits
+    pub decryption_handles_lo: TransferHandles, // 96 bytes
+
+    /// The decryption handles that allow decryption of the hi-bits
+    pub decryption_handles_hi: TransferHandles, // 96 bytes
+
+    /// The public encryption keys associated with the transfer: source, dest, and auditor
+    pub transfer_public_keys: TransferPubKeys, // 96 bytes
+
+    /// The final spendable ciphertext after the transfer
+    pub new_spendable_ct: PodElGamalCT, // 64 bytes
+
+    /// Proof that certifies that the decryption handles are generated correctly
+    pub proof: ValidityProof, // 160 bytes
+
+    /// Ephemeral state between the two transfer instruction data
+    pub ephemeral_state: TransferEphemeralState, // 128 bytes
+}
+
+/// The joint data that is shared between the two transfer instructions.
+///
+/// Identical ephemeral data should be included in the two transfer instructions and this should be
+/// checked by the ZK Token program.
+#[derive(Clone, Copy, Pod, Zeroable, PartialEq)]
+#[repr(C)]
+pub struct TransferEphemeralState {
+    pub spendable_comm_verification: PodPedersenComm, // 32 bytes
+    pub x: PodScalar,                                 // 32 bytes
+    pub z: PodScalar,                                 // 32 bytes
+    pub t_x_blinding: PodScalar,                      // 32 bytes
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl Verifiable for TransferValidityProofData {
+    fn verify(&self) -> Result<(), ProofError> {
+        self.proof.verify(
+            &self.new_spendable_ct.try_into()?,
+            &self.decryption_handles_lo,
+            &self.decryption_handles_hi,
+            &self.transfer_public_keys,
+            &self.ephemeral_state,
+        )
+    }
+}
+
+/// Just a grouping struct for the two proofs that are needed for a transfer instruction. The two
+/// proofs have to be generated together as they share joint data.
+pub struct TransferProofs {
+    pub range_proof: PodRangeProof128,
+    pub validity_proof: ValidityProof,
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_arch = "bpf"))]
+impl TransferProofs {
+    #[allow(clippy::too_many_arguments)]
+    #[allow(clippy::many_single_char_names)]
+    pub fn new(
+        source_sk: &ElGamalSK,
+        source_pk: &ElGamalPK,
+        dest_pk: &ElGamalPK,
+        auditor_pk: &ElGamalPK,
+        transfer_amt: (u64, u64),
+        lo_open: &PedersenOpen,
+        hi_open: &PedersenOpen,
+        new_spendable_balance: u64,
+        new_spendable_ct: &ElGamalCT,
+    ) -> (Self, TransferEphemeralState) {
+        // TODO: should also commit to pubkeys and commitments later
+        let mut transcript_validity_proof = merlin::Transcript::new(b"TransferValidityProof");
+
+        let H = PedersenBase::default().H;
+        let D = new_spendable_ct.decrypt_handle.get_point();
+        let s = source_sk.get_scalar();
+
+        // Generate proof for the new spendable ciphertext
+        let r_new = Scalar::random(&mut OsRng);
+        let y = Scalar::random(&mut OsRng);
+        let R = RistrettoPoint::multiscalar_mul(vec![y, r_new], vec![D, H]).compress();
+
+        transcript_validity_proof.append_point(b"R", &R);
+        let c = transcript_validity_proof.challenge_scalar(b"c");
+
+        let z = s + c * y;
+        let new_spendable_open = PedersenOpen(c * r_new);
+
+        let spendable_comm_verification =
+            Pedersen::commit_with(new_spendable_balance, &new_spendable_open);
+
+        // Generate proof for the transfer amounts
+        let t_1_blinding = PedersenOpen::random(&mut OsRng);
+        let t_2_blinding = PedersenOpen::random(&mut OsRng);
+
+        let u = transcript_validity_proof.challenge_scalar(b"u");
+        let P_joint = RistrettoPoint::multiscalar_mul(
+            vec![Scalar::one(), u, u * u],
+            vec![
+                source_pk.get_point(),
+                dest_pk.get_point(),
+                auditor_pk.get_point(),
+            ],
+        );
+        let T_joint = (new_spendable_open.get_scalar() * P_joint).compress();
+        let T_1 = (t_1_blinding.get_scalar() * P_joint).compress();
+        let T_2 = (t_2_blinding.get_scalar() * P_joint).compress();
+
+        transcript_validity_proof.append_point(b"T_1", &T_1);
+        transcript_validity_proof.append_point(b"T_2", &T_2);
+
+        // define the validity proof
+        let validity_proof = ValidityProof {
+            R: R.into(),
+            z: z.into(),
+            T_joint: T_joint.into(),
+            T_1: T_1.into(),
+            T_2: T_2.into(),
+        };
+
+        // generate the range proof
+        let mut transcript_range_proof = Transcript::new(b"TransferRangeProof");
+        let (range_proof, x, z) = RangeProof::create_with(
+            vec![new_spendable_balance, transfer_amt.0, transfer_amt.1],
+            vec![64, 32, 32],
+            vec![&new_spendable_open, lo_open, hi_open],
+            &t_1_blinding,
+            &t_2_blinding,
+            &mut transcript_range_proof,
+        );
+
+        // define ephemeral state
+        let ephemeral_state = TransferEphemeralState {
+            spendable_comm_verification: spendable_comm_verification.into(),
+            x: x.into(),
+            z: z.into(),
+            t_x_blinding: range_proof.t_x_blinding.into(),
+        };
+
+        (
+            Self {
+                range_proof: range_proof.try_into().expect("valid range_proof"),
+                validity_proof,
+            },
+            ephemeral_state,
+        )
+    }
+}
+
+/// Proof components for transfer instructions.
+///
+/// These two components should be output by a RangeProof creation function.
+#[allow(non_snake_case)]
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct ValidityProof {
+    // Proof component for the spendable ciphertext components: R
+    pub R: PodCompressedRistretto, // 32 bytes
+    // Proof component for the spendable ciphertext components: z
+    pub z: PodScalar, // 32 bytes
+    // Proof component for the transaction amount components: T_src
+    pub T_joint: PodCompressedRistretto, // 32 bytes
+    // Proof component for the transaction amount components: T_1
+    pub T_1: PodCompressedRistretto, // 32 bytes
+    // Proof component for the transaction amount components: T_2
+    pub T_2: PodCompressedRistretto, // 32 bytes
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_arch = "bpf"))]
+impl ValidityProof {
+    pub fn verify(
+        self,
+        new_spendable_ct: &ElGamalCT,
+        decryption_handles_lo: &TransferHandles,
+        decryption_handles_hi: &TransferHandles,
+        transfer_public_keys: &TransferPubKeys,
+        ephemeral_state: &TransferEphemeralState,
+    ) -> Result<(), ProofError> {
+        let mut transcript = Transcript::new(b"TransferValidityProof");
+
+        let source_pk: ElGamalPK = transfer_public_keys.source_pk.try_into()?;
+        let dest_pk: ElGamalPK = transfer_public_keys.dest_pk.try_into()?;
+        let auditor_pk: ElGamalPK = transfer_public_keys.auditor_pk.try_into()?;
+
+        // verify Pedersen commitment in the ephemeral state
+        let C_ephemeral: CompressedRistretto = ephemeral_state.spendable_comm_verification.into();
+
+        let C = new_spendable_ct.message_comm.get_point();
+        let D = new_spendable_ct.decrypt_handle.get_point();
+
+        let R = self.R.into();
+        let z: Scalar = self.z.into();
+
+        transcript.validate_and_append_point(b"R", &R)?;
+        let c = transcript.challenge_scalar(b"c");
+
+        let R = R.decompress().ok_or(ProofError::VerificationError)?;
+
+        let spendable_comm_verification =
+            RistrettoPoint::multiscalar_mul(vec![Scalar::one(), -z, c], vec![C, D, R]).compress();
+
+        if C_ephemeral != spendable_comm_verification {
+            return Err(ProofError::VerificationError);
+        }
+
+        // derive joint public key
+        let u = transcript.challenge_scalar(b"u");
+        let P_joint = RistrettoPoint::vartime_multiscalar_mul(
+            vec![Scalar::one(), u, u * u],
+            vec![
+                source_pk.get_point(),
+                dest_pk.get_point(),
+                auditor_pk.get_point(),
+            ],
+        );
+
+        // check well-formedness of decryption handles
+        let t_x_blinding: Scalar = ephemeral_state.t_x_blinding.into();
+        let T_1: CompressedRistretto = self.T_1.into();
+        let T_2: CompressedRistretto = self.T_2.into();
+
+        let x = ephemeral_state.x.into();
+        let z: Scalar = ephemeral_state.z.into();
+
+        let handle_source_lo: PedersenDecHandle = decryption_handles_lo.source.try_into()?;
+        let handle_dest_lo: PedersenDecHandle = decryption_handles_lo.dest.try_into()?;
+        let handle_auditor_lo: PedersenDecHandle = decryption_handles_lo.auditor.try_into()?;
+
+        let D_joint: CompressedRistretto = self.T_joint.into();
+        let D_joint = D_joint.decompress().ok_or(ProofError::VerificationError)?;
+
+        let D_joint_lo = RistrettoPoint::vartime_multiscalar_mul(
+            vec![Scalar::one(), u, u * u],
+            vec![
+                handle_source_lo.get_point(),
+                handle_dest_lo.get_point(),
+                handle_auditor_lo.get_point(),
+            ],
+        );
+
+        let handle_source_hi: PedersenDecHandle = decryption_handles_hi.source.try_into()?;
+        let handle_dest_hi: PedersenDecHandle = decryption_handles_hi.dest.try_into()?;
+        let handle_auditor_hi: PedersenDecHandle = decryption_handles_hi.auditor.try_into()?;
+
+        let D_joint_hi = RistrettoPoint::vartime_multiscalar_mul(
+            vec![Scalar::one(), u, u * u],
+            vec![
+                handle_source_hi.get_point(),
+                handle_dest_hi.get_point(),
+                handle_auditor_hi.get_point(),
+            ],
+        );
+
+        // TODO: combine Pedersen commitment verification above to here for efficiency
+        // TODO: might need to add an additional proof-of-knowledge here (additional 64 byte)
+        let mega_check = RistrettoPoint::optional_multiscalar_mul(
+            vec![-t_x_blinding, x, x * x, z * z, z * z * z, z * z * z * z],
+            vec![
+                Some(P_joint),
+                T_1.decompress(),
+                T_2.decompress(),
+                Some(D_joint),
+                Some(D_joint_lo),
+                Some(D_joint_hi),
+            ],
+        )
+        .ok_or(ProofError::VerificationError)?;
+
+        if mega_check.is_identity() {
+            Ok(())
+        } else {
+            Err(ProofError::VerificationError)
+        }
+    }
+}
+
+/// The ElGamal public keys needed for a transfer
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct TransferPubKeys {
+    pub source_pk: PodElGamalPK,  // 32 bytes
+    pub dest_pk: PodElGamalPK,    // 32 bytes
+    pub auditor_pk: PodElGamalPK, // 32 bytes
+}
+
+/// The transfer amount commitments needed for a transfer
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct TransferComms {
+    pub lo: PodPedersenComm, // 32 bytes
+    pub hi: PodPedersenComm, // 32 bytes
+}
+
+/// The decryption handles needed for a transfer
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct TransferHandles {
+    pub source: PodPedersenDecHandle,  // 32 bytes
+    pub dest: PodPedersenDecHandle,    // 32 bytes
+    pub auditor: PodPedersenDecHandle, // 32 bytes
+}
+
+/// Split u64 number into two u32 numbers
+#[cfg(not(target_arch = "bpf"))]
+pub fn split_u64_into_u32(amt: u64) -> (u32, u32) {
+    let lo = amt as u32;
+    let hi = (amt >> 32) as u32;
+
+    (lo, hi)
+}
+
+/// Constant for 2^32
+#[cfg(not(target_arch = "bpf"))]
+const TWO_32: u64 = 4294967296;
+
+#[cfg(not(target_arch = "bpf"))]
+pub fn combine_u32_comms(comm_lo: PedersenComm, comm_hi: PedersenComm) -> PedersenComm {
+    comm_lo + comm_hi * Scalar::from(TWO_32)
+}
+
+#[cfg(not(target_arch = "bpf"))]
+pub fn combine_u32_handles(
+    handle_lo: PedersenDecHandle,
+    handle_hi: PedersenDecHandle,
+) -> PedersenDecHandle {
+    handle_lo + handle_hi * Scalar::from(TWO_32)
+}
+
+#[cfg(not(target_arch = "bpf"))]
+pub fn combine_u32_ciphertexts(ct_lo: ElGamalCT, ct_hi: ElGamalCT) -> ElGamalCT {
+    ct_lo + ct_hi * Scalar::from(TWO_32)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::encryption::elgamal::ElGamal;
+
+    #[test]
+    fn test_transfer_correctness() {
+        // ElGamal keys for source, destination, and auditor accounts
+        let (source_pk, source_sk) = ElGamal::keygen();
+        let (dest_pk, _) = ElGamal::keygen();
+        let (auditor_pk, _) = ElGamal::keygen();
+
+        // create source account spendable ciphertext
+        let spendable_balance: u64 = 77;
+        let spendable_ct = source_pk.encrypt(spendable_balance);
+
+        // transfer amount
+        let transfer_amount: u64 = 55;
+
+        // create transfer data
+        let transfer_data = TransferData::new(
+            transfer_amount,
+            spendable_balance,
+            spendable_ct,
+            source_pk,
+            &source_sk,
+            dest_pk,
+            auditor_pk,
+        );
+
+        // verify range proof
+        assert!(transfer_data.range_proof.verify().is_ok());
+
+        // verify ciphertext validity proof
+        assert!(transfer_data.validity_proof.verify().is_ok());
+    }
+}

zk-token-sdk/src/instruction/update_account_pk.rs

@@ -0,0 +1,271 @@
+use {
+    crate::pod::*,
+    bytemuck::{Pod, Zeroable},
+};
+#[cfg(not(target_arch = "bpf"))]
+use {
+    crate::{
+        encryption::{
+            elgamal::{ElGamalCT, ElGamalPK, ElGamalSK},
+            pedersen::PedersenBase,
+        },
+        errors::ProofError,
+        instruction::Verifiable,
+        transcript::TranscriptProtocol,
+    },
+    curve25519_dalek::{
+        ristretto::RistrettoPoint,
+        scalar::Scalar,
+        traits::{IsIdentity, MultiscalarMul},
+    },
+    merlin::Transcript,
+    rand::rngs::OsRng,
+    std::convert::TryInto,
+};
+
+/// This struct includes the cryptographic proof *and* the account data information needed to verify
+/// the proof
+///
+/// - The pre-instruction should call UpdateAccountPKData::verify(&self)
+/// - The actual program should check that `current_ct` is consistent with what is
+///   currently stored in the confidential token account
+///
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct UpdateAccountPkData {
+    /// Current ElGamal encryption key
+    pub current_pk: PodElGamalPK, // 32 bytes
+
+    /// Current encrypted available balance
+    pub current_ct: PodElGamalCT, // 64 bytes
+
+    /// New ElGamal encryption key
+    pub new_pk: PodElGamalPK, // 32 bytes
+
+    /// New encrypted available balance
+    pub new_ct: PodElGamalCT, // 64 bytes
+
+    /// Proof that the current and new ciphertexts are consistent
+    pub proof: UpdateAccountPkProof, // 160 bytes
+}
+
+impl UpdateAccountPkData {
+    #[cfg(not(target_arch = "bpf"))]
+    pub fn new(
+        current_balance: u64,
+        current_ct: ElGamalCT,
+        current_pk: ElGamalPK,
+        current_sk: &ElGamalSK,
+        new_pk: ElGamalPK,
+        new_sk: &ElGamalSK,
+    ) -> Self {
+        let new_ct = new_pk.encrypt(current_balance);
+
+        let proof =
+            UpdateAccountPkProof::new(current_balance, current_sk, new_sk, &current_ct, &new_ct);
+
+        Self {
+            current_pk: current_pk.into(),
+            current_ct: current_ct.into(),
+            new_ct: new_ct.into(),
+            new_pk: new_pk.into(),
+            proof,
+        }
+    }
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl Verifiable for UpdateAccountPkData {
+    fn verify(&self) -> Result<(), ProofError> {
+        let current_ct = self.current_ct.try_into()?;
+        let new_ct = self.new_ct.try_into()?;
+        self.proof.verify(&current_ct, &new_ct)
+    }
+}
+
+/// This struct represents the cryptographic proof component that certifies that the current_ct and
+/// new_ct encrypt equal values
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+#[allow(non_snake_case)]
+pub struct UpdateAccountPkProof {
+    pub R_0: PodCompressedRistretto, // 32 bytes
+    pub R_1: PodCompressedRistretto, // 32 bytes
+    pub z_sk_0: PodScalar,           // 32 bytes
+    pub z_sk_1: PodScalar,           // 32 bytes
+    pub z_x: PodScalar,              // 32 bytes
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_arch = "bpf"))]
+impl UpdateAccountPkProof {
+    fn transcript_new() -> Transcript {
+        Transcript::new(b"UpdateAccountPkProof")
+    }
+
+    fn new(
+        current_balance: u64,
+        current_sk: &ElGamalSK,
+        new_sk: &ElGamalSK,
+        current_ct: &ElGamalCT,
+        new_ct: &ElGamalCT,
+    ) -> Self {
+        let mut transcript = Self::transcript_new();
+
+        // add a domain separator to record the start of the protocol
+        transcript.update_account_public_key_proof_domain_sep();
+
+        // extract the relevant scalar and Ristretto points from the input
+        let s_0 = current_sk.get_scalar();
+        let s_1 = new_sk.get_scalar();
+        let x = Scalar::from(current_balance);
+
+        let D_0 = current_ct.decrypt_handle.get_point();
+        let D_1 = new_ct.decrypt_handle.get_point();
+
+        let G = PedersenBase::default().G;
+
+        // generate a random masking factor that also serves as a nonce
+        let r_sk_0 = Scalar::random(&mut OsRng);
+        let r_sk_1 = Scalar::random(&mut OsRng);
+        let r_x = Scalar::random(&mut OsRng);
+
+        let R_0 = (r_sk_0 * D_0 + r_x * G).compress();
+        let R_1 = (r_sk_1 * D_1 + r_x * G).compress();
+
+        // record R_0, R_1 on transcript and receive a challenge scalar
+        transcript.append_point(b"R_0", &R_0);
+        transcript.append_point(b"R_1", &R_1);
+        let c = transcript.challenge_scalar(b"c");
+        let _w = transcript.challenge_scalar(b"w"); // for consistency of transcript
+
+        // compute the masked secret keys and amount
+        let z_sk_0 = c * s_0 + r_sk_0;
+        let z_sk_1 = c * s_1 + r_sk_1;
+        let z_x = c * x + r_x;
+
+        UpdateAccountPkProof {
+            R_0: R_0.into(),
+            R_1: R_1.into(),
+            z_sk_0: z_sk_0.into(),
+            z_sk_1: z_sk_1.into(),
+            z_x: z_x.into(),
+        }
+    }
+
+    fn verify(&self, current_ct: &ElGamalCT, new_ct: &ElGamalCT) -> Result<(), ProofError> {
+        let mut transcript = Self::transcript_new();
+
+        // add a domain separator to record the start of the protocol
+        transcript.update_account_public_key_proof_domain_sep();
+
+        // extract the relevant scalar and Ristretto points from the input
+        let C_0 = current_ct.message_comm.get_point();
+        let D_0 = current_ct.decrypt_handle.get_point();
+
+        let C_1 = new_ct.message_comm.get_point();
+        let D_1 = new_ct.decrypt_handle.get_point();
+
+        let R_0 = self.R_0.into();
+        let R_1 = self.R_1.into();
+        let z_sk_0 = self.z_sk_0.into();
+        let z_sk_1: Scalar = self.z_sk_1.into();
+        let z_x = self.z_x.into();
+
+        let G = PedersenBase::default().G;
+
+        // generate a challenge scalar
+        transcript.validate_and_append_point(b"R_0", &R_0)?;
+        transcript.validate_and_append_point(b"R_1", &R_1)?;
+        let c = transcript.challenge_scalar(b"c");
+        let w = transcript.challenge_scalar(b"w");
+
+        // decompress R_0, R_1 or return verification error
+        let R_0 = R_0.decompress().ok_or(ProofError::VerificationError)?;
+        let R_1 = R_1.decompress().ok_or(ProofError::VerificationError)?;
+
+        // check the required algebraic relation
+        let check = RistrettoPoint::multiscalar_mul(
+            vec![
+                z_sk_0,
+                z_x,
+                -c,
+                -Scalar::one(),
+                w * z_sk_1,
+                w * z_x,
+                -w * c,
+                -w * Scalar::one(),
+            ],
+            vec![D_0, G, C_0, R_0, D_1, G, C_1, R_1],
+        );
+
+        if check.is_identity() {
+            Ok(())
+        } else {
+            Err(ProofError::VerificationError)
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::encryption::elgamal::ElGamal;
+
+    #[test]
+    fn test_update_account_public_key_correctness() {
+        let (current_pk, current_sk) = ElGamal::keygen();
+        let (new_pk, new_sk) = ElGamal::keygen();
+
+        // If current_ct and new_ct encrypt same values, then the proof verification should succeed
+        let balance: u64 = 77;
+        let current_ct = current_pk.encrypt(balance);
+        let new_ct = new_pk.encrypt(balance);
+
+        let proof = UpdateAccountPkProof::new(balance, &current_sk, &new_sk, &current_ct, &new_ct);
+        assert!(proof.verify(&current_ct, &new_ct).is_ok());
+
+        // If current_ct and new_ct encrypt different values, then the proof verification should fail
+        let new_ct = new_pk.encrypt(55_u64);
+
+        let proof = UpdateAccountPkProof::new(balance, &current_sk, &new_sk, &current_ct, &new_ct);
+        assert!(proof.verify(&current_ct, &new_ct).is_err());
+
+        // A zeroed cipehrtext should be considered as an account balance of 0
+        let balance: u64 = 0;
+        let zeroed_ct_as_current_ct: ElGamalCT = PodElGamalCT::zeroed().try_into().unwrap();
+        let new_ct: ElGamalCT = new_pk.encrypt(balance);
+        let proof = UpdateAccountPkProof::new(
+            balance,
+            &current_sk,
+            &new_sk,
+            &zeroed_ct_as_current_ct,
+            &new_ct,
+        );
+        assert!(proof.verify(&zeroed_ct_as_current_ct, &new_ct).is_ok());
+
+        let current_ct: ElGamalCT = PodElGamalCT::zeroed().try_into().unwrap();
+        let zeroed_ct_as_new_ct: ElGamalCT = PodElGamalCT::zeroed().try_into().unwrap();
+        let proof = UpdateAccountPkProof::new(
+            balance,
+            &current_sk,
+            &new_sk,
+            &current_ct,
+            &zeroed_ct_as_new_ct,
+        );
+        assert!(proof.verify(&current_ct, &zeroed_ct_as_new_ct).is_ok());
+
+        let zeroed_ct_as_current_ct: ElGamalCT = PodElGamalCT::zeroed().try_into().unwrap();
+        let zeroed_ct_as_new_ct: ElGamalCT = PodElGamalCT::zeroed().try_into().unwrap();
+        let proof = UpdateAccountPkProof::new(
+            balance,
+            &current_sk,
+            &new_sk,
+            &zeroed_ct_as_current_ct,
+            &zeroed_ct_as_new_ct,
+        );
+        assert!(proof
+            .verify(&zeroed_ct_as_current_ct, &zeroed_ct_as_new_ct)
+            .is_ok());
+    }
+}

zk-token-sdk/src/instruction/withdraw.rs

@@ -0,0 +1,207 @@
+use {
+    crate::pod::*,
+    bytemuck::{Pod, Zeroable},
+};
+#[cfg(not(target_arch = "bpf"))]
+use {
+    crate::{
+        encryption::{
+            elgamal::{ElGamalCT, ElGamalPK, ElGamalSK},
+            pedersen::{PedersenBase, PedersenOpen},
+        },
+        errors::ProofError,
+        instruction::Verifiable,
+        range_proof::RangeProof,
+        transcript::TranscriptProtocol,
+    },
+    curve25519_dalek::{ristretto::RistrettoPoint, scalar::Scalar, traits::MultiscalarMul},
+    merlin::Transcript,
+    rand::rngs::OsRng,
+    std::convert::TryInto,
+};
+
+/// This struct includes the cryptographic proof *and* the account data information needed to verify
+/// the proof
+///
+/// - The pre-instruction should call WithdrawData::verify_proof(&self)
+/// - The actual program should check that `current_ct` is consistent with what is
+///   currently stored in the confidential token account TODO: update this statement
+///
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+pub struct WithdrawData {
+    /// The source account available balance *after* the withdraw (encrypted by
+    /// `source_pk`
+    pub final_balance_ct: PodElGamalCT, // 64 bytes
+
+    /// Proof that the account is solvent
+    pub proof: WithdrawProof, // 736 bytes
+}
+
+impl WithdrawData {
+    #[cfg(not(target_arch = "bpf"))]
+    pub fn new(
+        amount: u64,
+        source_pk: ElGamalPK,
+        source_sk: &ElGamalSK,
+        current_balance: u64,
+        current_balance_ct: ElGamalCT,
+    ) -> Self {
+        // subtract withdraw amount from current balance
+        //
+        // panics if current_balance < amount
+        let final_balance = current_balance - amount;
+
+        // encode withdraw amount as an ElGamal ciphertext and subtract it from
+        // current source balance
+        let amount_encoded = source_pk.encrypt_with(amount, &PedersenOpen::default());
+        let final_balance_ct = current_balance_ct - amount_encoded;
+
+        let proof = WithdrawProof::new(source_sk, final_balance, &final_balance_ct);
+
+        Self {
+            final_balance_ct: final_balance_ct.into(),
+            proof,
+        }
+    }
+}
+
+#[cfg(not(target_arch = "bpf"))]
+impl Verifiable for WithdrawData {
+    fn verify(&self) -> Result<(), ProofError> {
+        let final_balance_ct = self.final_balance_ct.try_into()?;
+        self.proof.verify(&final_balance_ct)
+    }
+}
+
+/// This struct represents the cryptographic proof component that certifies the account's solvency
+/// for withdrawal
+#[derive(Clone, Copy, Pod, Zeroable)]
+#[repr(C)]
+#[allow(non_snake_case)]
+pub struct WithdrawProof {
+    /// Wrapper for range proof: R component
+    pub R: PodCompressedRistretto, // 32 bytes
+    /// Wrapper for range proof: z component
+    pub z: PodScalar, // 32 bytes
+    /// Associated range proof
+    pub range_proof: PodRangeProof64, // 672 bytes
+}
+
+#[allow(non_snake_case)]
+#[cfg(not(target_arch = "bpf"))]
+impl WithdrawProof {
+    fn transcript_new() -> Transcript {
+        Transcript::new(b"WithdrawProof")
+    }
+
+    pub fn new(source_sk: &ElGamalSK, final_balance: u64, final_balance_ct: &ElGamalCT) -> Self {
+        let mut transcript = Self::transcript_new();
+
+        // add a domain separator to record the start of the protocol
+        transcript.withdraw_proof_domain_sep();
+
+        // extract the relevant scalar and Ristretto points from the input
+        let H = PedersenBase::default().H;
+        let D = final_balance_ct.decrypt_handle.get_point();
+        let s = source_sk.get_scalar();
+
+        // new pedersen opening
+        let r_new = Scalar::random(&mut OsRng);
+
+        // generate a random masking factor that also serves as a nonce
+        let y = Scalar::random(&mut OsRng);
+
+        let R = RistrettoPoint::multiscalar_mul(vec![y, r_new], vec![D, H]).compress();
+
+        // record R on transcript and receive a challenge scalar
+        transcript.append_point(b"R", &R);
+        let c = transcript.challenge_scalar(b"c");
+
+        // compute the masked secret key
+        let z = s + c * y;
+
+        // compute the new Pedersen commitment and opening
+        let new_open = PedersenOpen(c * r_new);
+
+        let range_proof = RangeProof::create(
+            vec![final_balance],
+            vec![64],
+            vec![&new_open],
+            &mut transcript,
+        );
+
+        WithdrawProof {
+            R: R.into(),
+            z: z.into(),
+            range_proof: range_proof.try_into().expect("range proof"),
+        }
+    }
+
+    pub fn verify(&self, final_balance_ct: &ElGamalCT) -> Result<(), ProofError> {
+        let mut transcript = Self::transcript_new();
+
+        // Add a domain separator to record the start of the protocol
+        transcript.withdraw_proof_domain_sep();
+
+        // Extract the relevant scalar and Ristretto points from the input
+        let C = final_balance_ct.message_comm.get_point();
+        let D = final_balance_ct.decrypt_handle.get_point();
+
+        let R = self.R.into();
+        let z: Scalar = self.z.into();
+
+        // generate a challenge scalar
+        transcript.validate_and_append_point(b"R", &R)?;
+        let c = transcript.challenge_scalar(b"c");
+
+        // decompress R or return verification error
+        let R = R.decompress().ok_or(ProofError::VerificationError)?;
+
+        // compute new Pedersen commitment to verify range proof with
+        let new_comm = RistrettoPoint::multiscalar_mul(vec![Scalar::one(), -z, c], vec![C, D, R]);
+
+        let range_proof: RangeProof = self.range_proof.try_into()?;
+        range_proof.verify(vec![&new_comm.compress()], vec![64_usize], &mut transcript)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::encryption::elgamal::ElGamal;
+
+    #[test]
+    #[ignore]
+    fn test_withdraw_correctness() {
+        // generate and verify proof for the proper setting
+        let (source_pk, source_sk) = ElGamal::keygen();
+
+        let current_balance: u64 = 77;
+        let current_balance_ct = source_pk.encrypt(current_balance);
+
+        let withdraw_amount: u64 = 55;
+
+        let data = WithdrawData::new(
+            withdraw_amount,
+            source_pk,
+            &source_sk,
+            current_balance,
+            current_balance_ct,
+        );
+        assert!(data.verify().is_ok());
+
+        // generate and verify proof with wrong balance
+        let wrong_balance: u64 = 99;
+        let data = WithdrawData::new(
+            withdraw_amount,
+            source_pk,
+            &source_sk,
+            wrong_balance,
+            current_balance_ct,
+        );
+        assert!(data.verify().is_err());
+
+        // TODO: test for ciphertexts that encrypt numbers outside the 0, 2^64 range
+    }
+}