From 90f28314d2239961b4cc7728bf6416d4aedd4adf Mon Sep 17 00:00:00 2001
From: "John L. Jegutanis" <erasmospunk@gmail.com>
Date: Thu, 5 Apr 2018 00:49:56 +0300
Subject: [PATCH] Add DROP_CLIENT env variable (#432)

This will disconnect any client based on their version string,
using a regular expression.
Useful for dropping buggy/unsupported clients.
---
 docs/environment.rst     |  6 ++++++
 lib/env_base.py          | 11 +++++++++++
 server/controller.py     |  3 +++
 server/env.py            |  2 ++
 server/session.py        |  4 ++++
 tests/server/test_env.py | 11 +++++++++++
 6 files changed, 37 insertions(+)

diff --git a/docs/environment.rst b/docs/environment.rst
index 1a0cdec..1342e57 100644
--- a/docs/environment.rst
+++ b/docs/environment.rst
@@ -201,6 +201,12 @@ These environment variables are optional:
 
   If you are not sure what this means leave it unset.
 
+.. envvar:: DROP_CLIENT
+
+  Set a regular expression to disconnect any client based on their
+  version string. For example to drop versions from 1.0 to 1.2 use
+  the regex ``1\.[0-2]\.\d+``.
+
 
 Resource Usage Limits
 =====================
diff --git a/lib/env_base.py b/lib/env_base.py
index a9e9dd5..a94eabb 100644
--- a/lib/env_base.py
+++ b/lib/env_base.py
@@ -53,6 +53,17 @@ class EnvBase(lib_util.LoggedClass):
             raise cls.Error('cannot convert envvar {} value {} to an integer'
                             .format(envvar, value))
 
+    @classmethod
+    def custom(cls, envvar, default, parse):
+        value = environ.get(envvar)
+        if value is None:
+            return default
+        try:
+            return parse(value)
+        except Exception as e:
+            raise cls.Error('cannot parse envvar {} value {}'
+                            .format(envvar, value)) from e
+
     @classmethod
     def obsolete(cls, envvars):
         bad = [envvar for envvar in envvars if environ.get(envvar)]
diff --git a/server/controller.py b/server/controller.py
index f6a9b5e..46b29dd 100644
--- a/server/controller.py
+++ b/server/controller.py
@@ -254,6 +254,9 @@ class Controller(ServerBase):
         self.logger.info('max subscriptions per session: {:,d}'
                          .format(self.env.max_session_subs))
         self.logger.info('bands: {}'.format(self.bands))
+        if self.env.drop_client is not None:
+            self.logger.info('drop clients matching: {}'
+                             .format(self.env.drop_client.pattern))
         await self.start_external_servers()
 
     async def start_external_servers(self):
diff --git a/server/env.py b/server/env.py
index ff74f24..cd7966d 100644
--- a/server/env.py
+++ b/server/env.py
@@ -8,6 +8,7 @@
 '''Class for handling environment configuration and defaults.'''
 
 
+import re
 import resource
 from collections import namedtuple
 from ipaddress import ip_address
@@ -67,6 +68,7 @@ class Env(EnvBase):
         self.max_session_subs = self.integer('MAX_SESSION_SUBS', 50000)
         self.bandwidth_limit = self.integer('BANDWIDTH_LIMIT', 2000000)
         self.session_timeout = self.integer('SESSION_TIMEOUT', 600)
+        self.drop_client = self.custom("DROP_CLIENT", None, re.compile)
 
         # Identities
         clearnet_identity = self.clearnet_identity()
diff --git a/server/session.py b/server/session.py
index a26af3e..18ff5aa 100644
--- a/server/session.py
+++ b/server/session.py
@@ -358,6 +358,10 @@ class ElectrumX(SessionBase):
         protocol_version: the protocol version spoken by the client
         '''
         if client_name:
+            if self.env.drop_client is not None and \
+                    self.env.drop_client.match(client_name):
+                raise RPCError('unsupported client: {}'
+                               .format(client_name), JSONRPC.FATAL_ERROR)
             self.client = str(client_name)[:17]
             try:
                 self.client_version = tuple(int(part) for part
diff --git a/tests/server/test_env.py b/tests/server/test_env.py
index 456e8c1..b18a6b5 100644
--- a/tests/server/test_env.py
+++ b/tests/server/test_env.py
@@ -2,6 +2,7 @@
 
 import os
 import random
+import re
 
 import pytest
 
@@ -340,3 +341,13 @@ def test_tor_identity():
     assert ident.host == tor_host
     assert ident.tcp_port == 234
     assert ident.ssl_port == 432
+
+def test_ban_versions():
+    e = Env()
+    assert e.drop_client is None
+    ban_re = '1\.[0-2]\.\d+?[_\w]*'
+    os.environ['DROP_CLIENT'] = ban_re
+    e = Env()
+    assert e.drop_client == re.compile(ban_re)
+    assert e.drop_client.match("1.2.3_buggy_client")
+    assert e.drop_client.match("1.3.0_good_client") is None