Unlock account with timeout for geth compatibility (#1854) (#1858)

* Unlock account with timeout for geth compatibility

* Fixed test

ethcore/src/account_provider.rs

@@ -17,9 +17,9 @@
 //! Account management.
 
 use std::fmt;
-use std::sync::RwLock;
 use std::collections::HashMap;
-use util::{Address as H160, H256, H520};
+use std::time::{Instant, Duration};
+use util::{Address as H160, H256, H520, Mutex};
 use ethstore::{SecretStore, Error as SSError, SafeAccount, EthStore};
 use ethstore::dir::{KeyDirectory};
 use ethstore::ethkey::{Address as SSAddress, Message as SSMessage, Secret as SSSecret, Random, Generator};
@@ -32,6 +32,8 @@ enum Unlock {
 	/// Account unlocked permantently can always sign message.
 	/// Use with caution.
 	Perm,
+	/// Account unlocked with a timeout
+	Timed((Instant, u32)),
 }
 
 /// Data associated with account.
@@ -126,7 +128,7 @@ impl KeyDirectory for NullDir {
 /// Account management.
 /// Responsible for unlocking accounts.
 pub struct AccountProvider {
-	unlocked: RwLock<HashMap<SSAddress, AccountData>>,
+	unlocked: Mutex<HashMap<SSAddress, AccountData>>,
 	sstore: Box<SecretStore>,
 }
 
@@ -134,7 +136,7 @@ impl AccountProvider {
 	/// Creates new account provider.
 	pub fn new(sstore: Box<SecretStore>) -> Self {
 		AccountProvider {
-			unlocked: RwLock::new(HashMap::new()),
+			unlocked: Mutex::new(HashMap::new()),
 			sstore: sstore,
 		}
 	}
@@ -142,7 +144,7 @@ impl AccountProvider {
 	/// Creates not disk backed provider.
 	pub fn transient_provider() -> Self {
 		AccountProvider {
-			unlocked: RwLock::new(HashMap::new()),
+			unlocked: Mutex::new(HashMap::new()),
 			sstore: Box::new(EthStore::open(Box::new(NullDir)).unwrap())
 		}
 	}
@@ -176,12 +178,10 @@ impl AccountProvider {
 		let _ = try!(self.sstore.sign(&account, &password, &Default::default()));
 
 		// check if account is already unlocked pernamently, if it is, do nothing
-		{
-			let unlocked = self.unlocked.read().unwrap();
-			if let Some(data) = unlocked.get(&account) {
-				if let Unlock::Perm = data.unlock {
-					return Ok(())
-				}
+		let mut unlocked = self.unlocked.lock().unwrap();
+		if let Some(data) = unlocked.get(&account) {
+			if let Unlock::Perm = data.unlock {
+				return Ok(())
 			}
 		}
 
@@ -190,7 +190,6 @@ impl AccountProvider {
 			password: password,
 		};
 
-		let mut unlocked = self.unlocked.write().unwrap();
 		unlocked.insert(account, data);
 		Ok(())
 	}
@@ -205,10 +204,15 @@ impl AccountProvider {
 		self.unlock_account(account, password, Unlock::Temp)
 	}
 
+	/// Unlocks account temporarily with a timeout.
+	pub fn unlock_account_timed<A>(&self, account: A, password: String, duration_ms: u32) -> Result<(), Error> where Address: From<A> {
+		self.unlock_account(account, password, Unlock::Timed((Instant::now(), duration_ms)))
+	}
+
 	/// Checks if given account is unlocked
 	pub fn is_unlocked<A>(&self, account: A) -> bool where Address: From<A> {
 		let account = Address::from(account).into();
-		let unlocked = self.unlocked.read().unwrap();
+		let unlocked = self.unlocked.lock().unwrap();
 		unlocked.get(&account).is_some()
 	}
 
@@ -218,15 +222,20 @@ impl AccountProvider {
 		let message = Message::from(message).into();
 
 		let data = {
-			let unlocked = self.unlocked.read().unwrap();
-			try!(unlocked.get(&account).ok_or(Error::NotUnlocked)).clone()
+			let mut unlocked = self.unlocked.lock().unwrap();
+			let data = try!(unlocked.get(&account).ok_or(Error::NotUnlocked)).clone();
+			if let Unlock::Temp = data.unlock {
+				unlocked.remove(&account).expect("data exists: so key must exist: qed");
+			}
+			if let Unlock::Timed((ref start, ref duration)) = data.unlock {
+				if start.elapsed() > Duration::from_millis(*duration as u64) {
+					unlocked.remove(&account).expect("data exists: so key must exist: qed");
+					return Err(Error::NotUnlocked);
+				}
+			}
+			data
 		};
 
-		if let Unlock::Temp = data.unlock {
-			let mut unlocked = self.unlocked.write().unwrap();
-			unlocked.remove(&account).expect("data exists: so key must exist: qed");
-		}
-
 		let signature = try!(self.sstore.sign(&account, &data.password, &message));
 		Ok(H520(signature.into()))
 	}
@@ -244,6 +253,7 @@ impl AccountProvider {
 mod tests {
 	use super::AccountProvider;
 	use ethstore::ethkey::{Generator, Random};
+	use std::time::Duration;
 
 	#[test]
 	fn unlock_account_temp() {
@@ -269,4 +279,16 @@ mod tests {
 		assert!(ap.sign(kp.address(), [0u8; 32]).is_ok());
 		assert!(ap.sign(kp.address(), [0u8; 32]).is_ok());
 	}
+
+	#[test]
+	fn unlock_account_timer() {
+		let kp = Random.generate().unwrap();
+		let ap = AccountProvider::transient_provider();
+		assert!(ap.insert_account(kp.secret().clone(), "test").is_ok());
+		assert!(ap.unlock_account_timed(kp.address(), "test1".into(), 2000).is_err());
+		assert!(ap.unlock_account_timed(kp.address(), "test".into(), 2000).is_ok());
+		assert!(ap.sign(kp.address(), [0u8; 32]).is_ok());
+		::std::thread::sleep(Duration::from_millis(2000));
+		assert!(ap.sign(kp.address(), [0u8; 32]).is_err());
+	}
 }