peer_server: fill in more scaffolding for rate limiter

2 files changed, 17 insertions, 1 deletions

diff --git a/Cargo.lock b/Cargo.lock
index ce8f1bd..180b7ea 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -489,7 +489,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 [[package]]
 name = "libc"
 version = "0.2.40"
-source = "git+https://github.com/rust-lang/libc#d0a57265999c78aa56c0202d97911e16e43baea3"
+source = "git+https://github.com/rust-lang/libc#5e025d65d40bde298dc36c289f11c5717fab6fa3"
 
 [[package]]
 name = "libc"
diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 255553f..13f5997 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -18,6 +18,7 @@ use tokio_core::reactor::Handle;
 
 use std::collections::VecDeque;
 use std::convert::TryInto;
+use std::net::IpAddr;
 use std::rc::Rc;
 use std::time::Instant;
 
@@ -187,6 +188,21 @@ impl PeerServer {
         let (mac_in, mac_out) = packet.split_at(116);
         self.cookie.verify_mac1(&mac_in[..], &mac_out[..16])?;
 
+        if self.under_load() {
+            let mac2_verified = match addr.ip() {
+                IpAddr::V4(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).is_ok(),
+                IpAddr::V6(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).is_ok(),
+            };
+
+            if !mac2_verified {
+                bail!("would send cookie request now");
+            }
+
+            if !self.rate_limiter.allow(&addr.ip()) {
+                bail!("rejected by rate limiter.");
+            }
+        }
+
         debug!("got handshake initiation request (0x01)");
 
         let handshake = Peer::process_incoming_handshake(