From ec7ae91290ccadb90d797972d9c87760aca33182 Mon Sep 17 00:00:00 2001
From: David Palm <dvdplm@gmail.com>
Date: Wed, 24 Mar 2021 20:26:01 +0100
Subject: [PATCH] Add --denylist ChainIDontWant OtherChainIDontWant Add --log
 error|warn|info|debug|trace

---
 backend/src/aggregator.rs    | 11 ++++++--
 backend/src/chain.rs         |  2 +-
 backend/src/main.rs          | 50 ++++++++++++++++++++++++++++++++----
 backend/src/util/location.rs |  2 +-
 4 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/backend/src/aggregator.rs b/backend/src/aggregator.rs
index 29cdc36..01776e1 100644
--- a/backend/src/aggregator.rs
+++ b/backend/src/aggregator.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use actix::prelude::*;
 
 use crate::node::connector::Initialize;
@@ -14,6 +14,8 @@ pub struct Aggregator {
     chains: DenseMap<ChainEntry>,
     feeds: DenseMap<Addr<FeedConnector>>,
     serializer: FeedMessageSerializer,
+    /// Denylist for networks we do not want to allow connecting.
+    denylist: HashSet<String>
 }
 
 pub struct ChainEntry {
@@ -24,13 +26,14 @@ pub struct ChainEntry {
 }
 
 impl Aggregator {
-    pub fn new() -> Self {
+    pub fn new(denylist: HashSet<String>) -> Self {
         Aggregator {
             labels: HashMap::new(),
             networks: HashMap::new(),
             chains: DenseMap::new(),
             feeds: DenseMap::new(),
             serializer: FeedMessageSerializer::new(),
+            denylist,
         }
     }
 
@@ -176,6 +179,10 @@ impl Handler<AddNode> for Aggregator {
     type Result = ();
 
     fn handle(&mut self, msg: AddNode, ctx: &mut Self::Context) {
+        if self.denylist.contains(&*msg.node.chain) {
+            log::debug!(target: "Aggregator::AddNode", "'{}' is on the denylist.", msg.node.chain);
+            return;
+        }
         let AddNode { node, conn_id, rec } = msg;
 
         let cid = self.lazy_chain(&node.chain, &None, ctx);
diff --git a/backend/src/chain.rs b/backend/src/chain.rs
index abd85ff..38d5fb3 100644
--- a/backend/src/chain.rs
+++ b/backend/src/chain.rs
@@ -18,7 +18,7 @@ pub type Label = Arc<str>;
 
 pub struct Chain {
     cid: ChainId,
-    /// Who to inform if we Chain drops itself
+    /// Who to inform if the Chain drops itself
     aggregator: Addr<Aggregator>,
     /// Label of this chain, along with count of nodes that use this label
     label: (Label, usize),
diff --git a/backend/src/main.rs b/backend/src/main.rs
index 5ec0e59..68e44fb 100644
--- a/backend/src/main.rs
+++ b/backend/src/main.rs
@@ -1,4 +1,6 @@
 use std::net::Ipv4Addr;
+use std::collections::HashSet;
+use std::iter::FromIterator;
 
 use actix::prelude::*;
 use actix_http::ws::Codec;
@@ -25,16 +27,52 @@ const AUTHORS: &'static str = env!("CARGO_PKG_AUTHORS");
 const NAME: &'static str = "Substrate Telemetry Backend";
 const ABOUT: &'static str = "This is the Telemetry Backend that injects and provide the data sent by Substrate/Polkadot nodes";
 
-#[derive(Clap)]
+#[derive(Clap, Debug)]
 #[clap(name = NAME, version = VERSION, author = AUTHORS, about = ABOUT)]
 struct Opts {
     #[clap(
         short = 'l',
         long = "listen",
         default_value = "127.0.0.1:8000",
-        about = "This is the socket address Telemetry is listening to. This is restricted localhost (127.0.0.1) by default and should be fine for most use cases. If you are using Telemetry in a container, you likely want to set this to '0.0.0.0:8000'"
+        about = "This is the socket address Telemetry is listening to. This is restricted to localhost (127.0.0.1) by default and should be fine for most use cases. If you are using Telemetry in a container, you likely want to set this to '0.0.0.0:8000'"
     )]
     socket: std::net::SocketAddr,
+    #[clap(
+        required = false,
+        long = "denylist",
+        default_value = "Earth",
+        about = "Space delimited list of chains that are not allowed to connect to telemetry. Case sensitive."
+    )]
+    denylist: Vec<String>,
+    #[clap(
+        arg_enum,
+        required = false,
+        long = "log",
+        default_value = "info",
+        about = "Log level. Defaults to 'info'. Valid values are: error, warn, info, debug and trace"
+    )]
+    log_level: LogLevel,
+}
+
+#[derive(Clap, Debug, PartialEq)]
+enum LogLevel {
+    Error,
+    Warn,
+    Info,
+    Debug,
+    Trace,
+}
+
+impl Into<log::LevelFilter> for &LogLevel {
+    fn into(self) -> log::LevelFilter {
+        match self {
+            LogLevel::Error => log::LevelFilter::Error,
+            LogLevel::Warn => log::LevelFilter::Warn,
+            LogLevel::Info => log::LevelFilter::Info,
+            LogLevel::Debug => log::LevelFilter::Debug,
+            LogLevel::Trace => log::LevelFilter::Trace,
+        }
+    }
 }
 
 /// Entry point for connecting nodes
@@ -127,10 +165,12 @@ async fn health(aggregator: web::Data<Addr<Aggregator>>) -> Result<HttpResponse,
 /// This can be changed using the `PORT` and `BIND` ENV variables.
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
-    SimpleLogger::new().with_level(log::LevelFilter::Info).init().expect("Must be able to start a logger");
-
     let opts: Opts = Opts::parse();
-    let aggregator = Aggregator::new().start();
+    let log_level = &opts.log_level;
+    SimpleLogger::new().with_level(log_level.into()).init().expect("Must be able to start a logger");
+
+    let denylist = HashSet::from_iter(opts.denylist);
+    let aggregator = Aggregator::new(denylist).start();
     let factory = LocatorFactory::new();
     let locator = SyncArbiter::start(4, move || factory.create());
 
diff --git a/backend/src/util/location.rs b/backend/src/util/location.rs
index d4bad3d..c7ff7df 100644
--- a/backend/src/util/location.rs
+++ b/backend/src/util/location.rs
@@ -176,4 +176,4 @@ mod tests {
 
         assert!(location.is_none());
     }
-}
\ No newline at end of file
+}