improve socket channel close handling, and test the node banning (roughly)

backend/telemetry_shard/src/blocked_addrs.rs

@@ -15,9 +15,9 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use std::collections::HashMap;
-use std::time::{Duration, Instant};
 use std::net::IpAddr;
-use std::sync::{ Mutex, Arc };
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
 
 /// Keep track of nodes that have been blocked.
 #[derive(Debug, Clone)]
@@ -26,7 +26,7 @@ pub struct BlockedAddrs(Arc<BlockAddrsInner>);
 #[derive(Debug)]
 struct BlockAddrsInner {
     block_duration: Duration,
-    inner: Mutex<HashMap<IpAddr, (&'static str, Instant)>>
+    inner: Mutex<HashMap<IpAddr, (&'static str, Instant)>>,
 }
 
 impl BlockedAddrs {
@@ -35,7 +35,7 @@ impl BlockedAddrs {
     pub fn new(block_duration: Duration) -> BlockedAddrs {
         BlockedAddrs(Arc::new(BlockAddrsInner {
             block_duration,
-            inner: Mutex::new(HashMap::new())
+            inner: Mutex::new(HashMap::new()),
         }))
     }
 
@@ -52,8 +52,8 @@ impl BlockedAddrs {
         let mut map = self.0.inner.lock().unwrap();
 
         let (reason, time) = match map.get(addr) {
-            Some(&(reason,time)) => (reason, time),
-            None => return None
+            Some(&(reason, time)) => (reason, time),
+            None => return None,
         };
 
         if time + self.0.block_duration < Instant::now() {
@@ -63,4 +63,4 @@ impl BlockedAddrs {
             Some(reason)
         }
     }
-}
+}

backend/telemetry_shard/src/main.rs

@@ -16,14 +16,15 @@
 
 #[warn(missing_docs)]
 mod aggregator;
+mod blocked_addrs;
 mod connection;
 mod json_message;
 mod real_ip;
-mod blocked_addrs;
 
 use std::{collections::HashSet, net::IpAddr, time::Duration};
 
 use aggregator::{Aggregator, FromWebsocket};
+use blocked_addrs::BlockedAddrs;
 use common::byte_size::ByteSize;
 use common::http_utils;
 use common::node_message;
@@ -33,7 +34,6 @@ use http::Uri;
 use hyper::{Method, Response};
 use simple_logger::SimpleLogger;
 use structopt::StructOpt;
-use blocked_addrs::BlockedAddrs;
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
@@ -79,7 +79,7 @@ struct Opts {
     /// How many seconds is a "/feed" connection that violates the '--max-node-data-per-second'
     /// value prevented from reconnecting to this shard for, in seconds.
     #[structopt(long, default_value = "600")]
-    node_block_seconds: u64
+    node_block_seconds: u64,
 }
 
 #[tokio::main]
@@ -118,10 +118,7 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
                     let real_addr = real_ip::real_ip(addr, req.headers());
 
                     if let Some(reason) = block_list.blocked_reason(&real_addr) {
-                        return Ok(Response::builder()
-                            .status(403)
-                            .body(reason.into())
-                            .unwrap())
+                        return Ok(Response::builder().status(403).body(reason.into()).unwrap());
                     }
 
                     Ok(http_utils::upgrade_to_websocket(
@@ -136,7 +133,7 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
                                     tx_to_aggregator,
                                     max_nodes_per_connection,
                                     bytes_per_second,
-                                    block_list
+                                    block_list,
                                 )
                                 .await;
                             log::info!("Closing /submit connection from {:?}", addr);
@@ -167,7 +164,7 @@ async fn handle_node_websocket_connection<S>(
     mut tx_to_aggregator: S,
     max_nodes_per_connection: usize,
     bytes_per_second: ByteSize,
-    block_list: BlockedAddrs
+    block_list: BlockedAddrs,
 ) -> (S, http_utils::WsSender)
 where
     S: futures::Sink<FromWebsocket, Error = anyhow::Error> + Unpin + Send + 'static,