Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[zk-token-sdk] replace hard-coded constants with constant variables #32274

Merged
merged 9 commits into from
Jun 28, 2023
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