From 9bbaa03932322a66b4a1b808947203186e71cbed Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Thu, 17 Nov 2016 15:03:12 +0300
Subject: [PATCH] check for sanity

---
 ethcrypto/src/lib.rs                 | 43 ++++++++++++++++++++++++++--
 ethstore/src/account/safe_account.rs |  2 +-
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/ethcrypto/src/lib.rs b/ethcrypto/src/lib.rs
index 103e750e6..c98d14027 100644
--- a/ethcrypto/src/lib.rs
+++ b/ethcrypto/src/lib.rs
@@ -34,16 +34,43 @@ pub const KEY_LENGTH: usize = 32;
 pub const KEY_ITERATIONS: usize = 10240;
 pub const KEY_LENGTH_AES: usize = KEY_LENGTH / 2;
 
+#[derive(PartialEq, Debug)]
+pub enum ScryptError {
+	// log(N) < r / 16
+	InvalidN,
+	// p <= (2^31-1 * 32)/(128 * r)
+	InvalidP,
+}
+
+impl fmt::Display for ScryptError {
+	fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
+		let s = match *self {
+			ScryptError::InvalidN => "Invalid N argument of the scrypt encryption" ,
+			ScryptError::InvalidP => "Invalid p argument of the scrypt encryption",
+		};
+
+		write!(f, "{}", s)
+	}
+}
+
 #[derive(PartialEq, Debug)]
 pub enum Error {
 	Secp(SecpError),
+	Scrypt(ScryptError),
 	InvalidMessage,
 }
 
+impl From<ScryptError> for Error {
+	fn from(err: ScryptError) -> Self {
+		Error::Scrypt(err)
+	}
+}
+
 impl fmt::Display for Error {
 	fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
 		let s = match *self {
 			Error::Secp(ref err) => err.to_string(),
+			Error::Scrypt(ref err) => err.to_string(),
 			Error::InvalidMessage => "Invalid message".into(),
 		};
 
@@ -80,13 +107,23 @@ pub fn derive_key_iterations(password: &str, salt: &[u8; 32], c: u32) -> (Vec<u8
 	(derived_right_bits.to_vec(), derived_left_bits.to_vec())
 }
 
-pub fn derive_key_scrypt(password: &str, salt: &[u8; 32], n: u32, p: u32, r: u32) -> (Vec<u8>, Vec<u8>) {
+pub fn derive_key_scrypt(password: &str, salt: &[u8; 32], n: u32, p: u32, r: u32) -> Result<(Vec<u8>, Vec<u8>), Error> {
+	// sanity checks
+	let log_n = (32 - n.leading_zeros() - 1) as u8;
+	if log_n as u32 >= r * 16 {
+		return Err(Error::Scrypt(ScryptError::InvalidN));
+	}
+
+	if p as u64 > ((u32::max_value() as u64 - 1) * 32)/(128 * (r as u64)) {
+		return Err(Error::Scrypt(ScryptError::InvalidP));
+	}
+
 	let mut derived_key = vec![0u8; KEY_LENGTH];
-	let scrypt_params = ScryptParams::new(n.trailing_zeros() as u8, r, p);
+	let scrypt_params = ScryptParams::new(log_n, r, p);
 	scrypt(password.as_bytes(), salt, &scrypt_params, &mut derived_key);
 	let derived_right_bits = &derived_key[0..KEY_LENGTH_AES];
 	let derived_left_bits = &derived_key[KEY_LENGTH_AES..KEY_LENGTH];
-	(derived_right_bits.to_vec(), derived_left_bits.to_vec())
+	Ok((derived_right_bits.to_vec(), derived_left_bits.to_vec()))
 }
 
 pub fn derive_mac(derived_left_bits: &[u8], cipher_text: &[u8]) -> Vec<u8> {
diff --git a/ethstore/src/account/safe_account.rs b/ethstore/src/account/safe_account.rs
index 5dab35251..336e72875 100644
--- a/ethstore/src/account/safe_account.rs
+++ b/ethstore/src/account/safe_account.rs
@@ -113,7 +113,7 @@ impl Crypto {
 
 		let (derived_left_bits, derived_right_bits) = match self.kdf {
 			Kdf::Pbkdf2(ref params) => crypto::derive_key_iterations(password, &params.salt, params.c),
-			Kdf::Scrypt(ref params) => crypto::derive_key_scrypt(password, &params.salt, params.n, params.p, params.r),
+			Kdf::Scrypt(ref params) => try!(crypto::derive_key_scrypt(password, &params.salt, params.n, params.p, params.r)),
 		};
 
 		let mac = crypto::derive_mac(&derived_right_bits, &self.ciphertext).keccak256();