SecretStore: remove session on master node (#5545)

* ECDKG protocol prototype

* added test for enc/dec math

* get rid of decryption_session

* added licenses

* fix after merge

* get rid of unused serde dependency

* doc

* decryption session [without commutative enc]

* failed_dec_session

* fixed tests

* added commen

* added more decryption session tests

* helper to localize an issue

* more computations to localize error

* decryption_session::SessionParams

* added tests for EC math to localize problem

* secretstore network transport

* encryption_session_works_over_network

* network errors processing

* connecting to KeyServer

* licenses

* get rid of debug println-s

* fixed secretstore args

* encryption results are stored in KS database

* decryption protocol works over network

* enc/dec Session traits

* fixing warnings

* fix after merge

* on-chain ACL checker proto

* fixed compilation

* fixed compilation

* finally fixed <odd>-of-N-scheme

* temporary commented test

* 1-of-N works in math

* scheme 1-of-N works

* updated AclStorage with real contract ABI

* remove unnecessary unsafety

* fixed grumbles

* wakeup on access denied

* encrypt secretstore messages

* 'shadow' decryption

* fix grumbles

* lost files

* secretstore cli-options

* decryption seccion when ACL check failed on master

* disallow regenerating key for existing document

* removed obsolete TODO

* fix after merge

* switched to tokio_io

* fix after merge

* fix after merge

* fix after merge

* fix after merge

* fix after merge

* fixed test

* fix after merge

* encryption session errors are now fatal

* session timeouts

* autorestart decryption session

* remove sessions on completion

* exclude disconnected nodes from decryption session

* test for enc/dec session over network with 1 node

* remove debug printlns

* fixed 1-of-1 scheme

* drop for KeyServerHttpListener

* Use standard encryption and decryption (as in RPC)

* added some tests

* moved DEFAULT_MAC to ethcrypto

* rpc_secretstore_encrypt_and_decrypt

* serialization with "0x" prefix (RPC compatibility)

* secretstore RPC API

* fix after merge

* fixed typo

* secretstore_shadowDecrypt RPC

* enable secretstore RPCs by default

* fixed test

* SecStore RPCs available without SecStore feature

* fixed grumbles

* lost files

* added password argument to Parity RPCs

* update docs

* remove enc/dec session on master node

* lost file

* pass weak instead of arc

secret_store/src/key_server_cluster/cluster.rs

@@ -16,7 +16,7 @@
 
 use std::io;
 use std::time;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::collections::{BTreeMap, BTreeSet, VecDeque};
 use std::collections::btree_map::Entry;
@@ -27,8 +27,8 @@ use parking_lot::{RwLock, Mutex};
 use tokio_io::IoFuture;
 use tokio_core::reactor::{Handle, Remote, Interval};
 use tokio_core::net::{TcpListener, TcpStream};
-use ethkey::{Secret, KeyPair, Signature, Random, Generator};
+use ethkey::{Public, Secret, KeyPair, Signature, Random, Generator};
-use key_server_cluster::{Error, NodeId, SessionId, AclStorage, KeyStorage};
+use key_server_cluster::{Error, NodeId, SessionId, AclStorage, KeyStorage, DocumentEncryptedKeyShadow};
 use key_server_cluster::message::{self, Message, ClusterMessage, EncryptionMessage, DecryptionMessage};
 use key_server_cluster::decryption_session::{SessionImpl as DecryptionSessionImpl, SessionState as DecryptionSessionState,
 	SessionParams as DecryptionSessionParams, Session as DecryptionSession, DecryptionSessionId};
@@ -236,6 +236,28 @@ pub struct Connection {
 	last_message_time: Mutex<time::Instant>,
 }
 
+/// Encryption session implementation, which removes session from cluster on drop.
+struct EncryptionSessionWrapper {
+	/// Wrapped session.
+	session: Arc<EncryptionSession>,
+	/// Session Id.
+	session_id: SessionId,
+	/// Cluster data reference.
+	cluster: Weak<ClusterData>,
+}
+
+/// Decryption session implementation, which removes session from cluster on drop.
+struct DecryptionSessionWrapper {
+	/// Wrapped session.
+	session: Arc<DecryptionSession>,
+	/// Session Id.
+	session_id: SessionId,
+	/// Session sub id.
+	access_key: Secret,
+	/// Cluster data reference.
+	cluster: Weak<ClusterData>,
+}
+
 impl ClusterCore {
 	pub fn new(handle: Handle, config: ClusterConfiguration) -> Result<Arc<Self>, Error> {
 		let listen_address = make_socket_address(&config.listen_address.0, config.listen_address.1)?;
@@ -1011,9 +1033,9 @@ impl ClusterClient for ClusterClientImpl {
 		connected_nodes.insert(self.data.self_key_pair.public().clone());
 
 		let cluster = Arc::new(ClusterView::new(self.data.clone(), connected_nodes.clone()));
-		let session = self.data.sessions.new_encryption_session(self.data.self_key_pair.public().clone(), session_id, cluster)?;
+		let session = self.data.sessions.new_encryption_session(self.data.self_key_pair.public().clone(), session_id.clone(), cluster)?;
 		session.initialize(threshold, connected_nodes)?;
-		Ok(session)
+		Ok(EncryptionSessionWrapper::new(Arc::downgrade(&self.data), session_id, session))
 	}
 
 	fn new_decryption_session(&self, session_id: SessionId, requestor_signature: Signature, is_shadow_decryption: bool) -> Result<Arc<DecryptionSession>, Error> {
@@ -1022,9 +1044,9 @@ impl ClusterClient for ClusterClientImpl {
 
 		let access_key = Random.generate()?.secret().clone();
 		let cluster = Arc::new(ClusterView::new(self.data.clone(), connected_nodes.clone()));
-		let session = self.data.sessions.new_decryption_session(self.data.self_key_pair.public().clone(), session_id, access_key, cluster)?;
+		let session = self.data.sessions.new_decryption_session(self.data.self_key_pair.public().clone(), session_id, access_key.clone(), cluster)?;
 		session.initialize(requestor_signature, is_shadow_decryption)?;
-		Ok(session)
+		Ok(DecryptionSessionWrapper::new(Arc::downgrade(&self.data), session_id, access_key, session))
 	}
 
 	#[cfg(test)]
@@ -1043,6 +1065,64 @@ impl ClusterClient for ClusterClientImpl {
 	}
 }
 
+impl EncryptionSessionWrapper {
+	pub fn new(cluster: Weak<ClusterData>, session_id: SessionId, session: Arc<EncryptionSession>) -> Arc<Self> {
+		Arc::new(EncryptionSessionWrapper {
+			session: session,
+			session_id: session_id,
+			cluster: cluster,
+		})
+	}
+}
+
+impl EncryptionSession for EncryptionSessionWrapper {
+	fn state(&self) -> EncryptionSessionState {
+		self.session.state()
+	}
+
+	fn wait(&self, timeout: Option<time::Duration>) -> Result<Public, Error> {
+		self.session.wait(timeout)
+	}
+
+	#[cfg(test)]
+	fn joint_public_key(&self) -> Option<Result<Public, Error>> {
+		self.session.joint_public_key()
+	}
+}
+
+impl Drop for EncryptionSessionWrapper {
+	fn drop(&mut self) {
+		if let Some(cluster) = self.cluster.upgrade() {
+			cluster.sessions.remove_encryption_session(&self.session_id);
+		}
+	}
+}
+
+impl DecryptionSessionWrapper {
+	pub fn new(cluster: Weak<ClusterData>, session_id: SessionId, access_key: Secret, session: Arc<DecryptionSession>) -> Arc<Self> {
+		Arc::new(DecryptionSessionWrapper {
+			session: session,
+			session_id: session_id,
+			access_key: access_key,
+			cluster: cluster,
+		})
+	}
+}
+
+impl DecryptionSession for DecryptionSessionWrapper {
+	fn wait(&self) -> Result<DocumentEncryptedKeyShadow, Error> {
+		self.session.wait()
+	}
+}
+
+impl Drop for DecryptionSessionWrapper {
+	fn drop(&mut self) {
+		if let Some(cluster) = self.cluster.upgrade() {
+			cluster.sessions.remove_decryption_session(&self.session_id, &self.access_key);
+		}
+	}
+}
+
 fn make_socket_address(address: &str, port: u16) -> Result<SocketAddr, Error> {
 	let ip_address: IpAddr = address.parse().map_err(|_| Error::InvalidNodeAddress)?;
 	Ok(SocketAddr::new(ip_address, port))