From 1a922ecb2926e55410b7c515949474fa412abc9c Mon Sep 17 00:00:00 2001
From: Izalia Mae <izalia@barkshark.tk>
Date: Tue, 21 May 2019 12:29:55 -0400
Subject: [PATCH] add whitelist feature

---
 relay.yaml.example |  5 ++++-
 relay/actor.py     | 12 ++++++++++--
 relay/database.py  |  7 ++++++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/relay.yaml.example b/relay.yaml.example
index 4f0b69c..1a0da3d 100644
--- a/relay.yaml.example
+++ b/relay.yaml.example
@@ -17,4 +17,7 @@ ap:
   blocked_instances:
   - 'bad-instance.example.com'
   - 'another-bad-instance.example.com'
-  
+  whitelist_enabled: false
+  whitelist:
+  - 'good-instance.example.com'
+  - 'another.good-instance.example.com'
\ No newline at end of file
diff --git a/relay/actor.py b/relay/actor.py
index b2117fd..f2d1379 100644
--- a/relay/actor.py
+++ b/relay/actor.py
@@ -35,7 +35,7 @@ from . import app, CONFIG
 from .remote_actor import fetch_actor
 
 
-AP_CONFIG = CONFIG.get('ap', {'host': 'localhost','blocked_instances':[]})
+AP_CONFIG = CONFIG.get('ap', {'host': 'localhost','blocked_instances':[], 'whitelist_enabled': False, 'whitelist': []})
 CACHE_SIZE = CONFIG.get('cache-size', 16384)
 
 
@@ -106,10 +106,15 @@ async def push_message_to_actor(actor, message, our_key_id):
 
 async def follow_remote_actor(actor_uri):
     actor = await fetch_actor(actor_uri)
+    
     if not actor:
         logging.info('failed to fetch actor at: %r', actor_uri)
         return
 
+    if AP_CONFIG['whitelist_enabled'] is True and urlsplit(actor_uri).hostname not in AP_CONFIG['whitelist']:
+        logging.info('refusing to follow non-whitelisted actor: %r', actor_uri)
+        return
+
     logging.info('following: %r', actor_uri)
 
     message = {
@@ -294,7 +299,10 @@ async def inbox(request):
     if 'actor' not in data or not request['validated']:
         raise aiohttp.web.HTTPUnauthorized(body='access denied', content_type='text/plain')
 
-    if data['type'] != 'Follow' and 'https://{}/inbox'.format(instance) not in DATABASE['relay-list']:
+    elif data['type'] != 'Follow' and 'https://{}/inbox'.format(instance) not in DATABASE['relay-list']:
+        raise aiohttp.web.HTTPUnauthorized(body='access denied', content_type='text/plain')
+
+    elif AP_CONFIG['whitelist_enabled'] is True and instance not in AP_CONFIG['whitelist']:
         raise aiohttp.web.HTTPUnauthorized(body='access denied', content_type='text/plain')
 
     actor = await fetch_actor(data["actor"])
diff --git a/relay/database.py b/relay/database.py
index 388fe08..6a99d49 100644
--- a/relay/database.py
+++ b/relay/database.py
@@ -5,6 +5,7 @@ import simplejson as json
 
 
 from . import CONFIG
+AP_CONFIG = CONFIG.get('ap', {'blocked_instances':[], 'whitelist_enabled': False, 'whitelist': []})
 
 
 try:
@@ -16,7 +17,11 @@ except:
 
 following = DATABASE.get('relay-list', [])
 for inbox in following:
-    if urllib.parse.urlsplit(inbox).hostname in CONFIG['ap']['blocked_instances']:
+    if urllib.parse.urlsplit(inbox).hostname in AP_CONFIG['blocked_instances']:
+        following.remove(inbox)
+        DATABASE['relay-list'] = following
+
+    elif AP_CONFIG['whitelist_enabled'] is True and urllib.parse.urlsplit(inbox).hostname not in AP_CONFIG['whitelist']:
         following.remove(inbox)
         DATABASE['relay-list'] = following