From 653b7c5e6b28250b8747ae303aa38240766951b1 Mon Sep 17 00:00:00 2001
From: David Hill <dhill@mindcry.org>
Date: Tue, 6 Nov 2018 19:05:28 -0500
Subject: [PATCH] [config]: add --maxsameip to limit # of conns to same IP

---
 config.go |  3 +++
 doc.go    |  2 ++
 server.go | 35 +++++++++++++++++++++++++++++++++--
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/config.go b/config.go
index 5a29748764..872aea23fe 100644
--- a/config.go
+++ b/config.go
@@ -39,6 +39,7 @@ const (
 	defaultLogLevel              = "info"
 	defaultLogDirname            = "logs"
 	defaultLogFilename           = "dcrd.log"
+	defaultMaxSameIP             = 5
 	defaultMaxPeers              = 125
 	defaultBanDuration           = time.Hour * 24
 	defaultBanThreshold          = 100
@@ -100,6 +101,7 @@ type config struct {
 	ConnectPeers         []string      `long:"connect" description:"Connect only to the specified peers at startup"`
 	DisableListen        bool          `long:"nolisten" description:"Disable listening for incoming connections -- NOTE: Listening is automatically disabled if the --connect or --proxy options are used without also specifying listen interfaces via --listen"`
 	Listeners            []string      `long:"listen" description:"Add an interface/port to listen for connections (default all interfaces port: 9108, testnet: 19108)"`
+	MaxSameIP            int           `long:"maxsameip" description:"Max number of connections with the same IP -- 0 to disable"`
 	MaxPeers             int           `long:"maxpeers" description:"Max number of inbound and outbound peers"`
 	DisableBanning       bool          `long:"nobanning" description:"Disable banning of misbehaving peers"`
 	BanDuration          time.Duration `long:"banduration" description:"How long to ban misbehaving peers.  Valid time units are {s, m, h}.  Minimum 1 second"`
@@ -434,6 +436,7 @@ func loadConfig() (*config, []string, error) {
 		HomeDir:              defaultHomeDir,
 		ConfigFile:           defaultConfigFile,
 		DebugLevel:           defaultLogLevel,
+		MaxSameIP:            defaultMaxSameIP,
 		MaxPeers:             defaultMaxPeers,
 		BanDuration:          defaultBanDuration,
 		BanThreshold:         defaultBanThreshold,
diff --git a/doc.go b/doc.go
index ebe4bca3a6..316dc28872 100644
--- a/doc.go
+++ b/doc.go
@@ -35,6 +35,8 @@ Application Options:
                             listen interfaces via --listen
       --listen=             Add an interface/port to listen for connections
                             (default all interfaces port: 9108, testnet: 19108)
+      --maxsameip=          Max number of connections with the same IP -- 0 to
+                            disable (default: 5)
       --maxpeers=           Max number of inbound and outbound peers (125)
       --nobanning           Disable banning of misbehaving peers
       --banduration=        How long to ban misbehaving peers.  Valid time units
diff --git a/server.go b/server.go
index ad2393379a..f97dbef8c8 100644
--- a/server.go
+++ b/server.go
@@ -122,6 +122,27 @@ type peerState struct {
 	outboundGroups  map[string]int
 }
 
+// ConnectionsWithIP returns the number of connections with the given IP.
+func (ps *peerState) ConnectionsWithIP(ip net.IP) int {
+	var total int
+	for _, p := range ps.inboundPeers {
+		if ip.Equal(p.NA().IP) {
+			total++
+		}
+	}
+	for _, p := range ps.outboundPeers {
+		if ip.Equal(p.NA().IP) {
+			total++
+		}
+	}
+	for _, p := range ps.persistentPeers {
+		if ip.Equal(p.NA().IP) {
+			total++
+		}
+	}
+	return total
+}
+
 // Count returns the count of all known peers.
 func (ps *peerState) Count() int {
 	return len(ps.inboundPeers) + len(ps.outboundPeers) +
@@ -1277,11 +1298,21 @@ func (s *server) handleAddPeerMsg(state *peerState, sp *serverPeer) bool {
 		delete(state.banned, host)
 	}
 
-	// TODO: Check for max peers from a single IP.
+	// Limit max number of connections from a single IP.  However, allow
+	// whitelisted inbound peers and localhost connections regardless.
+	isInboundWhitelisted := sp.isWhitelisted && sp.Inbound()
+	peerIP := sp.NA().IP
+	if cfg.MaxSameIP > 0 && !isInboundWhitelisted && !peerIP.IsLoopback() &&
+		state.ConnectionsWithIP(peerIP)+1 > cfg.MaxSameIP {
+		srvrLog.Infof("Max connections with %s reached [%d] - "+
+			"disconnecting peer", sp, cfg.MaxSameIP)
+		sp.Disconnect()
+		return false
+	}
 
 	// Limit max number of total peers.  However, allow whitelisted inbound
 	// peers regardless.
-	if state.Count() >= cfg.MaxPeers && !(sp.Inbound() && sp.isWhitelisted) {
+	if state.Count()+1 > cfg.MaxPeers && !isInboundWhitelisted {
 		srvrLog.Infof("Max peers reached [%d] - disconnecting peer %s",
 			cfg.MaxPeers, sp)
 		sp.Disconnect()