// Copyright 2015-2017 Parity Technologies (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

use std::collections::{BTreeSet, BTreeMap, VecDeque};
use std::fmt::{Debug, Formatter, Error as FmtError};
use std::sync::Arc;
use parking_lot::Mutex;
use ethkey::{Public, Secret};
use key_server_cluster::{Error, NodeId, SessionId};
use key_server_cluster::math;
use key_server_cluster::cluster::Cluster;
use key_server_cluster::message::{Message, InitializeSession, ConfirmInitialization, CompleteInitialization,
	KeysDissemination, Complaint, ComplaintResponse, PublicKeyShare};

/// Encryption (distributed key generation) session.
/// Based on "ECDKG: A Distributed Key Generation Protocol Based on Elliptic Curve Discrete Logarithm" paper:
/// http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.124.4128&rep=rep1&type=pdf
/// Brief overview:
/// 1) initialization: master node (which has received request for generating joint public + secret) initializes the session on all other nodes
/// 2) key dissemination (KD): all nodes are generating secret + public values and send these to appropriate nodes
/// 3) key verification (KV): all nodes are checking values, received for other nodes and complaining if keys are wrong
/// 4) key check phase (KC): nodes are processing complaints, received from another nodes
/// 5) key generation phase (KG): nodes are exchanging with information, enough to generate joint public key
pub struct Session {
	/// Unique session id.
	id: SessionId,
	/// Public identifier of this node.
	self_node_id: NodeId,
	/// Cluster which allows this node to send messages to other nodes in the cluster.
	cluster: Arc<Cluster>,
	/// Mutable session data.
	data: Mutex<SessionData>,
}

#[derive(Debug)]
/// Mutable data of encryption (distributed key generation) session.
struct SessionData {
	/// Current state of the session.
	state: SessionState,

	// === Values, filled when session initialization just starts ===
	/// Reference to the node, which has started this session.
	master: Option<NodeId>,

	// === Values, filled when session initialization is completed ===
	/// Threshold value for this DKG. Only `threshold + 1` will be able to collectively recreate joint secret,
	/// and thus - decrypt message, encrypted with joint public.
	threshold: Option<usize>,
	/// Random point, jointly generated by every node in the cluster.
	derived_point: Option<Public>,
	/// Nodes-specific data.
	nodes: BTreeMap<NodeId, NodeData>,

	// === Values, filled during KD phase ===
	/// Value of polynom1[0], generated by this node.
	secret_coeff: Option<Secret>,

	// === Values, filled during KG phase ===
	/// Secret share, which this node holds. Persistent + private.
	secret_share: Option<Secret>,

	/// === Values, filled when DKG session is completed successfully ===
	/// Jointly generated public key, which can be used to encrypt secret. Public.
	joint_public: Option<Public>,
}

#[derive(Debug, Clone)]
/// Mutable node-specific data.
struct NodeData {
	/// Random unique scalar. Persistent.
	pub id_number: Secret,

	// === Values, filled during KD phase ===
	/// Secret value1, which has been sent to this node.
	pub secret1_sent: Option<Secret>,
	/// Secret value2, which has been sent to this node.
	pub secret2_sent: Option<Secret>,
	/// Secret value1, which has been received from this node.
	pub secret1: Option<Secret>,
	/// Secret value2, which has been received from this node.
	pub secret2: Option<Secret>,
	/// Public values, which have been received from this node.
	pub publics: Option<Vec<Public>>,

	/// === Values, filled during KC phase ===
	/// Nodes, complaining against this node.
	pub complaints: BTreeSet<NodeId>,

	/// === Values, filled during KG phase ===
	/// Public share, which has been received from this node.
	pub public_share: Option<Public>,
}

#[derive(Debug, Clone, PartialEq)]
/// Schedule for visiting other nodes of cluster.
pub struct EveryOtherNodeVisitor {
	/// Already visited nodes.
	visited: BTreeSet<NodeId>,
	/// Not yet visited nodes.
	unvisited: VecDeque<NodeId>,
	/// Nodes, which are currently visited.
	in_progress: BTreeSet<NodeId>,
}

#[derive(Debug, Clone, PartialEq)]
/// Encryption (distributed key generation) session state.
pub enum SessionState {
	// === Initialization states ===
	/// Every node starts in this state.
	WaitingForInitialization,
	/// Master node asks every other node to confirm initialization.
	/// Derived point is generated by all nodes in the cluster.
	WaitingForInitializationConfirm(EveryOtherNodeVisitor),
	/// Slave nodes are in this state until initialization completion is reported by master node.
	WaitingForInitializationComplete,

	// === KD phase states ===
	/// Node is waiting for generated keys from every other node.
	WaitingForKeysDissemination,

	// === KC phase states ===
	/// Keys check currently occurs.
	KeyCheck,

	// === KG phase states ===
	/// Node is waiting for joint public key share to be received from every other node.
	WaitingForPublicKeyShare,

	// === Final states of the session ===
	/// Joint public key generation is completed.
	Finished,
	/// Joint public key generation is failed.
	Failed,
}

impl Session {
	/// Create new encryption session.
	pub fn new(id: SessionId, self_node_id: Public, cluster: Arc<Cluster>) -> Self {
		Session {
			id: id,
			self_node_id: self_node_id,
			cluster: cluster,
			data: Mutex::new(SessionData {
				state: SessionState::WaitingForInitialization,
				master: None,
				threshold: None,
				derived_point: None,
				nodes: BTreeMap::new(),
				secret_coeff: None,
				secret_share: None,
				joint_public: None,
			}),
		}
	}

	/// Get this node Id.
	pub fn node(&self) -> &NodeId {
		&self.self_node_id
	}

