diff --git a/CHECKLIST.md b/CHECKLIST.md index b72082c..3f4a21e 100644 --- a/CHECKLIST.md +++ b/CHECKLIST.md @@ -48,7 +48,7 @@ - [x] Add new filters to the chain, verifying with the `FilterHash` - [ ] Optimizations - [x] Hashmap the `BlockHash` to `FilterHash` relationship in memory - - [ ] Persist SPKs that have already been proven to be in a filter + - [ ] Persist SPKs that have already been proven to be in a filter? #### Main thread @@ -81,7 +81,7 @@ - [x] Should serve CPF - [ ] Set up "timer" - [x] Check for DOS - - [ ] Message counter + - [x] Message counter - [ ] `Ping` if peer has not been heard from - [ ] `Disconnect` peers with high latency - [ ] Add BIP-324 with V1 fallback diff --git a/src/peers/counter.rs b/src/peers/counter.rs new file mode 100644 index 0000000..d7eaa89 --- /dev/null +++ b/src/peers/counter.rs @@ -0,0 +1,83 @@ +// Very simple denial of service protection so a peer cannot spam us with unsolicited messages. +#[derive(Debug, Clone)] +pub(crate) struct MessageCounter { + version: i8, + verack: i8, + header: i32, + filter_header: i32, + filters: i64, + addrs: i32, + block: i32, +} + +impl MessageCounter { + pub(crate) fn new() -> Self { + Self { + version: 1, + verack: 1, + header: 0, + filter_header: 0, + filters: 0, + addrs: 0, + block: 0, + } + } + + pub(crate) fn got_version(&mut self) { + self.version -= 1; + } + + pub(crate) fn got_verack(&mut self) { + self.verack -= 1; + } + + pub(crate) fn got_header(&mut self) { + self.header -= 1; + } + + pub(crate) fn got_filter_header(&mut self) { + self.filter_header -= 1; + } + + pub(crate) fn got_filter(&mut self) { + self.filters -= 1; + } + + pub(crate) fn got_addrs(&mut self) { + self.addrs -= 1; + } + + pub(crate) fn got_block(&mut self) { + self.block -= 1; + } + + pub(crate) fn sent_header(&mut self) { + self.header += 1; + } + + pub(crate) fn sent_filter_header(&mut self) { + self.filter_header += 1; + } + + pub(crate) fn sent_filters(&mut self) { + self.filters += 1000; + } + + pub(crate) fn sent_addrs(&mut self) { + self.addrs += 5; + } + + pub(crate) fn sent_block(&mut self) { + self.block += 1; + } + + pub(crate) fn unsolicited(&self) -> bool { + self.version < 0 + || self.header < 0 + || self.filters < 0 + || self.verack < 0 + || self.filter_header < 0 + || self.addrs < 0 + || self.block < 0 + } +} diff --git a/src/peers/mod.rs b/src/peers/mod.rs index 6f01ebc..e5712c0 100644 --- a/src/peers/mod.rs +++ b/src/peers/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod counter; pub(crate) mod dns; pub(crate) mod outbound_messages; pub(crate) mod peer; diff --git a/src/peers/peer.rs b/src/peers/peer.rs index b4a92fe..a4a062a 100644 --- a/src/peers/peer.rs +++ b/src/peers/peer.rs @@ -16,7 +16,7 @@ use crate::{ prelude::default_port_from_network, }; -use super::reader::Reader; +use super::{counter::MessageCounter, reader::Reader}; pub(crate) struct Peer { nonce: u32, @@ -25,12 +25,7 @@ pub(crate) struct Peer { main_thread_sender: Sender, main_thread_recv: Receiver, network: Network, -} - -enum State { - NotConnected, - Connected, - Verack, + message_counter: MessageCounter, } impl Peer { @@ -43,6 +38,7 @@ impl Peer { main_thread_recv: Receiver, ) -> Self { let default_port = default_port_from_network(&network); + let message_counter = MessageCounter::new(); Self { nonce, ip_addr, @@ -50,6 +46,7 @@ impl Peer { main_thread_sender, main_thread_recv, network, + message_counter, } } @@ -92,6 +89,10 @@ impl Peer { if read_handle.is_finished() { return Ok(()); } + if self.message_counter.unsolicited() { + println!("oh no"); + return Ok(()); + } select! { // The peer sent us a message peer_message = rx.recv() => { @@ -141,6 +142,7 @@ impl Peer { ) -> Result<(), PeerError> { match message { PeerMessage::Version(version) => { + self.message_counter.got_version(); self.main_thread_sender .send(PeerThreadMessage { nonce: self.nonce, @@ -160,6 +162,7 @@ impl Peer { Ok(()) } PeerMessage::Addr(addrs) => { + self.message_counter.got_addrs(); self.main_thread_sender .send(PeerThreadMessage { nonce: self.nonce, @@ -170,6 +173,7 @@ impl Peer { Ok(()) } PeerMessage::Headers(headers) => { + self.message_counter.got_header(); self.main_thread_sender .send(PeerThreadMessage { nonce: self.nonce, @@ -180,6 +184,7 @@ impl Peer { Ok(()) } PeerMessage::FilterHeaders(cf_headers) => { + self.message_counter.got_filter_header(); self.main_thread_sender .send(PeerThreadMessage { nonce: self.nonce, @@ -190,6 +195,7 @@ impl Peer { Ok(()) } PeerMessage::Filter(filter) => { + self.message_counter.got_filter(); self.main_thread_sender .send(PeerThreadMessage { nonce: self.nonce, @@ -200,6 +206,7 @@ impl Peer { Ok(()) } PeerMessage::Block(block) => { + self.message_counter.got_block(); self.main_thread_sender .send(PeerThreadMessage { nonce: self.nonce, @@ -219,7 +226,10 @@ impl Peer { .map_err(|_| PeerError::ThreadChannel)?; Ok(()) } - PeerMessage::Verack => Ok(()), + PeerMessage::Verack => { + self.message_counter.got_verack(); + Ok(()) + } PeerMessage::Ping(nonce) => { writer .write_all(&message_generator.new_pong(nonce)) @@ -249,12 +259,14 @@ impl Peer { ) -> Result<(), PeerError> { match request { MainThreadMessage::GetAddr => { + self.message_counter.sent_addrs(); writer .write_all(&message_generator.new_get_addr()) .await .map_err(|_| PeerError::BufferWrite)?; } MainThreadMessage::GetHeaders(config) => { + self.message_counter.sent_header(); let message = message_generator.new_get_headers(config.locators, config.stop_hash); writer .write_all(&message) @@ -262,6 +274,7 @@ impl Peer { .map_err(|_| PeerError::BufferWrite)?; } MainThreadMessage::GetFilterHeaders(config) => { + self.message_counter.sent_filter_header(); let message = message_generator.new_cf_headers(config); writer .write_all(&message) @@ -269,6 +282,7 @@ impl Peer { .map_err(|_| PeerError::BufferWrite)?; } MainThreadMessage::GetFilters(config) => { + self.message_counter.sent_filters(); let message = message_generator.new_filters(config); writer .write_all(&message) @@ -276,6 +290,7 @@ impl Peer { .map_err(|_| PeerError::BufferWrite)?; } MainThreadMessage::GetBlock(message) => { + self.message_counter.sent_block(); let message = message_generator.new_block(message); writer .write_all(&message)