From 39701fa5604894f3dcb0dc4223fb5f0750bd6cd4 Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Thu, 11 May 2023 04:43:25 +0900
Subject: [PATCH] [zk-token-sdk] Refactor encryption error types (#31569)

* reorganize encryption error types

* remove encryption errors module

* cargo fmt
---
 .../src/encryption/auth_encryption.rs         | 11 +++++++-
 zk-token-sdk/src/encryption/discrete_log.rs   | 10 +++++++-
 zk-token-sdk/src/encryption/elgamal.rs        |  8 +++++-
 zk-token-sdk/src/encryption/errors.rs         | 25 -------------------
 zk-token-sdk/src/encryption/mod.rs            |  1 -
 5 files changed, 26 insertions(+), 29 deletions(-)
 delete mode 100644 zk-token-sdk/src/encryption/errors.rs

diff --git a/zk-token-sdk/src/encryption/auth_encryption.rs b/zk-token-sdk/src/encryption/auth_encryption.rs
index 52badaa383..4081b33749 100644
--- a/zk-token-sdk/src/encryption/auth_encryption.rs
+++ b/zk-token-sdk/src/encryption/auth_encryption.rs
@@ -3,12 +3,12 @@
 //! This module is a simple wrapper of the `Aes128GcmSiv` implementation.
 #[cfg(not(target_os = "solana"))]
 use {
-    crate::encryption::errors::AuthenticatedEncryptionError,
     aes_gcm_siv::{
         aead::{Aead, NewAead},
         Aes128GcmSiv,
     },
     rand::{rngs::OsRng, CryptoRng, Rng, RngCore},
+    thiserror::Error,
 };
 use {
     arrayref::{array_ref, array_refs},
@@ -33,6 +33,15 @@ use {
     zeroize::Zeroize,
 };
 
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum AuthenticatedEncryptionError {
+    #[error("key derivation method not supported")]
+    DerivationMethodNotSupported,
+
+    #[error("pubkey does not exist")]
+    PubkeyDoesNotExist,
+}
+
 struct AuthenticatedEncryption;
 impl AuthenticatedEncryption {
     #[cfg(not(target_os = "solana"))]
diff --git a/zk-token-sdk/src/encryption/discrete_log.rs b/zk-token-sdk/src/encryption/discrete_log.rs
index 2ebe099016..668dbecce9 100644
--- a/zk-token-sdk/src/encryption/discrete_log.rs
+++ b/zk-token-sdk/src/encryption/discrete_log.rs
@@ -17,7 +17,6 @@
 #![cfg(not(target_os = "solana"))]
 
 use {
-    crate::encryption::errors::DiscreteLogError,
     curve25519_dalek::{
         constants::RISTRETTO_BASEPOINT_POINT as G,
         ristretto::RistrettoPoint,
@@ -27,11 +26,20 @@ use {
     itertools::Itertools,
     serde::{Deserialize, Serialize},
     std::{collections::HashMap, thread},
+    thiserror::Error,
 };
 
 const TWO16: u64 = 65536; // 2^16
 const TWO17: u64 = 131072; // 2^17
 
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum DiscreteLogError {
+    #[error("discrete log number of threads not power-of-two")]
+    DiscreteLogThreads,
+    #[error("discrete log batch size too large")]
+    DiscreteLogBatchSize,
+}
+
 /// Type that captures a discrete log challenge.
 ///
 /// The goal of discrete log is to find x such that x * generator = target.
diff --git a/zk-token-sdk/src/encryption/elgamal.rs b/zk-token-sdk/src/encryption/elgamal.rs
index e31fc7a93f..b5cbde910b 100644
--- a/zk-token-sdk/src/encryption/elgamal.rs
+++ b/zk-token-sdk/src/encryption/elgamal.rs
@@ -16,7 +16,6 @@
 use {
     crate::encryption::{
         discrete_log::DiscreteLog,
-        errors::ElGamalError,
         pedersen::{Pedersen, PedersenCommitment, PedersenOpening, G, H},
     },
     core::ops::{Add, Mul, Sub},
@@ -39,6 +38,7 @@ use {
     },
     std::convert::TryInto,
     subtle::{Choice, ConstantTimeEq},
+    thiserror::Error,
     zeroize::Zeroize,
 };
 #[cfg(not(target_os = "solana"))]
@@ -52,6 +52,12 @@ use {
     },
 };
 
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum ElGamalError {
+    #[error("key derivation method not supported")]
+    DerivationMethodNotSupported,
+}
+
 /// Algorithm handle for the twisted ElGamal encryption scheme
 pub struct ElGamal;
 impl ElGamal {
diff --git a/zk-token-sdk/src/encryption/errors.rs b/zk-token-sdk/src/encryption/errors.rs
deleted file mode 100644
index 4f05bbdca1..0000000000
--- a/zk-token-sdk/src/encryption/errors.rs
+++ /dev/null
@@ -1,25 +0,0 @@
-//! Errors related to the twisted ElGamal encryption scheme.
-use thiserror::Error;
-
-#[derive(Error, Clone, Debug, Eq, PartialEq)]
-pub enum DiscreteLogError {
-    #[error("discrete log number of threads not power-of-two")]
-    DiscreteLogThreads,
-    #[error("discrete log batch size too large")]
-    DiscreteLogBatchSize,
-}
-
-#[derive(Error, Clone, Debug, Eq, PartialEq)]
-pub enum ElGamalError {
-    #[error("key derivation method not supported")]
-    DerivationMethodNotSupported,
-}
-
-#[derive(Error, Clone, Debug, Eq, PartialEq)]
-pub enum AuthenticatedEncryptionError {
-    #[error("key derivation method not supported")]
-    DerivationMethodNotSupported,
-
-    #[error("pubkey does not exist")]
-    PubkeyDoesNotExist,
-}
diff --git a/zk-token-sdk/src/encryption/mod.rs b/zk-token-sdk/src/encryption/mod.rs
index 0025ec735d..a90b88a5fa 100644
--- a/zk-token-sdk/src/encryption/mod.rs
+++ b/zk-token-sdk/src/encryption/mod.rs
@@ -13,5 +13,4 @@
 pub mod auth_encryption;
 pub mod discrete_log;
 pub mod elgamal;
-pub mod errors;
 pub mod pedersen;