Skip to content

Commit

Permalink
[zk-token-sdk] replace hard-coded constants with constant variables (#…
Browse files Browse the repository at this point in the history
…32274)

* add ristretto and scalar byte length constants

* add serialization and deserialization helper functions

* remove hard-coded constants in the `sigma` module

* remove hard-coded constants in the `encryption` module

* remove hard-coded constants in the `zk-token-elgamal` module

* Apply suggestions from code review

Co-authored-by: Tyera <[email protected]>

* fix docs for range proof constants

* Apply suggestions from code review

Co-authored-by: Tyera <[email protected]>

* clippy

---------

Co-authored-by: Tyera <[email protected]>
  • Loading branch information
samkim-crypto and Tyera authored Jun 28, 2023
1 parent 5dee2e4 commit 91186d3
Show file tree
Hide file tree
Showing 19 changed files with 423 additions and 228 deletions.
43 changes: 28 additions & 15 deletions zk-token-sdk/src/encryption/auth_encryption.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,19 @@ use {
zeroize::Zeroize,
};

/// Byte length of an authenticated encryption secret key
const AE_KEY_LEN: usize = 16;

/// Byte length of an authenticated encryption nonce component
const NONCE_LEN: usize = 12;

/// Byte lenth of an authenticated encryption ciphertext component
const CIPHERTEXT_LEN: usize = 24;

/// Byte length of a complete authenticated encryption ciphertext component that includes the
/// ciphertext and nonce components
const AE_CIPHERTEXT_LEN: usize = 36;

