use std::cmp::Ordering; use async_executor::Executor; use async_std::sync::Arc; use async_trait::async_trait; use chrono::Utc; use log::debug; use rand::{rngs::OsRng, RngCore}; use darkfi::{ net, serial::{SerialDecodable, SerialEncodable}, Result, }; use crate::{buffers::Buffers, settings, Privmsg}; #[derive(SerialDecodable, SerialEncodable, Clone, Debug)] struct Inv { id: u64, invs: Vec, } #[derive(SerialDecodable, SerialEncodable, Clone, Debug)] pub struct LastTerm { pub term: u64, } impl Inv { fn new(invs: Vec) -> Self { let id = OsRng.next_u64(); Self { id, invs } } } #[derive(SerialDecodable, SerialEncodable, Clone, Debug)] struct GetData { invs: Vec, term: Option, } impl GetData { fn new(invs: Vec, term: Option) -> Self { Self { invs, term } } } #[derive(SerialDecodable, SerialEncodable, Clone, Debug)] struct InvObject(String); pub struct ProtocolPrivmsg { jobsman: net::ProtocolJobsManagerPtr, notify: async_channel::Sender, msg_sub: net::MessageSubscription, inv_sub: net::MessageSubscription, getdata_sub: net::MessageSubscription, last_term_sub: net::MessageSubscription, p2p: net::P2pPtr, channel: net::ChannelPtr, buffers: Buffers, } impl ProtocolPrivmsg { pub async fn init( channel: net::ChannelPtr, notify: async_channel::Sender, p2p: net::P2pPtr, buffers: Buffers, ) -> net::ProtocolBasePtr { let message_subsytem = channel.get_message_subsystem(); message_subsytem.add_dispatch::().await; message_subsytem.add_dispatch::().await; message_subsytem.add_dispatch::().await; message_subsytem.add_dispatch::().await; let msg_sub = channel.clone().subscribe_msg::().await.expect("Missing Privmsg dispatcher!"); let inv_sub = channel.subscribe_msg::().await.expect("Missing Inv dispatcher!"); let getdata_sub = channel.clone().subscribe_msg::().await.expect("Missing GetData dispatcher!"); let last_term_sub = channel .clone() .subscribe_msg::() .await .expect("Missing LastTerm dispatcher!"); Arc::new(Self { notify, msg_sub, inv_sub, getdata_sub, last_term_sub, jobsman: net::ProtocolJobsManager::new("ProtocolPrivmsg", channel.clone()), p2p, channel, buffers, }) } async fn handle_receive_inv(self: Arc) -> Result<()> { debug!(target: "ircd", "ProtocolPrivmsg::handle_receive_inv() [START]"); let exclude_list = vec![self.channel.address()]; loop { let inv = self.inv_sub.receive().await?; let inv = (*inv).to_owned(); if !self.buffers.seen_ids.push(inv.id).await { continue } let mut inv_requested = vec![]; for inv_object in inv.invs.iter() { if !self.buffers.unread_msgs.inc_read_confirms(&inv_object.0).await { inv_requested.push(inv_object.clone()); } } if !inv_requested.is_empty() { self.channel.send(GetData::new(inv_requested, None)).await?; } self.update_unread_msgs().await?; self.p2p.broadcast_with_exclude(inv, &exclude_list).await?; } } async fn handle_receive_msg(self: Arc) -> Result<()> { debug!(target: "ircd", "ProtocolPrivmsg::handle_receive_msg() [START]"); let exclude_list = vec![self.channel.address()]; loop { let msg = self.msg_sub.receive().await?; let mut msg = (*msg).to_owned(); if !self.buffers.seen_ids.push(msg.id).await { continue } if msg.read_confirms >= settings::MAX_CONFIRM { self.add_to_msgs(&msg).await?; } else { msg.read_confirms += 1; let hash = self.add_to_unread_msgs(&msg).await; self.p2p.broadcast(Inv::new(vec![InvObject(hash)])).await?; } self.update_unread_msgs().await?; self.p2p.broadcast_with_exclude(msg, &exclude_list).await?; } } async fn handle_receive_last_term(self: Arc) -> Result<()> { debug!(target: "ircd", "ProtocolPrivmsg::handle_receive_last_term() [START]"); loop { let last_term = self.last_term_sub.receive().await?; let last_term = last_term.term; self.update_unread_msgs().await?; match self.buffers.privmsgs.last_term().await.cmp(&last_term) { Ordering::Greater => { for msg in self.buffers.privmsgs.fetch_msgs(last_term).await { self.channel.send(msg).await?; } } Ordering::Less => { self.channel.send(GetData::new(vec![], Some(last_term))).await?; } Ordering::Equal => continue, } } } async fn handle_receive_getdata(self: Arc) -> Result<()> { debug!(target: "ircd", "ProtocolPrivmsg::handle_receive_getdata() [START]"); loop { let getdata = self.getdata_sub.receive().await?; let getdata = (*getdata).to_owned(); for inv in getdata.invs { if let Some(msg) = self.buffers.unread_msgs.get(&inv.0).await { self.channel.send(msg.clone()).await?; } } if let Some(term) = getdata.term { for msg in self.buffers.privmsgs.fetch_msgs(term).await { self.channel.send(msg).await?; } } } } async fn add_to_unread_msgs(&self, msg: &Privmsg) -> String { self.buffers.unread_msgs.insert(msg).await } async fn update_unread_msgs(&self) -> Result<()> { for (hash, msg) in self.buffers.unread_msgs.load().await { if msg.timestamp + settings::UNREAD_MSG_EXPIRE_TIME < Utc::now().timestamp() { self.buffers.unread_msgs.remove(&hash).await; continue } if msg.read_confirms >= settings::MAX_CONFIRM { if let Some(msg) = self.buffers.unread_msgs.remove(&hash).await { self.add_to_msgs(&msg).await?; } } } Ok(()) } async fn add_to_msgs(&self, msg: &Privmsg) -> Result<()> { self.buffers.privmsgs.push(msg).await; self.notify.send(msg.clone()).await?; Ok(()) } } #[async_trait] impl net::ProtocolBase for ProtocolPrivmsg { /// Starts ping-pong keep-alive messages exchange. Runs ping-pong in the /// protocol task manager, then queues the reply. Sends out a ping and /// waits for pong reply. Waits for ping and replies with a pong. async fn start(self: Arc, executor: Arc>) -> Result<()> { // once a channel get started for m in self.buffers.privmsgs.load().await { self.channel.send(m).await?; } debug!(target: "ircd", "ProtocolPrivmsg::start() [START]"); self.jobsman.clone().start(executor.clone()); self.jobsman.clone().spawn(self.clone().handle_receive_msg(), executor.clone()).await; self.jobsman.clone().spawn(self.clone().handle_receive_inv(), executor.clone()).await; self.jobsman.clone().spawn(self.clone().handle_receive_getdata(), executor.clone()).await; self.jobsman.clone().spawn(self.clone().handle_receive_last_term(), executor.clone()).await; debug!(target: "ircd", "ProtocolPrivmsg::start() [END]"); Ok(()) } fn name(&self) -> &'static str { "ProtocolPrivmsg" } } impl net::Message for Privmsg { fn name() -> &'static str { "privmsg" } } impl net::Message for Inv { fn name() -> &'static str { "inv" } } impl net::Message for GetData { fn name() -> &'static str { "getdata" } } impl net::Message for LastTerm { fn name() -> &'static str { "last_term" } }