use crate::block::{Block, BlockType};
use crate::blockchain::{Blockchain, GENESIS_PERIOD};
use crate::consensus::SaitoMessage;
use crate::crypto::{hash, verify, SaitoHash, SaitoPublicKey};
use crate::hop::Hop;
use crate::mempool::Mempool;
use crate::network::{
Network, CHALLENGE_EXPIRATION_TIME, CHALLENGE_SIZE, INBOUND_PEER_CONNECTIONS_GLOBAL,
OUTBOUND_PEER_CONNECTIONS_GLOBAL, PEERS_DB_GLOBAL, PEERS_REQUEST_RESPONSES_GLOBAL,
PEERS_REQUEST_WAKERS_GLOBAL,
};
use crate::networking::message_types::handshake_challenge::HandshakeChallenge;
use crate::networking::message_types::request_block_message::RequestBlockMessage;
use crate::networking::message_types::request_blockchain_message::RequestBlockchainMessage;
use crate::networking::message_types::send_block_head_message::SendBlockHeadMessage;
use crate::networking::message_types::send_blockchain_message::{
SendBlockchainBlockData, SendBlockchainMessage, SyncType,
};
use crate::time::create_timestamp;
use crate::transaction::Transaction;
use crate::wallet::Wallet;
use async_recursion::async_recursion;
use futures::stream::SplitSink;
use std::collections::HashMap;
use std::convert::TryInto;
use std::error::Error;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll, Waker};
use tokio::sync::{broadcast, mpsc, RwLock};
use tokio_stream::wrappers::UnboundedReceiverStream;
use tracing::{error, info};
use uuid::Uuid;
use warp::ws::{Message, WebSocket};
use crate::networking::api_message::APIMessage;
use futures::{Future, FutureExt, SinkExt, StreamExt};
use tokio::net::TcpStream;
use tokio::sync::broadcast::Sender;
use tokio_tungstenite::{tungstenite, MaybeTlsStream, WebSocketStream};
pub type PeersDB = HashMap<SaitoHash, SaitoPeer>;
pub type RequestResponses = HashMap<(SaitoHash, u32), APIMessage>;
pub type RequestWakers = HashMap<(SaitoHash, u32), Waker>;
pub type OutboundPeersDB = HashMap<SaitoHash, OutboundPeer>;
pub type InboundPeersDB = HashMap<SaitoHash, InboundPeer>;
pub struct PeerFlags {
is_connected_or_connecting: bool,
has_completed_handshake: bool,
is_from_peer_list: bool,
}
pub struct SaitoPeer {
peer_flags: PeerFlags,
connection_id: SaitoHash,
publickey: Option<SaitoPublicKey>,
host: Option<[u8; 4]>,
port: Option<u16>,
request_count: u32,
wallet_lock: Arc<RwLock<Wallet>>,
mempool_lock: Arc<RwLock<Mempool>>,
blockchain_lock: Arc<RwLock<Blockchain>>,
broadcast_channel_sender: broadcast::Sender<SaitoMessage>,
}
pub struct InboundPeer {
pub sender: mpsc::UnboundedSender<std::result::Result<Message, warp::Error>>,
}
pub struct OutboundPeer {
pub write_sink:
SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::protocol::Message>,
}
pub struct PeerRequest {
connection_id: SaitoHash,
request_id: u32,
api_message_command: String,
}
impl PeerRequest {
pub async fn new(command: &str, message: Vec<u8>, peer: &mut SaitoPeer) -> Self {
peer.request_count += 1;
let api_message = APIMessage::new(command, peer.request_count - 1, message);
send_message_to_socket(api_message, &peer.connection_id).await;
PeerRequest {
connection_id: peer.connection_id,
request_id: peer.request_count - 1,
api_message_command: String::from(command),
}
}
}
impl Future for PeerRequest {
type Output = Result<APIMessage, Box<dyn Error>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let request_responses_lock = PEERS_REQUEST_RESPONSES_GLOBAL.clone();
let mut request_responses = request_responses_lock.write().unwrap();
match request_responses.remove(&(self.connection_id, self.request_id)) {
Some(response) => {
info!("HANDLING RESPONSE {}", self.api_message_command);
Poll::Ready(Ok(response))
}
None => {
let request_wakers_lock = PEERS_REQUEST_WAKERS_GLOBAL.clone();
let mut request_wakers = request_wakers_lock.write().unwrap();
request_wakers.insert((self.connection_id, self.request_id), cx.waker().clone());
Poll::Pending
}
}
}
}
pub async fn send_message_to_socket(api_message: APIMessage, connection_id: &SaitoHash) {
let outbound_peer_connection_db_global = OUTBOUND_PEER_CONNECTIONS_GLOBAL.clone();
let mut outbound_peer_connection_db = outbound_peer_connection_db_global.write().await;
let inbound_peer_connection_db_global = INBOUND_PEER_CONNECTIONS_GLOBAL.clone();
let mut inbound_peer_connection_db = inbound_peer_connection_db_global.write().await;
let outbound_peer_connection = outbound_peer_connection_db.get_mut(connection_id);
let inbound_peer_connection = inbound_peer_connection_db.get_mut(connection_id);
if inbound_peer_connection.is_some() && outbound_peer_connection.is_some() {
panic!("Peer has both inbound and outbound connections, this should never happen");
}
if inbound_peer_connection.is_none() && outbound_peer_connection.is_none() {
panic!("Peer has neither outbound nor inbound connection");
}
if outbound_peer_connection.is_some() {
outbound_peer_connection
.unwrap()
.write_sink
.send(api_message.serialize().into())
.await
.expect("unable to send to outbound_peer_connection... It may be that the socket\n
was closed but the outbound peer was not cleaned up. If this occurs, it should be investigated.");
}
if inbound_peer_connection.is_some() {
inbound_peer_connection
.unwrap()
.sender
.send(Ok(Message::binary(api_message.serialize())))
.expect("unable to send to outbound_peer_connection... It may be that the socket\n
was closed but the inbount peer was not cleaned up. If this occurs, it should be investigated.");
}
}
impl SaitoPeer {
pub fn new(
connection_id: SaitoHash,
host: Option<[u8; 4]>,
port: Option<u16>,
is_connected_or_connecting: bool,
has_completed_handshake: bool,
is_from_peer_list: bool,
wallet_lock: Arc<RwLock<Wallet>>,
mempool_lock: Arc<RwLock<Mempool>>,
blockchain_lock: Arc<RwLock<Blockchain>>,
broadcast_channel_sender: Sender<SaitoMessage>,
) -> SaitoPeer {
SaitoPeer {
peer_flags: PeerFlags {
is_connected_or_connecting,
has_completed_handshake,
is_from_peer_list,
},
connection_id,
host,
port,
publickey: None,
request_count: 0,
wallet_lock,
mempool_lock,
blockchain_lock,
broadcast_channel_sender,
}
}
pub fn get_is_from_peer_list(&self) -> bool {
self.peer_flags.is_from_peer_list
}
pub fn set_has_completed_handshake(&mut self, has_completed_handshake: bool) {
self.peer_flags.has_completed_handshake = has_completed_handshake;
}
pub fn get_has_completed_handshake(&self) -> bool {
self.peer_flags.has_completed_handshake
}
pub fn set_publickey(&mut self, publickey: SaitoPublicKey) {
self.publickey = Some(publickey)
}
pub fn get_publickey(&self) -> Option<SaitoPublicKey> {
self.publickey
}
pub fn get_broadcast_channel_sender(&self) -> &broadcast::Sender<SaitoMessage> {
&self.broadcast_channel_sender
}
pub async fn set_is_connected_or_connecting(&mut self, is_connected_or_connecting: bool) {
if !is_connected_or_connecting {
let outbound_peer_connection_db_global = OUTBOUND_PEER_CONNECTIONS_GLOBAL.clone();
let mut outbound_peer_connection_db = outbound_peer_connection_db_global.write().await;
outbound_peer_connection_db.remove(&self.connection_id);
let inbound_peer_connection_db_global = INBOUND_PEER_CONNECTIONS_GLOBAL.clone();
let mut inbound_peer_connection_db = inbound_peer_connection_db_global.write().await;
inbound_peer_connection_db.remove(&self.connection_id);
self.peer_flags.has_completed_handshake = false;
}
self.peer_flags.is_connected_or_connecting = is_connected_or_connecting;
}
pub fn get_is_connected_or_connecting(&self) -> bool {
self.peer_flags.is_connected_or_connecting
}
pub fn get_host(&self) -> Option<[u8; 4]> {
self.host
}
pub fn get_port(&self) -> Option<u16> {
self.port
}
pub fn get_connection_id(&self) -> SaitoHash {
self.connection_id
}
pub fn is_in_path(&self, path: &Vec<Hop>) -> bool {
for hop in path {
if self.publickey.unwrap() == hop.get_from() {
return true;
}
}
false
}
#[async_recursion]
pub async fn send_command(
&mut self,
command: &str,
message: Vec<u8>,
) -> Result<APIMessage, APIMessage> {
let peer_request = PeerRequest::new(command, message, self).await;
let response_message = peer_request
.await
.expect(&format!("Error returned from {}", command));
match response_message.get_message_name_as_string().as_str() {
"RESULT__" => Ok(response_message),
"ERROR___" => Err(response_message),
_ => {
panic!("Received non-response response");
}
}
}
pub async fn send_command_fire_and_forget(&mut self, command: &str, message: Vec<u8>) {
let _peer_request = PeerRequest::new(command, message, self).await;
}
pub async fn send_response_from_str(&mut self, message_id: u32, message_str: &str) {
send_message_to_socket(
APIMessage::new_from_string("RESULT__", message_id, message_str),
&self.connection_id,
)
.await;
}
pub async fn send_response(&mut self, message_id: u32, message: Vec<u8>) {
send_message_to_socket(
APIMessage::new("RESULT__", message_id, message),
&self.connection_id,
)
.await;
}
pub async fn send_error_response_from_str(&mut self, message_id: u32, message_str: &str) {
send_message_to_socket(
APIMessage::new_from_string("ERROR___", message_id, message_str),
&self.connection_id,
)
.await;
}
pub async fn send_error_response(&mut self, message_id: u32, message: Vec<u8>) {
send_message_to_socket(
APIMessage::new("ERROR___", message_id, message),
&self.connection_id,
)
.await;
}
pub async fn handle_peer_message(api_message_orig: APIMessage, connection_id: SaitoHash) {
match api_message_orig.get_message_name_as_string().as_str() {
"RESULT__" | "ERROR___" => {
let request_wakers_lock = PEERS_REQUEST_WAKERS_GLOBAL.clone();
let mut request_wakers = request_wakers_lock.write().unwrap();
let option_waker =
request_wakers.remove(&(connection_id, api_message_orig.message_id));
let request_responses_lock = PEERS_REQUEST_RESPONSES_GLOBAL.clone();
let mut request_responses = request_responses_lock.write().unwrap();
request_responses.insert(
(connection_id, api_message_orig.message_id),
api_message_orig,
);
if let Some(waker) = option_waker {
waker.wake();
}
}
_ => {
let peers_db_global = PEERS_DB_GLOBAL.clone();
let mut peer_db = peers_db_global.write().await;
let peer = peer_db.get_mut(&connection_id).unwrap();
SaitoPeer::handle_peer_command(peer, api_message_orig).await;
}
}
}
pub async fn do_reqblock(&self, block_hash: SaitoHash) {
let request_block_message = RequestBlockMessage::new(None, Some(block_hash), None);
let connection_id_clone = self.connection_id.clone();
let mempool_lock = self.mempool_lock.clone();
tokio::spawn(async move {
let peers_db_global = PEERS_DB_GLOBAL.clone();
let mut peer_db = peers_db_global.write().await;
let peer = peer_db.get_mut(&connection_id_clone).unwrap();
let result = peer
.send_command(&String::from("REQBLOCK"), request_block_message.serialize())
.await;
match result {
Ok(serialized_block_message) => {
let mut block =
Block::deserialize_for_net(serialized_block_message.get_message_data());
block.set_source_connection_id(peer.connection_id);
{
let mut mempool = mempool_lock.write().await;
mempool.add_block(block);
}
Mempool::send_blocks_to_blockchain(
peer.mempool_lock.clone(),
peer.blockchain_lock.clone(),
)
.await;
}
Err(error_message) => {
error!(
"REQBLOCK ERROR: {}",
error_message.get_message_data_as_string()
);
}
}
});
}
async fn handle_peer_command(peer: &mut SaitoPeer, api_message: APIMessage) {
let mempool_lock = peer.mempool_lock.clone();
let blockchain_lock = peer.blockchain_lock.clone();
let command = api_message.get_message_name_as_string();
info!("HANDLING COMMAND {}", command);
match command.as_str() {
"SHAKINIT" => {
if let Ok(serialized_handshake_challenge) =
build_serialized_challenge(&api_message, peer.wallet_lock.clone()).await
{
peer.send_response(api_message.message_id, serialized_handshake_challenge)
.await;
}
}
"SHAKCOMP" => match socket_handshake_verify(&api_message.get_message_data()) {
Some(deserialize_challenge) => {
peer.set_has_completed_handshake(true);
peer.set_publickey(deserialize_challenge.opponent_pubkey());
peer.send_response(
api_message.message_id,
String::from("OK").as_bytes().try_into().unwrap(),
)
.await;
}
None => {
error!("Error verifying peer handshake signature");
}
},
"REQBLOCK" => {
let api_message = build_request_block_response(&api_message, blockchain_lock).await;
send_message_to_socket(api_message, &peer.connection_id).await;
}
"REQBLKHD" => {
let message_id = api_message.message_id;
if let Some(bytes) = socket_send_block_header(&api_message, blockchain_lock).await {
let message_data = String::from("OK").as_bytes().try_into().unwrap();
peer.send_response(message_id, message_data).await;
let _result = peer.send_command("SNDBLKHD", bytes).await;
} else {
peer.send_error_response_from_str(message_id, "ERROR").await;
}
}
"REQCHAIN" => {
peer.send_response_from_str(api_message.message_id, "OK")
.await;
if let Some(send_blockchain_message) = build_send_blockchain_message(
&RequestBlockchainMessage::deserialize(&api_message.get_message_data()),
blockchain_lock,
)
.await
{
let connection_id_clone = peer.connection_id.clone();
tokio::spawn(async move {
let peers_db_global = PEERS_DB_GLOBAL.clone();
let mut peer_db = peers_db_global.write().await;
let peer = peer_db.get_mut(&connection_id_clone).unwrap();
let _result = peer
.send_command(
&String::from("SNDCHAIN"),
send_blockchain_message.serialize(),
)
.await;
});
} else {
peer.send_error_response_from_str(api_message.message_id, "UNKNOWN BLOCK HASH")
.await;
}
}
"SNDCHAIN" => {
peer.send_response_from_str(api_message.message_id, "OK")
.await;
let send_blockchain_message =
SendBlockchainMessage::deserialize(api_message.get_message_data());
for send_blockchain_block_data in
send_blockchain_message.get_blocks_data().into_iter()
{
peer.do_reqblock(send_blockchain_block_data.block_hash)
.await;
}
}
"SNDBLKHD" => {
let send_block_head_message =
SendBlockHeadMessage::deserialize(api_message.get_message_data());
send_block_head_message.get_block_hash();
let blockchain = blockchain_lock.read().await;
match blockchain
.get_block(send_block_head_message.get_block_hash())
.await
{
Some(_block) => {
info!(
"SNDBLKHD hash already known: {}",
hex::encode(send_block_head_message.get_block_hash()),
);
}
None => {
let message_data = String::from("OK").as_bytes().try_into().unwrap();
peer.send_response(api_message.get_message_id(), message_data)
.await;
peer.do_reqblock(send_block_head_message.get_block_hash().clone())
.await
}
}
}
"SNDTRANS" => {
if let Some(mut tx) = socket_receive_transaction(api_message.clone()) {
let wallet_lock_clone = peer.wallet_lock.clone();
let wallet = wallet_lock_clone.read().await;
tx.generate_metadata(wallet.get_publickey());
let blockchain = blockchain_lock.read().await;
let mut mempool = mempool_lock.write().await;
if !mempool.transaction_exists(tx.get_hash_for_signature()) {
if tx.validate(&blockchain.utxoset, &blockchain.staking) {
mempool.add_transaction(tx.clone()).await;
peer.send_response_from_str(api_message.message_id, "OK")
.await;
Network::propagate_transaction(peer.wallet_lock.clone(), tx).await;
} else {
peer.send_error_response_from_str(
api_message.message_id,
"INVALID TRANSACTION",
)
.await;
}
} else {
peer.send_error_response_from_str(
api_message.message_id,
"TRANSACTION ALREADY EXISTED IN MEMPOOL",
)
.await;
}
}
}
"SNDKYLST" => {
peer.send_error_response_from_str(api_message.message_id, "UNHANDLED COMMAND")
.await;
}
_ => {
error!(
"Unhandled command received by client... {}",
&api_message.get_message_name_as_string()
);
peer.send_error_response_from_str(api_message.message_id, "NO SUCH")
.await;
}
}
}
}
pub async fn handle_inbound_peer_connection(
ws: WebSocket,
peer_db_lock: Arc<RwLock<PeersDB>>,
wallet_lock: Arc<RwLock<Wallet>>,
mempool_lock: Arc<RwLock<Mempool>>,
blockchain_lock: Arc<RwLock<Blockchain>>,
broadcast_channel_sender: broadcast::Sender<SaitoMessage>,
) {
let (peer_ws_sender, mut peer_ws_rcv) = ws.split();
let (peer_sender, peer_rcv) = mpsc::unbounded_channel();
let peer_rcv = UnboundedReceiverStream::new(peer_rcv);
tokio::task::spawn(peer_rcv.forward(peer_ws_sender).map(|result| {
if let Err(e) = result {
error!("error sending websocket msg: {}", e);
}
}));
let connection_id: SaitoHash = hash(&Uuid::new_v4().as_bytes().to_vec());
let peer = SaitoPeer::new(
connection_id,
None,
None,
true,
false,
false,
wallet_lock.clone(),
mempool_lock.clone(),
blockchain_lock.clone(),
broadcast_channel_sender.clone(),
);
peer_db_lock
.write()
.await
.insert(connection_id.clone(), peer);
let inbound_peer_db_lock_global = INBOUND_PEER_CONNECTIONS_GLOBAL.clone();
inbound_peer_db_lock_global.write().await.insert(
connection_id.clone(),
InboundPeer {
sender: peer_sender,
},
);
tokio::task::spawn(async move {
while let Some(result) = peer_ws_rcv.next().await {
let msg = match result {
Ok(msg) => msg,
Err(e) => {
eprintln!("error receiving ws message for peer: {}", e);
break;
}
};
if !msg.as_bytes().is_empty() {
let api_message = APIMessage::deserialize(&msg.as_bytes().to_vec());
SaitoPeer::handle_peer_message(api_message, connection_id).await;
} else {
error!(
"Message of length 0... why?\n
This seems to occur if we aren't holding a reference to the sender/stream on the\n
other end of the connection. I suspect that when the stream goes out of scope,\n
it's deconstructor is being called and sends a 0 length message to indicate\n
that the stream has ended... I'm leaving this println here for now because\n
it would be very helpful to see this if starts to occur again. We may want to\n
treat this as a disconnect."
);
}
}
{
let mut peer_db = peer_db_lock.write().await;
let peer = peer_db.get_mut(&connection_id).unwrap();
peer.set_is_connected_or_connecting(false).await;
peer_db.remove(&connection_id);
}
});
}
pub async fn build_serialized_challenge(
message: &APIMessage,
wallet_lock: Arc<RwLock<Wallet>>,
) -> crate::Result<Vec<u8>> {
let wallet = wallet_lock.read().await;
let my_pubkey = wallet.get_publickey();
let my_privkey = wallet.get_privatekey();
let mut peer_octets: [u8; 4] = [0; 4];
peer_octets[0..4].clone_from_slice(&message.message_data[0..4]);
let peer_pubkey: SaitoPublicKey = message.message_data[4..37].try_into().unwrap();
let my_octets: [u8; 4] = [127, 0, 0, 1];
let challenge = HandshakeChallenge::new((my_octets, my_pubkey), (peer_octets, peer_pubkey));
let serialized_challenge = challenge.serialize_with_sig(my_privkey);
Ok(serialized_challenge)
}
pub fn socket_handshake_verify(message_data: &Vec<u8>) -> Option<HandshakeChallenge> {
let challenge = HandshakeChallenge::deserialize(message_data);
if challenge.timestamp() < create_timestamp() - CHALLENGE_EXPIRATION_TIME {
error!("Error validating timestamp for handshake complete");
return None;
}
if !verify(
&hash(&message_data[..CHALLENGE_SIZE + 64].to_vec()),
challenge.opponent_sig().unwrap(),
challenge.opponent_pubkey(),
) {
error!("Error with validating opponent sig");
return None;
}
if !verify(
&hash(&message_data[..CHALLENGE_SIZE].to_vec()),
challenge.challenger_sig().unwrap(),
challenge.challenger_pubkey(),
) {
error!("Error with validating challenger sig");
return None;
}
Some(challenge)
}
pub fn socket_receive_transaction(message: APIMessage) -> Option<Transaction> {
let tx = Transaction::deserialize_from_net(message.message_data);
Some(tx)
}
pub async fn build_request_block_response(
api_message: &APIMessage,
blockchain_lock: Arc<RwLock<Blockchain>>,
) -> APIMessage {
let request_block_message = RequestBlockMessage::deserialize(api_message.get_message_data());
let blockchain = blockchain_lock.read().await;
if request_block_message.get_block_id().is_some() {
APIMessage::new_from_string("ERROR___", api_message.message_id, "Unsupported: fetching blocks by id is not yet supported, please fetch the block by hash.")
} else if request_block_message.get_block_hash().is_some() {
let block_hash: SaitoHash = request_block_message.get_block_hash().unwrap();
match blockchain.get_block_sync(&block_hash) {
Some(target_block) => APIMessage::new(
"RESULT__",
api_message.message_id,
target_block.serialize_for_net(BlockType::Full),
),
None => APIMessage::new_from_string(
"ERROR___",
api_message.message_id,
"Unknown Block Hash",
),
}
} else {
APIMessage::new_from_string(
"ERROR___",
api_message.message_id,
"REQBLOCK requires ID or Hash of desired block",
)
}
}
pub async fn socket_send_block_header(
api_message: &APIMessage,
blockchain_lock: Arc<RwLock<Blockchain>>,
) -> Option<Vec<u8>> {
let block_hash: SaitoHash = api_message.message_data[0..32].try_into().unwrap();
let blockchain = blockchain_lock.read().await;
match blockchain.get_block_sync(&block_hash) {
Some(target_block) => Some(target_block.serialize_for_net(BlockType::Header)),
None => None,
}
}
pub async fn build_send_blockchain_message(
request_blockchain_message: &RequestBlockchainMessage,
blockchain_lock: Arc<RwLock<Blockchain>>,
) -> Option<SendBlockchainMessage> {
let block_zero_hash: SaitoHash = [0; 32];
let peers_latest_hash: &SaitoHash;
if request_blockchain_message.get_latest_block_id() == 0
&& request_blockchain_message.get_latest_block_hash() == &block_zero_hash
{
peers_latest_hash = &block_zero_hash;
} else {
let blockchain = blockchain_lock.read().await;
if !blockchain.contains_block_hash_at_block_id(
request_blockchain_message.get_latest_block_id(),
*request_blockchain_message.get_latest_block_hash(),
) {
return None;
}
peers_latest_hash = request_blockchain_message.get_latest_block_hash();
}
let blockchain = blockchain_lock.read().await;
let mut blocks_data: Vec<SendBlockchainBlockData> = vec![];
if let Some(latest_block) = blockchain.get_latest_block() {
let mut previous_block_hash: SaitoHash = latest_block.get_hash();
let mut this_block: &Block; let mut block_count = 0;
while &previous_block_hash != peers_latest_hash && block_count < GENESIS_PERIOD {
block_count += 1;
this_block = blockchain.get_block_sync(&previous_block_hash).unwrap();
blocks_data.push(SendBlockchainBlockData {
block_id: this_block.get_id(),
block_hash: this_block.get_hash(),
timestamp: this_block.get_timestamp(),
pre_hash: [0; 32],
number_of_transactions: 0,
});
previous_block_hash = this_block.get_previous_block_hash();
}
Some(SendBlockchainMessage::new(
SyncType::Full,
*peers_latest_hash,
blocks_data,
))
} else {
panic!("Blockchain does not have any blocks");
}
}