#[derive(Error, Clone, Debug, Eq, PartialEq)]
pub enum AuthenticatedEncryptionError {
#[error("key derivation method not supported")]
Expand All @@ -46,15 +59,15 @@ impl AuthenticatedEncryption {
/// This function is randomized. It internally samples a 128-bit key using `OsRng`.
#[cfg(not(target_os = "solana"))]
fn keygen() -> AeKey {
AeKey(OsRng.gen::<[u8; 16]>())
AeKey(OsRng.gen::<[u8; AE_KEY_LEN]>())
}

/// On input of an authenticated encryption key and an amount, the function returns a
/// corresponding authenticated encryption ciphertext.
#[cfg(not(target_os = "solana"))]
fn encrypt(key: &AeKey, balance: u64) -> AeCiphertext {
let mut plaintext = balance.to_le_bytes();
let nonce: Nonce = OsRng.gen::<[u8; 12]>();
let nonce: Nonce = OsRng.gen::<[u8; NONCE_LEN]>();

// The balance and the nonce have fixed length and therefore, encryption should not fail.
let ciphertext = Aes128GcmSiv::new(&key.0.into())
Expand Down Expand Up @@ -86,7 +99,7 @@ impl AuthenticatedEncryption {
}

#[derive(Debug, Zeroize)]
pub struct AeKey([u8; 16]);
pub struct AeKey([u8; AE_KEY_LEN]);
impl AeKey {
/// Deterministically derives an authenticated encryption key from a Solana signer and a public
/// seed.
Expand Down Expand Up @@ -144,7 +157,7 @@ impl AeKey {

impl EncodableKey for AeKey {
fn read<R: Read>(reader: &mut R) -> Result<Self, Box<dyn error::Error>> {
let bytes: [u8; 16] = serde_json::from_reader(reader)?;
let bytes: [u8; AE_KEY_LEN] = serde_json::from_reader(reader)?;
Ok(Self(bytes))
}

Expand All @@ -158,7 +171,7 @@ impl EncodableKey for AeKey {

impl SeedDerivable for AeKey {
fn from_seed(seed: &[u8]) -> Result<Self, Box<dyn error::Error>> {
const MINIMUM_SEED_LEN: usize = 16;
const MINIMUM_SEED_LEN: usize = AE_KEY_LEN;

if seed.len() < MINIMUM_SEED_LEN {
return Err(AuthenticatedEncryptionError::SeedLengthTooShort.into());
Expand All @@ -168,7 +181,7 @@ impl SeedDerivable for AeKey {
hasher.update(seed);
let result = hasher.finalize();

Ok(Self(result[..16].try_into()?))
Ok(Self(result[..AE_KEY_LEN].try_into()?))
}

fn from_seed_and_derivation_path(
Expand All @@ -191,8 +204,8 @@ impl SeedDerivable for AeKey {

/// For the purpose of encrypting balances for the spl token accounts, the nonce and ciphertext
/// sizes should always be fixed.
type Nonce = [u8; 12];
type Ciphertext = [u8; 24];
type Nonce = [u8; NONCE_LEN];
type Ciphertext = [u8; CIPHERTEXT_LEN];

/// Authenticated encryption nonce and ciphertext
#[derive(Debug, Default, Clone)]
Expand All @@ -205,20 +218,20 @@ impl AeCiphertext {
AuthenticatedEncryption::decrypt(key, self)
}

pub fn to_bytes(&self) -> [u8; 36] {
let mut buf = [0_u8; 36];
buf[..12].copy_from_slice(&self.nonce);
buf[12..].copy_from_slice(&self.ciphertext);
pub fn to_bytes(&self) -> [u8; AE_CIPHERTEXT_LEN] {
let mut buf = [0_u8; AE_CIPHERTEXT_LEN];
buf[..NONCE_LEN].copy_from_slice(&self.nonce);
buf[NONCE_LEN..].copy_from_slice(&self.ciphertext);
buf
}

pub fn from_bytes(bytes: &[u8]) -> Option<AeCiphertext> {
if bytes.len() != 36 {
if bytes.len() != AE_CIPHERTEXT_LEN {
return None;
}

let nonce = bytes[..32].try_into().ok()?;
let ciphertext = bytes[32..].try_into().ok()?;
let nonce = bytes[..NONCE_LEN].try_into().ok()?;
let ciphertext = bytes[NONCE_LEN..].try_into().ok()?;

Some(AeCiphertext { nonce, ciphertext })
}
Expand Down
8 changes: 6 additions & 2 deletions zk-token-sdk/src/encryption/discrete_log.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#![cfg(not(target_os = "solana"))]

use {
crate::RISTRETTO_POINT_LEN,
curve25519_dalek::{
constants::RISTRETTO_BASEPOINT_POINT as G,
ristretto::RistrettoPoint,
Expand All @@ -32,6 +33,9 @@ use {
const TWO16: u64 = 65536; // 2^16
const TWO17: u64 = 131072; // 2^17

/// Maximum number of threads permitted for discrete log computation
const MAX_THREAD: usize = 65536;

#[derive(Error, Clone, Debug, Eq, PartialEq)]
pub enum DiscreteLogError {
#[error("discrete log number of threads not power-of-two")]
Expand Down Expand Up @@ -61,7 +65,7 @@ pub struct DiscreteLog {
}

#[derive(Serialize, Deserialize, Default)]
pub struct DecodePrecomputation(HashMap<[u8; 32], u16>);
pub struct DecodePrecomputation(HashMap<[u8; RISTRETTO_POINT_LEN], u16>);

/// Builds a HashMap of 2^16 elements
#[allow(dead_code)]
Expand Down Expand Up @@ -110,7 +114,7 @@ impl DiscreteLog {
/// Adjusts number of threads in a discrete log instance.
pub fn num_threads(&mut self, num_threads: usize) -> Result<(), DiscreteLogError> {
// number of threads must be a positive power-of-two integer
if num_threads == 0 || (num_threads & (num_threads - 1)) != 0 || num_threads > 65536 {
if num_threads == 0 || (num_threads & (num_threads - 1)) != 0 || num_threads > MAX_THREAD {
return Err(DiscreteLogError::DiscreteLogThreads);
}

Expand Down
68 changes: 44 additions & 24 deletions zk-token-sdk/src/encryption/elgamal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,14 @@
//! discrete log to recover the originally encrypted value.
use {
crate::encryption::{
discrete_log::DiscreteLog,
pedersen::{Pedersen, PedersenCommitment, PedersenOpening, G, H},
crate::{
encryption::{
discrete_log::DiscreteLog,
pedersen::{
Pedersen, PedersenCommitment, PedersenOpening, G, H, PEDERSEN_COMMITMENT_LEN,
},
},
RISTRETTO_POINT_LEN, SCALAR_LEN,
},
base64::{prelude::BASE64_STANDARD, Engine},
core::ops::{Add, Mul, Sub},
Expand Down Expand Up @@ -50,6 +55,21 @@ use {
},
};

/// Byte length of a decrypt handle
const DECRYPT_HANDLE_LEN: usize = RISTRETTO_POINT_LEN;

/// Byte length of an ElGamal ciphertext
const ELGAMAL_CIPHERTEXT_LEN: usize = PEDERSEN_COMMITMENT_LEN + DECRYPT_HANDLE_LEN;

/// Byte length of an ElGamal public key
const ELGAMAL_PUBKEY_LEN: usize = RISTRETTO_POINT_LEN;

/// Byte length of an ElGamal secret key
const ELGAMAL_SECRET_KEY_LEN: usize = SCALAR_LEN;

/// Byte length of an ElGamal keypair
const ELGAMAL_KEYPAIR_LEN: usize = ELGAMAL_PUBKEY_LEN + ELGAMAL_SECRET_KEY_LEN;

#[derive(Error, Clone, Debug, Eq, PartialEq)]
pub enum ElGamalError {
#[error("key derivation method not supported")]
Expand Down Expand Up @@ -209,21 +229,21 @@ impl ElGamalKeypair {
&self.secret
}

pub fn to_bytes(&self) -> [u8; 64] {
let mut bytes = [0u8; 64];
bytes[..32].copy_from_slice(&self.public.to_bytes());
bytes[32..].copy_from_slice(self.secret.as_bytes());
pub fn to_bytes(&self) -> [u8; ELGAMAL_KEYPAIR_LEN] {
let mut bytes = [0u8; ELGAMAL_KEYPAIR_LEN];
bytes[..ELGAMAL_PUBKEY_LEN].copy_from_slice(&self.public.to_bytes());
bytes[ELGAMAL_PUBKEY_LEN..].copy_from_slice(self.secret.as_bytes());
bytes
}

pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
if bytes.len() != 64 {
if bytes.len() != ELGAMAL_KEYPAIR_LEN {
return None;
}

Some(Self {
public: ElGamalPubkey::from_bytes(&bytes[..32])?,
secret: ElGamalSecretKey::from_bytes(bytes[32..].try_into().ok()?)?,
public: ElGamalPubkey::from_bytes(&bytes[..ELGAMAL_PUBKEY_LEN])?,
secret: ElGamalSecretKey::from_bytes(bytes[ELGAMAL_PUBKEY_LEN..].try_into().ok()?)?,
})
}

Expand Down Expand Up @@ -317,12 +337,12 @@ impl ElGamalPubkey {
&self.0
}

pub fn to_bytes(&self) -> [u8; 32] {
pub fn to_bytes(&self) -> [u8; ELGAMAL_PUBKEY_LEN] {
self.0.compress().to_bytes()
}

pub fn from_bytes(bytes: &[u8]) -> Option<ElGamalPubkey> {
if bytes.len() != 32 {
if bytes.len() != ELGAMAL_PUBKEY_LEN {
return None;
}

Expand Down Expand Up @@ -428,7 +448,7 @@ impl ElGamalSecretKey {

/// Derive an ElGamal secret key from an entropy seed.
pub fn from_seed(seed: &[u8]) -> Result<Self, ElGamalError> {
const MINIMUM_SEED_LEN: usize = 32;
const MINIMUM_SEED_LEN: usize = ELGAMAL_SECRET_KEY_LEN;

if seed.len() < MINIMUM_SEED_LEN {
return Err(ElGamalError::SeedLengthTooShort);
Expand All @@ -453,11 +473,11 @@ impl ElGamalSecretKey {
ElGamal::decrypt_u32(self, ciphertext)
}

pub fn as_bytes(&self) -> &[u8; 32] {
pub fn as_bytes(&self) -> &[u8; ELGAMAL_SECRET_KEY_LEN] {
self.0.as_bytes()
}

pub fn to_bytes(&self) -> [u8; 32] {
pub fn to_bytes(&self) -> [u8; ELGAMAL_SECRET_KEY_LEN] {
self.0.to_bytes()
}

Expand Down Expand Up @@ -554,21 +574,21 @@ impl ElGamalCiphertext {
}
}

pub fn to_bytes(&self) -> [u8; 64] {
let mut bytes = [0u8; 64];
bytes[..32].copy_from_slice(&self.commitment.to_bytes());
bytes[32..].copy_from_slice(&self.handle.to_bytes());
pub fn to_bytes(&self) -> [u8; ELGAMAL_CIPHERTEXT_LEN] {
let mut bytes = [0u8; ELGAMAL_CIPHERTEXT_LEN];
bytes[..PEDERSEN_COMMITMENT_LEN].copy_from_slice(&self.commitment.to_bytes());
bytes[PEDERSEN_COMMITMENT_LEN..].copy_from_slice(&self.handle.to_bytes());
bytes
}

pub fn from_bytes(bytes: &[u8]) -> Option<ElGamalCiphertext> {
if bytes.len() != 64 {
if bytes.len() != ELGAMAL_CIPHERTEXT_LEN {
return None;
}

Some(ElGamalCiphertext {
commitment: PedersenCommitment::from_bytes(&bytes[..32])?,
handle: DecryptHandle::from_bytes(&bytes[32..])?,
commitment: PedersenCommitment::from_bytes(&bytes[..PEDERSEN_COMMITMENT_LEN])?,
handle: DecryptHandle::from_bytes(&bytes[PEDERSEN_COMMITMENT_LEN..])?,
})
}

Expand Down Expand Up @@ -676,12 +696,12 @@ impl DecryptHandle {
&self.0
}

pub fn to_bytes(&self) -> [u8; 32] {
pub fn to_bytes(&self) -> [u8; DECRYPT_HANDLE_LEN] {
self.0.compress().to_bytes()
}

pub fn from_bytes(bytes: &[u8]) -> Option<DecryptHandle> {
if bytes.len() != 32 {
if bytes.len() != DECRYPT_HANDLE_LEN {
return None;
}

Expand Down
15 changes: 9 additions & 6 deletions zk-token-sdk/src/encryption/grouped_elgamal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,13 @@
//!
use {
crate::encryption::{
discrete_log::DiscreteLog,
elgamal::{DecryptHandle, ElGamalCiphertext, ElGamalPubkey, ElGamalSecretKey},
pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
crate::{
encryption::{
discrete_log::DiscreteLog,
elgamal::{DecryptHandle, ElGamalCiphertext, ElGamalPubkey, ElGamalSecretKey},
pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
},
RISTRETTO_POINT_LEN,
},
curve25519_dalek::scalar::Scalar,
thiserror::Error,
Expand Down Expand Up @@ -163,7 +166,7 @@ impl<const N: usize> GroupedElGamalCiphertext<N> {
/// `(N+1) * 32`.
fn expected_byte_length() -> usize {
N.checked_add(1)
.and_then(|length| length.checked_mul(32))
.and_then(|length| length.checked_mul(RISTRETTO_POINT_LEN))
.unwrap()
}

Expand All @@ -181,7 +184,7 @@ impl<const N: usize> GroupedElGamalCiphertext<N> {
return None;
}

let mut iter = bytes.chunks(32);
let mut iter = bytes.chunks(RISTRETTO_POINT_LEN);
let commitment = PedersenCommitment::from_bytes(iter.next()?)?;

let mut handles = Vec::with_capacity(N);
Expand Down
15 changes: 11 additions & 4 deletions zk-token-sdk/src/encryption/pedersen.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#[cfg(not(target_os = "solana"))]
use rand::rngs::OsRng;
use {
crate::{RISTRETTO_POINT_LEN, SCALAR_LEN},
core::ops::{Add, Mul, Sub},
curve25519_dalek::{
constants::{RISTRETTO_BASEPOINT_COMPRESSED, RISTRETTO_BASEPOINT_POINT},
Expand All @@ -17,6 +18,12 @@ use {
zeroize::Zeroize,
};

/// Byte length of a Pedersen opening.
const PEDERSEN_OPENING_LEN: usize = SCALAR_LEN;

/// Byte length of a Pedersen commitment.
pub(crate) const PEDERSEN_COMMITMENT_LEN: usize = RISTRETTO_POINT_LEN;

lazy_static::lazy_static! {
/// Pedersen base point for encoding messages to be committed.
pub static ref G: RistrettoPoint = RISTRETTO_BASEPOINT_POINT;
Expand Down Expand Up @@ -82,11 +89,11 @@ impl PedersenOpening {
PedersenOpening(Scalar::random(&mut OsRng))
}

pub fn as_bytes(&self) -> &[u8; 32] {
pub fn as_bytes(&self) -> &[u8; PEDERSEN_OPENING_LEN] {
self.0.as_bytes()
}

pub fn to_bytes(&self) -> [u8; 32] {
pub fn to_bytes(&self) -> [u8; PEDERSEN_OPENING_LEN] {
self.0.to_bytes()
}

Expand Down Expand Up @@ -177,12 +184,12 @@ impl PedersenCommitment {
&self.0
}

pub fn to_bytes(&self) -> [u8; 32] {
pub fn to_bytes(&self) -> [u8; PEDERSEN_COMMITMENT_LEN] {
self.0.compress().to_bytes()
}

pub fn from_bytes(bytes: &[u8]) -> Option<PedersenCommitment> {
if bytes.len() != 32 {
if bytes.len() != PEDERSEN_COMMITMENT_LEN {
return None;
}

Expand Down
7 changes: 7 additions & 0 deletions zk-token-sdk/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,3 +38,10 @@ pub mod zk_token_elgamal;
pub mod zk_token_proof_instruction;
pub mod zk_token_proof_program;
pub mod zk_token_proof_state;

/// Byte length of a compressed Ristretto point or scalar in Curve255519
const UNIT_LEN: usize = 32;
/// Byte length of a compressed Ristretto point in Curve25519
const RISTRETTO_POINT_LEN: usize = UNIT_LEN;
/// Byte length of a scalar in Curve25519
const SCALAR_LEN: usize = UNIT_LEN;
Loading

0 comments on commit 91186d3

Please sign in to comment.