	/// Get current session state.
	pub fn state(&self) -> SessionState {
		self.data.lock().state.clone()
	}

	/// Get joint public key.
	pub fn joint_public_key(&self) -> Option<Public> {
		self.data.lock().joint_public.clone()
	}

	#[cfg(test)]
	/// Get derived point.
	pub fn derived_point(&self) -> Option<Public> {
		self.data.lock().derived_point.clone()
	}

	#[cfg(test)]
	/// Get qualified nodes.
	pub fn qualified_nodes(&self) -> BTreeSet<NodeId> {
		self.data.lock().nodes.keys().cloned().collect()
	}

	/// Start new session initialization. This must be called on master node.
	pub fn initialize(&self, threshold: usize, nodes: BTreeSet<NodeId>) -> Result<(), Error> {
		check_cluster_nodes(self.node(), &nodes)?;
		check_threshold(threshold, &nodes)?;

		let mut data = self.data.lock();

		// check state
		if data.state != SessionState::WaitingForInitialization {
			return Err(Error::InvalidStateForRequest);
		}

		// update state
		data.master = Some(self.node().clone());
		data.threshold = Some(threshold);
		for node_id in &nodes {
			// generate node identification parameter
			let node_id_number = math::generate_random_scalar()?;
			data.nodes.insert(node_id.clone(), NodeData::with_id_number(node_id_number));
		}

		let mut visit_policy = EveryOtherNodeVisitor::new(self.node(), data.nodes.keys().cloned());
		let next_node = visit_policy.next_node().expect("at least two nodes are in cluster");
		data.state = SessionState::WaitingForInitializationConfirm(visit_policy);

		// start initialization
		let derived_point = math::generate_random_point()?;
		self.cluster.send(&next_node, Message::InitializeSession(InitializeSession {
				session: self.id.clone(),
				derived_point: derived_point,
			}))
	}

	/// When session initialization message is received.
	pub fn on_initialize_session(&self, sender: NodeId, mut message: InitializeSession) -> Result<(), Error> {
		debug_assert!(self.id == message.session);
		debug_assert!(&sender != self.node());

		let mut data = self.data.lock();

		// check state
		if data.state != SessionState::WaitingForInitialization {
			return Err(Error::InvalidStateForRequest);
		}

		// update derived point with random scalar
		math::update_random_point(&mut message.derived_point)?;

		// send confirmation back to master node
		self.cluster.send(&sender, Message::ConfirmInitialization(ConfirmInitialization {
			session: self.id.clone(),
			derived_point: message.derived_point,
		}))?;

		// update state
		data.master = Some(sender);
		data.state = SessionState::WaitingForInitializationComplete;

		Ok(())
	}

	/// When session initialization confirmation message is reeived.
	pub fn on_confirm_initialization(&self, sender: NodeId, mut message: ConfirmInitialization) -> Result<(), Error> {
		debug_assert!(self.id == message.session);
		debug_assert!(&sender != self.node());

		let mut data = self.data.lock();
		debug_assert!(data.nodes.contains_key(&sender));

		// check state && select new node to be initialized
		let next_receiver = match data.state {
			SessionState::WaitingForInitializationConfirm(ref mut visit_policy) => {
				if !visit_policy.mark_visited(&sender) {
					return Err(Error::InvalidStateForRequest);
				}

				visit_policy.next_node()
			},
			_ => return Err(Error::InvalidStateForRequest),
		};

		// proceed message
		match next_receiver {
			Some(next_receiver) => {
				return self.cluster.send(&next_receiver, Message::InitializeSession(InitializeSession {
						session: self.id.clone(),
						derived_point: message.derived_point,
					}));
			},
			None => {
				// update point once again to make sure that derived point is not generated by last node
				math::update_random_point(&mut message.derived_point)?;

				// remember derived point
				data.derived_point = Some(message.derived_point.clone());

				// broadcast derived point && other session paraeters to every other node
				self.cluster.broadcast(Message::CompleteInitialization(CompleteInitialization {
					session: self.id.clone(),
					nodes: data.nodes.iter().map(|(id, data)| (id.clone(), data.id_number.clone())).collect(),
					threshold: data.threshold.expect("threshold is filled in initialization phase; KD phase follows initialization phase; qed"),
					derived_point: message.derived_point,
				}))?;
			},
		}

		// now it is time for keys dissemination (KD) phase
		drop(data);
		self.disseminate_keys()
	}

	/// When session initialization completion message is received.
	pub fn on_complete_initialization(&self, sender: NodeId, message: CompleteInitialization) -> Result<(), Error> {
		debug_assert!(self.id == message.session);
		debug_assert!(&sender != self.node());

		// check message
		let nodes_ids = message.nodes.keys().cloned().collect();
		check_cluster_nodes(self.node(), &nodes_ids)?;
		check_threshold(message.threshold, &nodes_ids)?;

		let mut data = self.data.lock();

		// check state
		if data.state != SessionState::WaitingForInitializationComplete {
			return Err(Error::InvalidStateForRequest);
		}
		if data.master != Some(sender) {
			return Err(Error::InvalidMessage);
		}

		// remember passed data
		data.threshold = Some(message.threshold);
		data.derived_point = Some(message.derived_point);
		data.nodes = message.nodes.into_iter().map(|(id, number)| (id, NodeData::with_id_number(number))).collect();

		// now it is time for keys dissemination (KD) phase
		drop(data);
		self.disseminate_keys()
	}

