From 8ae4322f2e3f191601f82225b1c184336900097e Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@mit.edu>
Date: Tue, 9 Jan 2018 19:12:46 -0800
Subject: [PATCH] server: adds hybrid black/whitelist filtering to user agents

---
 server.go | 68 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 1 deletion(-)

diff --git a/server.go b/server.go
index 15d77d2ec37..65ce1426e92 100644
--- a/server.go
+++ b/server.go
@@ -234,6 +234,14 @@ type server struct {
 	// The fee estimator keeps track of how long transactions are left in
 	// the mempool before they are mined into blocks.
 	feeEstimator *mempool.FeeEstimator
+
+	// agentBlacklist is a list of blacklisted substrings by which to filter
+	// user agents.
+	agentBlacklist []string
+
+	// agentWhitelist is a list of whitelisted user agent substrings, no
+	// whitelisting will be applied if the list is empty or nil.
+	agentWhitelist []string
 }
 
 // serverPeer extends the peer to maintain state shared by the server and
@@ -1368,6 +1376,12 @@ func (s *server) handleAddPeerMsg(state *peerState, sp *serverPeer) bool {
 		return false
 	}
 
+	// Disconnect peers with unwanted user agents.
+	if sp.HasUndesiredUserAgent(s.agentBlacklist, s.agentWhitelist) {
+		sp.Disconnect()
+		return false
+	}
+
 	// Ignore new peers if we're shutting down.
 	if atomic.LoadInt32(&s.shutdown) != 0 {
 		srvrLog.Infof("New peer %s ignored - server is shutting down", sp)
@@ -2319,7 +2333,10 @@ func setupRPCListeners() ([]net.Listener, error) {
 // newServer returns a new btcd server configured to listen on addr for the
 // bitcoin network type specified by chainParams.  Use start to begin accepting
 // connections from peers.
-func newServer(listenAddrs []string, db database.DB, chainParams *chaincfg.Params, interrupt <-chan struct{}) (*server, error) {
+func newServer(listenAddrs, agentBlacklist, agentWhitelist []string,
+	db database.DB, chainParams *chaincfg.Params,
+	interrupt <-chan struct{}) (*server, error) {
+
 	services := defaultServices
 	if cfg.NoPeerBloomFilters {
 		services &^= wire.SFNodeBloom
@@ -2343,6 +2360,9 @@ func newServer(listenAddrs []string, db database.DB, chainParams *chaincfg.Param
 		}
 	}
 
+	srvrLog.Infof("Using agent blacklist: %s and whitelist: %s",
+		agentBlacklist, agentWhitelist)
+
 	s := server{
 		chainParams:          chainParams,
 		addrManager:          amgr,
@@ -2361,6 +2381,8 @@ func newServer(listenAddrs []string, db database.DB, chainParams *chaincfg.Param
 		services:             services,
 		sigCache:             txscript.NewSigCache(cfg.SigCacheMaxSize),
 		hashCache:            txscript.NewHashCache(cfg.SigCacheMaxSize),
+		agentBlacklist:       agentBlacklist,
+		agentWhitelist:       agentWhitelist,
 	}
 
 	// Create the transaction and address indexes if needed.
@@ -2913,3 +2935,47 @@ func mergeCheckpoints(defaultCheckpoints, additional []chaincfg.Checkpoint) []ch
 	sort.Sort(checkpointSorter(checkpoints))
 	return checkpoints
 }
+
+// HasUndesiredUserAgent determines whether the server should continue to pursue
+// a connection with this peer based on its advertised user agent. It performs
+// the following steps:
+// 1) Reject the peer if it contains a blacklisted agent.
+// 2) If no whitelist is provided, accept all user agents.
+// 3) Accept the peer if it contains a whitelisted agent.
+// 4) Reject all other peers.
+func (sp *serverPeer) HasUndesiredUserAgent(blacklistedAgents,
+	whitelistedAgents []string) bool {
+
+	agent := sp.UserAgent()
+
+	// First, if peer's user agent contains any blacklisted substring, we
+	// will ignore the connection request.
+	for _, blacklistedAgent := range blacklistedAgents {
+		if strings.Contains(agent, blacklistedAgent) {
+			srvrLog.Debugf("Ignoring peer %s, user agent "+
+				"contains blacklisted user agent: %s", sp,
+				agent)
+			return true
+		}
+	}
+
+	// If no whitelist is provided, we will accept all user agents.
+	if len(whitelistedAgents) == 0 {
+		return false
+	}
+
+	// Peer's user agent passed blacklist. Now check to see if it contains
+	// one of our whitelisted user agents, if so accept.
+	for _, whitelistedAgent := range whitelistedAgents {
+		if strings.Contains(agent, whitelistedAgent) {
+			return false
+		}
+	}
+
+	// Otherwise, the peer's user agent was not included in our whitelist.
+	// Ignore just in case it could stall the initial block download.
+	srvrLog.Debugf("Ignoring peer %s, user agent: %s not found in "+
+		"whitelist", sp, agent)
+
+	return true
+}