limit num connections per ip range

p2p/peer_set_test.go

@@ -1,15 +1,22 @@
 package p2p
 
 import (
+	"math/rand"
+	"strings"
 	"testing"
 
 	"github.com/tendermint/tendermint/Godeps/_workspace/src/code.google.com/p/go-uuid/uuid"
+	. "github.com/tendermint/tendermint/common"
+	"github.com/tendermint/tendermint/types"
 )
 
 // Returns an empty dummy peer
 func randPeer() *Peer {
 	return &Peer{
 		Key: uuid.New(),
+		NodeInfo: &types.NodeInfo{
+			Host: Fmt("%v.%v.%v.%v", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256),
+		},
 	}
 }
 
@@ -18,7 +25,7 @@ func TestAddRemoveOne(t *testing.T) {
 
 	peer := randPeer()
 	added := peerSet.Add(peer)
-	if !added {
+	if err := added; err != nil {
 		t.Errorf("Failed to add new peer")
 	}
 	if peerSet.Size() != 1 {
@@ -38,10 +45,11 @@ func TestAddRemoveMany(t *testing.T) {
 	peerSet := NewPeerSet()
 
 	peers := []*Peer{}
-	for i := 0; i < 100; i++ {
+	N := 100
+	maxPeersPerIPRange = [4]int{N, N, N, N}
+	for i := 0; i < N; i++ {
 		peer := randPeer()
-		added := peerSet.Add(peer)
-		if !added {
+		if err := peerSet.Add(peer); err != nil {
 			t.Errorf("Failed to add new peer")
 		}
 		if peerSet.Size() != i+1 {
@@ -60,3 +68,105 @@ func TestAddRemoveMany(t *testing.T) {
 		}
 	}
 }
+
+func newPeerInIPRange(ipBytes ...string) *Peer {
+	ips := make([]string, 4)
+	for i, ipByte := range ipBytes {
+		ips[i] = ipByte
+	}
+
+	for i := len(ipBytes); i < 4; i++ {
+		ips[i] = Fmt("%v", rand.Int()%256)
+	}
+
+	ipS := strings.Join(ips, ".")
+	return &Peer{
+		Key: uuid.New(),
+		NodeInfo: &types.NodeInfo{
+			Host: ipS,
+		},
+	}
+}
+
+func TestIPRanges(t *testing.T) {
+	peerSet := NewPeerSet()
+
+	// test  /8
+	maxPeersPerIPRange = [4]int{2, 2, 2, 2}
+	peer := newPeerInIPRange("54")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54")
+	if err := peerSet.Add(peer); err == nil {
+		t.Errorf("Added peer when we shouldn't have")
+	}
+	peer = newPeerInIPRange("55")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+
+	// test  /16
+	peerSet = NewPeerSet()
+	maxPeersPerIPRange = [4]int{3, 2, 1, 1}
+	peer = newPeerInIPRange("54", "112")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54", "112")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54", "112")
+	if err := peerSet.Add(peer); err == nil {
+		t.Errorf("Added peer when we shouldn't have")
+	}
+	peer = newPeerInIPRange("54", "113")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+
+	// test  /24
+	peerSet = NewPeerSet()
+	maxPeersPerIPRange = [4]int{5, 3, 2, 1}
+	peer = newPeerInIPRange("54", "112", "11")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54", "112", "11")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54", "112", "11")
+	if err := peerSet.Add(peer); err == nil {
+		t.Errorf("Added peer when we shouldn't have")
+	}
+	peer = newPeerInIPRange("54", "112", "12")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+
+	// test  /32
+	peerSet = NewPeerSet()
+	maxPeersPerIPRange = [4]int{11, 7, 5, 2}
+	peer = newPeerInIPRange("54", "112", "11", "10")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54", "112", "11", "10")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+	peer = newPeerInIPRange("54", "112", "11", "10")
+	if err := peerSet.Add(peer); err == nil {
+		t.Errorf("Added peer when we shouldn't have")
+	}
+	peer = newPeerInIPRange("54", "112", "11", "11")
+	if err := peerSet.Add(peer); err != nil {
+		t.Errorf("Failed to add new peer")
+	}
+}