	/// When keys dissemination message is received.
	pub fn on_keys_dissemination(&self, sender: NodeId, message: KeysDissemination) -> Result<(), Error> {
		debug_assert!(self.id == message.session);
		debug_assert!(&sender != self.node());

		let mut data = self.data.lock();
		debug_assert!(data.nodes.contains_key(&sender));

		// check state
		if data.state != SessionState::WaitingForKeysDissemination {
			return Err(Error::InvalidStateForRequest);
		}

		// check message
		let threshold = data.threshold.expect("threshold is filled in initialization phase; KD phase follows initialization phase; qed");
		if message.publics.len() != threshold + 1 {
			return Err(Error::InvalidMessage);
		}

		// update node data
		{
			let node_data = data.nodes.get_mut(&sender).ok_or(Error::InvalidMessage)?;
			if node_data.secret1.is_some() || node_data.secret2.is_some() || node_data.publics.is_some() {
				return Err(Error::InvalidStateForRequest);
			}

			node_data.secret1 = Some(message.secret1);
			node_data.secret2 = Some(message.secret2);
			node_data.publics = Some(message.publics);
		}

		// check if we have received keys from every other node
		if data.nodes.iter().any(|(node_id, node_data)| node_id != self.node() && (node_data.publics.is_none() || node_data.secret1.is_none() || node_data.secret2.is_none())) {
			return Ok(())
		}

		// key verification (KV) phase: check that other nodes have passed correct secrets
		let derived_point = data.derived_point.clone().expect("derived point generated on initialization phase; KV phase follows initialization phase; qed");
		let number_id = data.nodes[self.node()].id_number.clone();
		for (node_id, node_data) in data.nodes.iter_mut().filter(|&(node_id, _)| node_id != self.node()) {
			let secret1 = node_data.secret1.as_ref().expect("keys received on KD phase; KV phase follows KD phase; qed");
			let secret2 = node_data.secret2.as_ref().expect("keys received on KD phase; KV phase follows KD phase; qed");
			let publics = node_data.publics.as_ref().expect("keys received on KD phase; KV phase follows KD phase; qed");
			let is_key_verification_ok = math::keys_verification(threshold, &derived_point, &number_id,
				secret1, secret2, publics)?;

			if !is_key_verification_ok {
				node_data.complaints.insert(self.node().clone());
				self.cluster.broadcast(Message::Complaint(Complaint {
					session: self.id.clone(),
					against: node_id.clone(),
				}))?;
			}
		}

		// update state
		data.state = SessionState::KeyCheck;

		Ok(())
	}

	/// When complaint is received.
	pub fn on_complaint(&self, sender: NodeId, message: Complaint) -> Result<(), Error> {
		debug_assert!(self.id == message.session);
		debug_assert!(&sender != self.node());

		let mut data = self.data.lock();
		debug_assert!(data.nodes.contains_key(&sender));

		// check state
		if data.state != SessionState::KeyCheck && data.state != SessionState::WaitingForKeysDissemination {
			// message can be received after timeout is passed
			// => encryption is already in progress
			// => we can not treat it as an error
			return Ok(());
		}

		// respond to complaint
		if &message.against == self.node() {
			let secret1_sent = data.nodes[&sender].secret1_sent.clone().expect("secrets were sent on KD phase; KC phase follows KD phase; qed");
			let secret2_sent = data.nodes[&sender].secret2_sent.clone().expect("secrets were sent on KD phase; KC phase follows KD phase; qed");

			// someone is complaining against us => let's respond
			return self.cluster.broadcast(Message::ComplaintResponse(ComplaintResponse {
				session: self.id.clone(),
				secret1: secret1_sent,
				secret2: secret2_sent,
			}));
		}

		// someone is complaining against someone else => let's remember this
		let threshold = data.threshold.expect("threshold is filled in initialization phase; KD phase follows initialization phase; qed");
		let is_critical_complaints_num = {
			let node = data.nodes.get_mut(&message.against).ok_or(Error::InvalidMessage)?;
			node.complaints.insert(sender);
			node.complaints.len() >= threshold + 1
		};

		if is_critical_complaints_num {
			// too many complaints => exclude from session
			Session::disqualify_node(&message.against, &*self.cluster, &mut *data);
		}

		Ok(())
	}

	/// When complaint response is received
	pub fn on_complaint_response(&self, sender: NodeId, message: ComplaintResponse) -> Result<(), Error> {
		debug_assert!(self.id == message.session);
		debug_assert!(&sender != self.node());

		let mut data = self.data.lock();
		debug_assert!(data.nodes.contains_key(&sender));

		// check state
		if data.state != SessionState::KeyCheck {
			// in theory this message can be received before KeyCheck (in WaitingForKeysDissemination state)
			// => but the fact that someone is complaining about keys means that sender has already sent its keys
			// => complaint response can't reach us before keys, as the cluster guarantees that messages are received in FIFO order

			// message can be received after timeout is passed
			// => encryption is already in progress
			// => we can not treat it as an error
			return Ok(());
		}

		// check keys again
		let is_key_verification_ok = {
			let threshold = data.threshold.expect("threshold is filled in initialization phase; KD phase follows initialization phase; qed");
			let derived_point = data.derived_point.as_ref().expect("derived point generated on initialization phase; KV phase follows initialization phase; qed");
			let number_id = &data.nodes[self.node()].id_number;
			let node_data = data.nodes.get(&sender).expect("cluster guarantees to deliver messages from qualified nodes only; qed");
			let publics = node_data.publics.as_ref().expect("keys received on KD phase; KV phase follows KD phase; qed");
			math::keys_verification(threshold, derived_point, number_id, &message.secret1, &message.secret2, publics)?
		};

		if !is_key_verification_ok {
			Session::disqualify_node(&sender, &*self.cluster, &mut *data);
		} else {
			let node_data = data.nodes.get_mut(&sender).expect("cluster guarantees to deliver messages from qualified nodes only; qed");
			node_data.secret1 = Some(message.secret1);
			node_data.secret2 = Some(message.secret2);
			node_data.complaints.remove(self.node());
		}

		Ok(())
	}

