From 555dab4fd681a719545628839acb5874c8ba2247 Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Sat, 11 Jul 2015 01:41:08 +0000
Subject: [PATCH] limit number of allowed connections per ip

---
 p2p/peer_set.go      |  2 --
 p2p/peer_set_test.go |  2 --
 p2p/pex_reactor.go   |  3 +++
 p2p/switch.go        | 21 ++++++++++++++++++---
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/p2p/peer_set.go b/p2p/peer_set.go
index 48ef71a64..4b43defec 100644
--- a/p2p/peer_set.go
+++ b/p2p/peer_set.go
@@ -35,8 +35,6 @@ func NewPeerSet() *PeerSet {
 }
 
 // Returns false if peer with key (uuid) is already in set.
-// TODO: we may want to do other things like restrict how many peers
-// we're willing to connect to behind a single IP
 func (ps *PeerSet) Add(peer *Peer) bool {
 	ps.mtx.Lock()
 	defer ps.mtx.Unlock()
diff --git a/p2p/peer_set_test.go b/p2p/peer_set_test.go
index b0d5dff89..d322c13ac 100644
--- a/p2p/peer_set_test.go
+++ b/p2p/peer_set_test.go
@@ -1,11 +1,9 @@
 package p2p
 
 import (
-	"math/rand"
 	"testing"
 
 	"github.com/tendermint/tendermint/Godeps/_workspace/src/code.google.com/p/go-uuid/uuid"
-	. "github.com/tendermint/tendermint/common"
 )
 
 // Returns an empty dummy peer
diff --git a/p2p/pex_reactor.go b/p2p/pex_reactor.go
index 5376e2527..abe25b7b1 100644
--- a/p2p/pex_reactor.go
+++ b/p2p/pex_reactor.go
@@ -209,6 +209,9 @@ func (pexR *PEXReactor) ensurePeers() {
 			}
 		}(item.(*NetAddress))
 	}
+
+	// TODO: if no addresses to dial, we should send a pexRequest to a random peer
+	// so we can get more peers
 }
 
 // implements events.Eventable
diff --git a/p2p/switch.go b/p2p/switch.go
index a7c683089..17dd17413 100644
--- a/p2p/switch.go
+++ b/p2p/switch.go
@@ -53,10 +53,12 @@ type Switch struct {
 
 var (
 	ErrSwitchDuplicatePeer = errors.New("Duplicate peer")
+	ErrSwitchMaxPeersPerIP = errors.New("IP has too many peers")
 )
 
 const (
 	peerDialTimeoutSeconds = 3
+	maxPeersPerIP          = 3
 )
 
 func NewSwitch() *Switch {
@@ -186,10 +188,21 @@ func (sw *Switch) AddPeerWithConnection(conn net.Conn, outbound bool) (*Peer, er
 	}
 	peer := newPeer(conn, peerNodeInfo, outbound, sw.reactorsByCh, sw.chDescs, sw.StopPeerForError)
 
+	// restrict the number of peers we're willing to connect to behind a single IP
+	var numPeersOnThisIP int
+	peers := sw.Peers().List()
+	for _, p := range peers {
+		if p.Host == peerNodeInfo.Host {
+			numPeersOnThisIP += 1
+		}
+	}
+	if numPeersOnThisIP == maxPeersPerIP {
+		log.Info("Ignoring peer as we have the max allowed for that IP", "IP", peerNodeInfo.Host, "peer", peer, "max", maxPeersPerIP)
+		return nil, ErrSwitchMaxPeersPerIP
+	}
+
 	// Add the peer to .peers
-	if sw.peers.Add(peer) {
-		log.Info("Added peer", "peer", peer)
-	} else {
+	if !sw.peers.Add(peer) {
 		log.Info("Ignoring duplicate peer", "peer", peer)
 		peer.stop() // will also close conn
 		return nil, ErrSwitchDuplicatePeer
@@ -198,6 +211,8 @@ func (sw *Switch) AddPeerWithConnection(conn net.Conn, outbound bool) (*Peer, er
 	if atomic.LoadUint32(&sw.running) == 1 {
 		sw.startInitPeer(peer)
 	}
+
+	log.Info("Added peer", "peer", peer)
 	return peer, nil
 }