From 81947d796b39c1a7a46c12c421271a9b912899b9 Mon Sep 17 00:00:00 2001
From: Neil Booth <kyuupichan@gmail.com>
Date: Wed, 11 Oct 2017 13:45:37 +0900
Subject: [PATCH] Enhance PEER_DISCOVERY environment variable

It is now tri-state - ON, OFF or SELF.

Fixes #287
---
 docs/ENVIRONMENT.rst     | 12 +++++++++---
 server/env.py            | 14 +++++++++++++-
 server/peers.py          | 13 +++++++------
 tests/server/test_env.py | 12 +++++++++++-
 4 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/docs/ENVIRONMENT.rst b/docs/ENVIRONMENT.rst
index d1f917e06..3553348ae 100644
--- a/docs/ENVIRONMENT.rst
+++ b/docs/ENVIRONMENT.rst
@@ -270,9 +270,15 @@ some of this.
 
 * **PEER_DISCOVERY**
 
-  If not defined, or non-empty, ElectrumX will occasionally connect to
-  and verify its network of peer servers.  Set to empty to disable
-  peer discovery.
+  This environment variable is case-insensitive and defaults to `on`.
+
+  If `on`, ElectrumX will occasionally connect to and verify its
+  network of peer servers.
+
+  If `off`, peer discovery is disabled and a hard-coded default list
+  of servers will be read in and served.  If set to `self` then peer
+  discovery is disabled and the server will only return itself in the
+  peers list.
 
 * **PEER_ANNOUNCE**
 
diff --git a/server/env.py b/server/env.py
index 145ade166..4ce80c69c 100644
--- a/server/env.py
+++ b/server/env.py
@@ -23,6 +23,9 @@
 class Env(EnvBase):
     '''Wraps environment configuration.'''
 
+    # Peer discovery
+    PD_OFF, PD_SELF, PD_ON = range(3)
+
     def __init__(self):
         super().__init__()
         self.obsolete(['UTXO_MB', 'HIST_MB', 'NETWORK'])
@@ -49,7 +52,7 @@ def __init__(self):
         self.anon_logs = self.boolean('ANON_LOGS', False)
         self.log_sessions = self.integer('LOG_SESSIONS', 3600)
         # Peer discovery
-        self.peer_discovery = self.boolean('PEER_DISCOVERY', True)
+        self.peer_discovery = self.peer_discovery_enum()
         self.peer_announce = self.boolean('PEER_ANNOUNCE', True)
         self.force_proxy = self.boolean('FORCE_PROXY', False)
         self.tor_proxy_host = self.default('TOR_PROXY_HOST', 'localhost')
@@ -143,3 +146,12 @@ def port(port_kind):
             ssl_port,
             '_tor',
         )
+
+    def peer_discovery_enum(self):
+        pd = self.default('PEER_DISCOVERY', 'on').strip().lower()
+        if pd in ('off', ''):
+            return self.PD_OFF
+        elif pd == 'self':
+            return self.PD_SELF
+        else:
+            return self.PD_ON
diff --git a/server/peers.py b/server/peers.py
index 95781438d..ccb607684 100644
--- a/server/peers.py
+++ b/server/peers.py
@@ -424,12 +424,13 @@ def on_peers_subscribe(self, is_tor):
     def import_peers(self):
         '''Import hard-coded peers from a file or the coin defaults.'''
         self.add_peers(self.myselves)
-        coin_peers = self.env.coin.PEERS
 
-        # Add the hard-coded ones
-        peers = [Peer.from_real_name(real_name, 'coins.py')
-                 for real_name in coin_peers]
-        self.add_peers(peers, limit=None)
+        # Add the hard-coded ones unless only returning self
+        if self.env.peer_discovery != self.env.PD_SELF:
+            coin_peers = self.env.coin.PEERS
+            peers = [Peer.from_real_name(real_name, 'coins.py')
+                     for real_name in coin_peers]
+            self.add_peers(peers, limit=None)
 
     def connect_to_irc(self):
         '''Connect to IRC if not disabled.'''
@@ -459,7 +460,7 @@ async def main_loop(self):
           3) Retrying old peers at regular intervals.
         '''
         self.connect_to_irc()
-        if not self.env.peer_discovery:
+        if self.env.peer_discovery != self.env.PD_ON:
             self.logger.info('peer discovery is disabled')
             return
 
diff --git a/tests/server/test_env.py b/tests/server/test_env.py
index 9d101f92e..1bc868376 100644
--- a/tests/server/test_env.py
+++ b/tests/server/test_env.py
@@ -199,7 +199,17 @@ def test_ANON_LOGS():
     assert_boolean('ANON_LOGS', 'anon_logs', False)
 
 def test_PEER_DISCOVERY():
-    assert_boolean('PEER_DISCOVERY', 'peer_discovery', True)
+    e = Env()
+    assert e.peer_discovery == Env.PD_ON
+    os.environ['PEER_DISCOVERY'] = ' '
+    e = Env()
+    assert e.peer_discovery == Env.PD_OFF
+    os.environ['PEER_DISCOVERY'] = 'ON'
+    e = Env()
+    assert e.peer_discovery == Env.PD_ON
+    os.environ['PEER_DISCOVERY'] = 'self'
+    e = Env()
+    assert e.peer_discovery == Env.PD_SELF
 
 def test_PEER_ANNOUNCE():
     assert_boolean('PEER_ANNOUNCE', 'peer_announce', True)