	/// When KC-phase timeout is expired, it is time to start KG phase.
	pub fn start_key_generation_phase(&self) -> Result<(), Error> {
		let mut data = self.data.lock();

		if data.state != SessionState::KeyCheck {
			return Err(Error::InvalidStateForRequest);
		}

		// calculate public share
		let self_public_share = {
			let self_secret_coeff = data.secret_coeff.as_ref().expect("secret_coeff is generated on KD phase; KG phase follows KD phase; qed");
			math::compute_public_share(self_secret_coeff)?
		};

		// calculate self secret + public shares
		let self_secret_share = {
			let secret_values_iter = data.nodes.values()
				.map(|n| n.secret1.as_ref().expect("keys received on KD phase; KG phase follows KD phase; qed"));
			math::compute_secret_share(secret_values_iter)?
		};

		// update state
		data.state = SessionState::WaitingForPublicKeyShare;
		data.secret_share = Some(self_secret_share);
		let self_node = data.nodes.get_mut(self.node()).expect("node is always qualified by himself; qed");
		self_node.public_share = Some(self_public_share.clone());

		// broadcast self public key share
		self.cluster.broadcast(Message::PublicKeyShare(PublicKeyShare {
			session: self.id.clone(),
			public_share: self_public_share,
		}))
	}

	/// When public key share is received.
	pub fn on_public_key_share(&self, sender: NodeId, message: PublicKeyShare) -> Result<(), Error> {
		let mut data = self.data.lock();

		// check state
		if data.state != SessionState::WaitingForPublicKeyShare {
			return Err(Error::InvalidStateForRequest);
		}

		// update node data with received public share
		{
			let node_data = &mut data.nodes.get_mut(&sender).ok_or(Error::InvalidMessage)?;
				if node_data.public_share.is_some() {
				return Err(Error::InvalidMessage);
			}

			node_data.public_share = Some(message.public_share);
		}

		// if there's also nodes, which has not sent us their public shares - do nothing
		if data.nodes.iter().any(|(node_id, node_data)| node_id != self.node() && node_data.public_share.is_none()) {
			return Ok(());
		}

		// else - calculate joint public key && finish session
		data.joint_public = {
			let public_shares = data.nodes.values().map(|n| n.public_share.as_ref().expect("keys received on KD phase; KG phase follows KD phase; qed"));
			Some(math::compute_joint_public(public_shares)?)
		};
		data.state = SessionState::Finished;

		Ok(())
	}

	/// Keys dissemination (KD) phase
	fn disseminate_keys(&self) -> Result<(), Error> {
		let mut data = self.data.lock();

		// pick 2t + 2 random numbers as polynomial coefficients for 2 polynoms
		let threshold = data.threshold.expect("threshold is filled on initialization phase; KD phase follows initialization phase; qed");
		let polynom1 = math::generate_random_polynom(threshold)?;
		let polynom2 = math::generate_random_polynom(threshold)?;
		data.secret_coeff = Some(polynom1[0].clone());
		
		// compute t+1 public values
		let publics = math::public_values_generation(threshold,
			data.derived_point.as_ref().expect("keys dissemination occurs after derived point is agreed; qed"),
			&polynom1,
			&polynom2)?;

		// compute secret values for every other node
		for (node, node_data) in data.nodes.iter_mut() {
			let secret1 = math::compute_polynom(&polynom1, &node_data.id_number)?;
			let secret2 = math::compute_polynom(&polynom2, &node_data.id_number)?;

			// send a message containing secret1 && secret2 to other node
			if node != self.node() {
				node_data.secret1_sent = Some(secret1.clone());
				node_data.secret2_sent = Some(secret2.clone());

				self.cluster.send(&node, Message::KeysDissemination(KeysDissemination {
					session: self.id.clone(),
					secret1: secret1,
					secret2: secret2,
					publics: publics.clone(),
				}))?;
			} else {
				node_data.secret1 = Some(secret1);
				node_data.secret2 = Some(secret2);
				node_data.publics = Some(publics.clone());
			}
		}

		// update state
		data.state = SessionState::WaitingForKeysDissemination;

		Ok(())
	}

	/// Disqualify node
	fn disqualify_node(node: &NodeId, cluster: &Cluster, data: &mut SessionData) {
		let threshold = data.threshold
			.expect("threshold is filled on initialization phase; node can only be disqualified during KC phase; KC phase follows initialization phase; qed");
		
		// blacklist node
		cluster.blacklist(&node);
		// too many complaints => exclude from session
		data.nodes.remove(&node);
		// check if secret still can be reconstructed
		if data.nodes.len() < threshold + 1 {
			// not enough nodes => session is failed
			data.state = SessionState::Failed;
		}
	}
}

impl EveryOtherNodeVisitor {
	pub fn new<I>(self_id: &NodeId, nodes: I) -> Self where I: Iterator<Item=NodeId> {
		EveryOtherNodeVisitor {
			visited: BTreeSet::new(),
			unvisited: nodes.filter(|n| n != self_id).collect(),
			in_progress: BTreeSet::new(),
		}
	}

	pub fn next_node(&mut self) -> Option<NodeId> {
		let next_node = self.unvisited.pop_front();
		if let Some(ref next_node) = next_node {
			self.in_progress.insert(next_node.clone());
		}
		next_node
	}

	pub fn mark_visited(&mut self, node: &NodeId) -> bool {
		if !self.in_progress.remove(node) {
			return false;
		}
		self.visited.insert(node.clone())
	}
}

