[zk-token-sdk] replace hard-coded constants with constant variables (s…

…olana-labs#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 <tyera@solana.com>

* fix docs for range proof constants

* Apply suggestions from code review

Co-authored-by: Tyera <tyera@solana.com>

* clippy

---------

Co-authored-by: Tyera <tyera@solana.com>

zk-token-sdk/src/encryption/auth_encryption.rs

@@ -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")]
@@ -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())
@@ -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.
@@ -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))
     }
 
@@ -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());
@@ -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(
@@ -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)]
@@ -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 })
     }

zk-token-sdk/src/encryption/discrete_log.rs

@@ -17,6 +17,7 @@
 #![cfg(not(target_os = "solana"))]
 
 use {
+    crate::RISTRETTO_POINT_LEN,
     curve25519_dalek::{
         constants::RISTRETTO_BASEPOINT_POINT as G,
         ristretto::RistrettoPoint,
@@ -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")]
@@ -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)]
@@ -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);
         }
 

zk-token-sdk/src/encryption/elgamal.rs

@@ -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},
@@ -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")]
@@ -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()?)?,
         })
     }
 
@@ -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;
         }
 
@@ -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);
@@ -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()
     }
 
@@ -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..])?,
         })
     }
 
@@ -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;
         }
 

zk-token-sdk/src/encryption/grouped_elgamal.rs

@@ -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,
@@ -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()
     }
 
@@ -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);

zk-token-sdk/src/encryption/pedersen.rs

@@ -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},
@@ -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;
@@ -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()
     }
 
@@ -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;
         }
 

zk-token-sdk/src/lib.rs

@@ -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;