impl NodeData {
	fn with_id_number(node_id_number: Secret) -> Self {
		NodeData {
			id_number: node_id_number,
			complaints: BTreeSet::new(),
			secret1_sent: None,
			secret2_sent: None,
			secret1: None,
			secret2: None,
			publics: None,
			public_share: None,
		}
	}
}

impl Debug for Session {
	fn fmt(&self, f: &mut Formatter) -> Result<(), FmtError> {
		write!(f, "Encryption session {} on {}", self.id, self.self_node_id)
	}
}

pub fn check_cluster_nodes(self_node_id: &NodeId, nodes: &BTreeSet<NodeId>) -> Result<(), Error> {
	// at least two nodes must be in cluster
	if nodes.len() < 2 {
		return Err(Error::InvalidNodesCount);
	}
	// this node must be a part of cluster
	if !nodes.contains(self_node_id) {
		return Err(Error::InvalidNodesConfiguration);
	}

	Ok(())
}

pub fn check_threshold(threshold: usize, nodes: &BTreeSet<NodeId>) -> Result<(), Error> {
	// at least threshold + 1 nodes are required to collectively decrypt message
	if threshold >= nodes.len() {
		return Err(Error::InvalidThreshold);
	}

	Ok(())
}

#[cfg(test)]
mod tests {
	use std::sync::Arc;
	use std::collections::{BTreeSet, BTreeMap};
	use ethkey::{Random, Generator};
	use key_server_cluster::{NodeId, SessionId, Error};
	use key_server_cluster::message::{self, Message};
	use key_server_cluster::cluster::tests::DummyCluster;
	use key_server_cluster::encryption_session::{Session, SessionState};
	use key_server_cluster::math;
	use key_server_cluster::math::tests::do_encryption_and_decryption;

	#[derive(Debug)]
	struct Node {
		pub cluster: Arc<DummyCluster>,
		pub session: Session,
	}

	#[derive(Debug)]
	struct MessageLoop {
		pub session_id: SessionId,
		pub nodes: BTreeMap<NodeId, Node>,
	}

	impl MessageLoop {
		pub fn new(nodes_num: usize) -> Self {
			let mut nodes = BTreeMap::new();
			let session_id = SessionId::default();
			for _ in 0..nodes_num {
				let key_pair = Random.generate().unwrap();
				let node_id = key_pair.public().clone();
				let cluster = Arc::new(DummyCluster::new(node_id.clone()));
				let session = Session::new(session_id.clone(), node_id.clone(), cluster.clone());
				nodes.insert(node_id, Node { cluster: cluster, session: session });
			}

			let nodes_ids: Vec<_> = nodes.keys().cloned().collect();
			for node in nodes.values() {
				for node_id in &nodes_ids {
					node.cluster.add_node(node_id.clone());
				}
			}

			MessageLoop {
				session_id: session_id,
				nodes: nodes,
			}
		}

		pub fn master(&self) -> &Session {
			&self.nodes.values().nth(0).unwrap().session
		}

		pub fn first_slave(&self) -> &Session {
			&self.nodes.values().nth(1).unwrap().session
		}

		pub fn second_slave(&self) -> &Session {
			&self.nodes.values().nth(2).unwrap().session
		}

		pub fn third_slave(&self) -> &Session {
			&self.nodes.values().nth(3).unwrap().session
		}

		pub fn take_message(&mut self) -> Option<(NodeId, NodeId, Message)> {
			self.nodes.values()
				.filter_map(|n| n.cluster.take_message().map(|m| (n.session.node().clone(), m.0, m.1)))
				.nth(0)
		}

		pub fn process_message(&mut self, msg: (NodeId, NodeId, Message)) -> Result<(), Error> {
			match msg.2 {
				Message::InitializeSession(message) => self.nodes[&msg.1].session.on_initialize_session(msg.0, message),
				Message::ConfirmInitialization(message) => self.nodes[&msg.1].session.on_confirm_initialization(msg.0, message),
				Message::CompleteInitialization(message) => self.nodes[&msg.1].session.on_complete_initialization(msg.0, message),
				Message::KeysDissemination(message) => self.nodes[&msg.1].session.on_keys_dissemination(msg.0, message),
				Message::Complaint(message) => self.nodes[&msg.1].session.on_complaint(msg.0, message),
				Message::ComplaintResponse(message) => self.nodes[&msg.1].session.on_complaint_response(msg.0, message),
				Message::PublicKeyShare(message) => self.nodes[&msg.1].session.on_public_key_share(msg.0, message),
				_ => panic!("unexpected"),
			}
		}

		pub fn take_and_process_message(&mut self) -> Result<(), Error> {
			let msg = self.take_message().unwrap();
			self.process_message(msg)
		}

		pub fn take_and_process_all_messages(&mut self) -> Result<(), Error> {
			while let Some(msg) = self.take_message() {
				self.process_message(msg)?;
			}
			Ok(())
		}
	}

	fn make_simple_cluster(threshold: usize, num_nodes: usize) -> Result<(SessionId, NodeId, NodeId, MessageLoop), Error> {
		let l = MessageLoop::new(num_nodes);
		l.master().initialize(threshold, l.nodes.keys().cloned().collect())?;

		let session_id = l.session_id.clone();
		let master_id = l.master().node().clone();
		let slave_id = l.first_slave().node().clone();
		Ok((session_id, master_id, slave_id, l))
	}

	#[test]
	fn fails_to_initialize_in_cluster_of_single_node() {
		assert_eq!(make_simple_cluster(0, 1).unwrap_err(), Error::InvalidNodesCount);
	}

	#[test]
	fn fails_to_initialize_if_not_a_part_of_cluster() {
		let node_id = math::generate_random_point().unwrap();
		let cluster = Arc::new(DummyCluster::new(node_id.clone()));
		let session = Session::new(SessionId::default(), node_id.clone(), cluster);
		let cluster_nodes: BTreeSet<_> = (0..2).map(|_| math::generate_random_point().unwrap()).collect();
		assert_eq!(session.initialize(0, cluster_nodes).unwrap_err(), Error::InvalidNodesConfiguration);
	}

	#[test]
	fn fails_to_initialize_if_threshold_is_wrong() {
		assert_eq!(make_simple_cluster(2, 2).unwrap_err(), Error::InvalidThreshold);
	}

	#[test]
	fn fails_to_initialize_when_already_initialized() {
		let (_, _, _, l) = make_simple_cluster(0, 2).unwrap();
		assert_eq!(l.master().initialize(0, l.nodes.keys().cloned().collect()).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn fails_to_accept_initialization_when_already_initialized() {
		let (sid, m, _, mut l) = make_simple_cluster(0, 2).unwrap();
		l.take_and_process_message().unwrap();
		assert_eq!(l.first_slave().on_initialize_session(m, message::InitializeSession {
			session: sid,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn slave_updates_derived_point_on_initialization() {
		let (_, _, _, mut l) = make_simple_cluster(0, 2).unwrap();
		let passed_point = match l.take_message().unwrap() {
			(f, t, Message::InitializeSession(message)) => {
				let point = message.derived_point.clone();
				l.process_message((f, t, Message::InitializeSession(message))).unwrap();
				point
			},
			_ => panic!("unexpected"),
		};

		match l.take_message().unwrap() {
			(_, _, Message::ConfirmInitialization(message)) => assert!(passed_point != message.derived_point),
			_ => panic!("unexpected"),
		}
	}

	#[test]
	fn fails_to_accept_initialization_confirmation_if_already_accepted_from_the_same_node() {
		let (sid, _, s, mut l) = make_simple_cluster(0, 3).unwrap();
		l.take_and_process_message().unwrap();
		l.take_and_process_message().unwrap();
		l.take_and_process_message().unwrap();
		assert_eq!(l.master().on_confirm_initialization(s, message::ConfirmInitialization {
			session: sid,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn fails_to_accept_initialization_confirmation_if_initialization_already_completed() {
		let (sid, _, s, mut l) = make_simple_cluster(0, 2).unwrap();
		l.take_and_process_message().unwrap();
		l.take_and_process_message().unwrap();
		assert_eq!(l.master().on_confirm_initialization(s, message::ConfirmInitialization {
			session: sid,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn master_updates_derived_point_on_initialization_completion() {
		let (_, _, _, mut l) = make_simple_cluster(0, 2).unwrap();
		l.take_and_process_message().unwrap();
		let passed_point = match l.take_message().unwrap() {
			(f, t, Message::ConfirmInitialization(message)) => {
				let point = message.derived_point.clone();
				l.process_message((f, t, Message::ConfirmInitialization(message))).unwrap();
				point
			},
			_ => panic!("unexpected"),
		};

		assert!(passed_point != l.master().derived_point().unwrap());
	}

	#[test]
	fn fails_to_complete_initialization_in_cluster_of_single_node() {
		let (sid, m, s, l) = make_simple_cluster(0, 2).unwrap();
		let mut nodes = BTreeMap::new();
		nodes.insert(s, math::generate_random_scalar().unwrap());
		assert_eq!(l.first_slave().on_complete_initialization(m, message::CompleteInitialization {
			session: sid,
			nodes: nodes,
			threshold: 0,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidNodesCount);
	}

	#[test]
	fn fails_to_complete_initialization_if_not_a_part_of_cluster() {
		let (sid, m, _, l) = make_simple_cluster(0, 2).unwrap();
		let mut nodes = BTreeMap::new();
		nodes.insert(m, math::generate_random_scalar().unwrap());
		nodes.insert(math::generate_random_point().unwrap(), math::generate_random_scalar().unwrap());
		assert_eq!(l.first_slave().on_complete_initialization(m, message::CompleteInitialization {
			session: sid,
			nodes: nodes,
			threshold: 0,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidNodesConfiguration);
	}

	#[test]
	fn fails_to_complete_initialization_if_threshold_is_wrong() {
		let (sid, m, s, l) = make_simple_cluster(0, 2).unwrap();
		let mut nodes = BTreeMap::new();
		nodes.insert(m, math::generate_random_scalar().unwrap());
		nodes.insert(s, math::generate_random_scalar().unwrap());
		assert_eq!(l.first_slave().on_complete_initialization(m, message::CompleteInitialization {
			session: sid,
			nodes: nodes,
			threshold: 2,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidThreshold);
	}

	#[test]
	fn fails_to_complete_initialization_if_not_waiting_for_it() {
		let (sid, m, s, l) = make_simple_cluster(0, 2).unwrap();
		let mut nodes = BTreeMap::new();
		nodes.insert(m, math::generate_random_scalar().unwrap());
		nodes.insert(s, math::generate_random_scalar().unwrap());
		assert_eq!(l.first_slave().on_complete_initialization(m, message::CompleteInitialization {
			session: sid,
			nodes: nodes,
			threshold: 0,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn fails_to_complete_initialization_from_non_master_node() {
		let (sid, m, s, mut l) = make_simple_cluster(0, 3).unwrap();
		l.take_and_process_message().unwrap();
		l.take_and_process_message().unwrap();
		l.take_and_process_message().unwrap();
		l.take_and_process_message().unwrap();
		let mut nodes = BTreeMap::new();
		nodes.insert(m, math::generate_random_scalar().unwrap());
		nodes.insert(s, math::generate_random_scalar().unwrap());
		nodes.insert(l.second_slave().node().clone(), math::generate_random_scalar().unwrap());
		assert_eq!(l.first_slave().on_complete_initialization(l.second_slave().node().clone(), message::CompleteInitialization {
			session: sid,
			nodes: nodes,
			threshold: 0,
			derived_point: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidMessage);
	}

	#[test]
	fn fails_to_accept_keys_dissemination_if_not_waiting_for_it() {
		let (sid, _, s, l) = make_simple_cluster(0, 2).unwrap();
		assert_eq!(l.master().on_keys_dissemination(s, message::KeysDissemination {
			session: sid,
			secret1: math::generate_random_scalar().unwrap(),
			secret2: math::generate_random_scalar().unwrap(),
			publics: vec![math::generate_random_point().unwrap()],
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn fails_to_accept_keys_dissemination_if_wrong_number_of_publics_passed() {
		let (sid, m, _, mut l) = make_simple_cluster(0, 3).unwrap();
		l.take_and_process_message().unwrap(); // m -> s1: InitializeSession
		l.take_and_process_message().unwrap(); // m -> s2: InitializeSession
		l.take_and_process_message().unwrap(); // s1 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // s2 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // m -> s1: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s2: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s1: KeysDissemination
		assert_eq!(l.first_slave().on_keys_dissemination(m, message::KeysDissemination {
			session: sid,
			secret1: math::generate_random_scalar().unwrap(),
			secret2: math::generate_random_scalar().unwrap(),
			publics: vec![math::generate_random_point().unwrap(), math::generate_random_point().unwrap()],
		}).unwrap_err(), Error::InvalidMessage);
	}

	#[test]
	fn fails_to_accept_keys_dissemination_second_time_from_the_same_node() {
		let (sid, m, _, mut l) = make_simple_cluster(0, 3).unwrap();
		l.take_and_process_message().unwrap(); // m -> s1: InitializeSession
		l.take_and_process_message().unwrap(); // m -> s2: InitializeSession
		l.take_and_process_message().unwrap(); // s1 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // s2 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // m -> s1: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s2: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s1: KeysDissemination
		assert_eq!(l.first_slave().on_keys_dissemination(m, message::KeysDissemination {
			session: sid,
			secret1: math::generate_random_scalar().unwrap(),
			secret2: math::generate_random_scalar().unwrap(),
			publics: vec![math::generate_random_point().unwrap()],
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn defends_if_receives_complain_on_himself() {
		let (sid, m, s, mut l) = make_simple_cluster(1, 3).unwrap();
		l.take_and_process_all_messages().unwrap();
		l.master().on_complaint(s, message::Complaint {
			session: sid,
			against: m,
		}).unwrap();
		match l.take_message().unwrap() {
			(_, _, Message::ComplaintResponse(_)) => (),
			_ => panic!("unexpected"),
		}
	}

	#[test]
	fn node_is_disqualified_if_enough_complaints_received() {
		let (sid, _, s, mut l) = make_simple_cluster(1, 4).unwrap();
		l.take_and_process_all_messages().unwrap();
		l.master().on_complaint(l.second_slave().node().clone(), message::Complaint {
			session: sid,
			against: s.clone(),
		}).unwrap();
		l.master().on_complaint(l.third_slave().node().clone(), message::Complaint {
			session: sid,
			against: s,
		}).unwrap();
		assert_eq!(l.master().qualified_nodes().len(), 3);
	}

	#[test]
	fn node_is_not_disqualified_if_enough_complaints_received_from_the_same_node() {
		let (sid, _, s, mut l) = make_simple_cluster(1, 4).unwrap();
		l.take_and_process_all_messages().unwrap();
		l.master().on_complaint(l.second_slave().node().clone(), message::Complaint {
			session: sid,
			against: s.clone(),
		}).unwrap();
		l.master().on_complaint(l.second_slave().node().clone(), message::Complaint {
			session: sid,
			against: s,
		}).unwrap();
		assert_eq!(l.master().qualified_nodes().len(), 4);
	}

	#[test]
	fn node_is_disqualified_if_responds_to_complain_with_invalid_data() {
		let (sid, _, _, mut l) = make_simple_cluster(1, 3).unwrap();
		l.take_and_process_message().unwrap(); // m -> s1: InitializeSession
		l.take_and_process_message().unwrap(); // m -> s2: InitializeSession
		l.take_and_process_message().unwrap(); // s1 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // s2 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // m -> s1: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s2: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s1: KeysDissemination
		l.take_and_process_message().unwrap(); // m -> s2: KeysDissemination
		l.take_and_process_message().unwrap(); // s1 -> m: KeysDissemination
		l.take_and_process_message().unwrap(); // s1 -> s2: KeysDissemination
		let s2 = l.second_slave().node().clone();
		l.master().on_keys_dissemination(s2.clone(), message::KeysDissemination {
			session: sid.clone(),
			secret1: math::generate_random_scalar().unwrap(),
			secret2: math::generate_random_scalar().unwrap(),
			publics: vec![math::generate_random_point().unwrap(), math::generate_random_point().unwrap()],
		}).unwrap();
		assert_eq!(l.master().qualified_nodes().len(), 3);
		l.master().on_complaint_response(s2, message::ComplaintResponse {
			session: sid,
			secret1: math::generate_random_scalar().unwrap(),
			secret2: math::generate_random_scalar().unwrap(),
		}).unwrap();
		assert_eq!(l.master().qualified_nodes().len(), 2);
	}

	#[test]
	fn node_is_not_disqualified_if_responds_to_complain_with_valid_data() {
		let (sid, _, _, mut l) = make_simple_cluster(1, 3).unwrap();
		l.take_and_process_message().unwrap(); // m -> s1: InitializeSession
		l.take_and_process_message().unwrap(); // m -> s2: InitializeSession
		l.take_and_process_message().unwrap(); // s1 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // s2 -> m: ConfirmInitialization
		l.take_and_process_message().unwrap(); // m -> s1: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s2: CompleteInitialization
		l.take_and_process_message().unwrap(); // m -> s1: KeysDissemination
		l.take_and_process_message().unwrap(); // m -> s2: KeysDissemination
		l.take_and_process_message().unwrap(); // s1 -> m: KeysDissemination
		l.take_and_process_message().unwrap(); // s1 -> s2: KeysDissemination
		let (f, t, msg) = match l.take_message() {
			Some((f, t, Message::KeysDissemination(msg))) => (f, t, msg),
			_ => panic!("unexpected"),
		};
		assert_eq!(&f, l.second_slave().node());
		assert_eq!(&t, l.master().node());
		l.master().on_keys_dissemination(f.clone(), message::KeysDissemination {
			session: sid.clone(),
			secret1: math::generate_random_scalar().unwrap(),
			secret2: math::generate_random_scalar().unwrap(),
			publics: msg.publics.clone(),
		}).unwrap();
		assert_eq!(l.master().qualified_nodes().len(), 3);
		l.master().on_complaint_response(f, message::ComplaintResponse {
			session: sid,
			secret1: msg.secret1,
			secret2: msg.secret2,
		}).unwrap();
		assert_eq!(l.master().qualified_nodes().len(), 3);
	}

	#[test]
	fn should_not_start_key_generation_when_not_in_key_check_state() {
		let (_, _, _, l) = make_simple_cluster(1, 3).unwrap();
		assert_eq!(l.master().start_key_generation_phase().unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn should_not_accept_public_key_share_when_is_not_waiting_for_it() {
		let (sid, _, s, l) = make_simple_cluster(1, 3).unwrap();
		assert_eq!(l.master().on_public_key_share(s, message::PublicKeyShare {
			session: sid,
			public_share: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidStateForRequest);
	}

	#[test]
	fn should_not_accept_public_key_share_when_receiving_twice() {
		let (sid, m, _, mut l) = make_simple_cluster(0, 3).unwrap();
		l.take_and_process_all_messages().unwrap();
		l.master().start_key_generation_phase().unwrap();
		l.first_slave().start_key_generation_phase().unwrap();
		let (f, t, msg) = match l.take_message() {
			Some((f, t, Message::PublicKeyShare(msg))) => (f, t, msg),
			_ => panic!("unexpected"),
		};
		assert_eq!(&f, l.master().node());
		assert_eq!(&t, l.first_slave().node());
		l.process_message((f, t, Message::PublicKeyShare(msg.clone()))).unwrap();
		assert_eq!(l.first_slave().on_public_key_share(m, message::PublicKeyShare {
			session: sid,
			public_share: math::generate_random_point().unwrap(),
		}).unwrap_err(), Error::InvalidMessage);
	}

	#[test]
	fn complete_enc_dec_session() {
		// TODO: when number of nodes, needed to decrypt message is odd, algorithm won't work
		// let test_cases = [/*(0, 2), */(1, 2), (1, 3), (2, 3), (1, 4), (2, 4), (3, 4), (1, 5), (2, 5), (3, 5), (4, 5),
		//	(1, 10), (2, 10), (3, 10), (4, 10), (5, 10), (6, 10), (7, 10), (8, 10), (9, 10)];
		let test_cases = [(3, 5)];
		for &(threshold, num_nodes) in &test_cases {
			let mut l = MessageLoop::new(num_nodes);
			l.master().initialize(threshold, l.nodes.keys().cloned().collect()).unwrap();
			assert_eq!(l.nodes.len(), num_nodes);

			// let nodes do initialization + keys dissemination
			while let Some((from, to, message)) = l.take_message() {
				l.process_message((from, to, message)).unwrap();
			}

			// check that all nodes are waiting for complaint timeout is passed
			for node in l.nodes.values() {
				let state = node.session.state();
				assert_eq!(state, SessionState::KeyCheck);

				// simulate timeout pass
				node.session.start_key_generation_phase().unwrap();
			}

			// let nodes do joint public generation
			while let Some((from, to, message)) = l.take_message() {
				l.process_message((from, to, message)).unwrap();
			}

			// check that all nodes has finished joint public generation
			let joint_public_key = l.master().joint_public_key().unwrap();
			for node in l.nodes.values() {
				let state = node.session.state();
				assert_eq!(state, SessionState::Finished);
				assert_eq!(node.session.joint_public_key().as_ref(), Some(&joint_public_key));
			}

			// now let's encrypt some secret (which is a point on EC)
			let document_secret_plain = Random.generate().unwrap().public().clone();
			let all_nodes_id_numbers: Vec<_> = l.master().data.lock().nodes.values().map(|n| n.id_number.clone()).collect();
			let all_nodes_secret_shares: Vec<_> = l.nodes.values().map(|n| n.session.data.lock().secret_share.as_ref().unwrap().clone()).collect();
			let document_secret_decrypted = do_encryption_and_decryption(threshold, &joint_public_key,
				&all_nodes_id_numbers,
				&all_nodes_secret_shares,
				None,
				document_secret_plain.clone()
			).0;
			assert_eq!(document_secret_plain, document_secret_decrypted);
		}
	}

	// TODO: add test where some nodes are disqualified from session
}