From 84f577bdd16b9b11601ebed1cb1f2db3f89f0397 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Tue, 9 Feb 2021 10:17:30 +0200
Subject: [PATCH] add subnet whitelisting for IPColocation

---
 score.go        | 10 +++++++++
 score_params.go |  8 ++++---
 score_test.go   | 60 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/score.go b/score.go
index 76a27c6..31c9f47 100644
--- a/score.go
+++ b/score.go
@@ -336,12 +336,22 @@ func (ps *peerScore) ipColocationFactor(p peer.ID) float64 {
 	}
 
 	var result float64
+loop:
 	for _, ip := range pstats.ips {
 		_, whitelisted := ps.params.IPColocationFactorWhitelist[ip]
 		if whitelisted {
 			continue
 		}
 
+		if len(ps.params.IPColocationFactorWhitelistSubnets) > 0 {
+			ipObj := net.ParseIP(ip)
+			for _, ipNet := range ps.params.IPColocationFactorWhitelistSubnets {
+				if ipNet.Contains(ipObj) {
+					continue loop
+				}
+			}
+		}
+
 		// P6 has a cliff (IPColocationFactorThreshold); it's only applied iff
 		// at least that many peers are connected to us from that source IP
 		// addr. It is quadratic, and the weight is negative (validated by
diff --git a/score_params.go b/score_params.go
index 3a20584..44c69bd 100644
--- a/score_params.go
+++ b/score_params.go
@@ -3,6 +3,7 @@ package pubsub
 import (
 	"fmt"
 	"math"
+	"net"
 	"time"
 
 	"github.com/libp2p/go-libp2p-core/peer"
@@ -69,9 +70,10 @@ type PeerScoreParams struct {
 	// The weight of the parameter MUST be negative, unless you want to disable for testing.
 	// Note: In order to simulate many IPs in a managable manner when testing, you can set the weight to 0
 	//       thus disabling the IP colocation penalty.
-	IPColocationFactorWeight    float64
-	IPColocationFactorThreshold int
-	IPColocationFactorWhitelist map[string]struct{}
+	IPColocationFactorWeight           float64
+	IPColocationFactorThreshold        int
+	IPColocationFactorWhitelist        map[string]struct{}
+	IPColocationFactorWhitelistSubnets []*net.IPNet
 
 	// P7: behavioural pattern penalties.
 	// This parameter has an associated counter which tracks misbehaviour as detected by the
diff --git a/score_test.go b/score_test.go
index 217983e..c5d7dbd 100644
--- a/score_test.go
+++ b/score_test.go
@@ -2,6 +2,7 @@ package pubsub
 
 import (
 	"math"
+	"net"
 	"sync"
 	"testing"
 	"time"
@@ -742,6 +743,65 @@ func TestScoreIPColocation(t *testing.T) {
 	}
 }
 
+func TestScoreIPColocationWhitelistSubnet(t *testing.T) {
+	// Create parameters with reasonable default values
+	mytopic := "mytopic"
+
+	_, ipNet, err := net.ParseCIDR("2.3.0.0/16")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	params := &PeerScoreParams{
+		AppSpecificScore:                   func(peer.ID) float64 { return 0 },
+		IPColocationFactorThreshold:        1,
+		IPColocationFactorWeight:           -1,
+		IPColocationFactorWhitelistSubnets: []*net.IPNet{ipNet},
+		Topics:                             make(map[string]*TopicScoreParams),
+	}
+
+	peerA := peer.ID("A")
+	peerB := peer.ID("B")
+	peerC := peer.ID("C")
+	peerD := peer.ID("D")
+	peers := []peer.ID{peerA, peerB, peerC, peerD}
+
+	ps := newPeerScore(params)
+	for _, p := range peers {
+		ps.AddPeer(p, "myproto")
+		ps.Graft(p, mytopic)
+	}
+
+	// peerA should have no penalty, but B, C, and D should be penalized for sharing an IP
+	setIPsForPeer(t, ps, peerA, "1.2.3.4")
+	setIPsForPeer(t, ps, peerB, "2.3.4.5")
+	setIPsForPeer(t, ps, peerC, "2.3.4.5", "3.4.5.6")
+	setIPsForPeer(t, ps, peerD, "2.3.4.5")
+
+	ps.refreshScores()
+	aScore := ps.Score(peerA)
+	bScore := ps.Score(peerB)
+	cScore := ps.Score(peerC)
+	dScore := ps.Score(peerD)
+
+	if aScore != 0 {
+		t.Errorf("expected peer A to have score 0.0, got %f", aScore)
+	}
+
+	if bScore != 0 {
+		t.Errorf("expected peer B to have score 0.0, got %f", aScore)
+	}
+
+	if cScore != 0 {
+		t.Errorf("expected peer C to have score 0.0, got %f", aScore)
+	}
+
+	if dScore != 0 {
+		t.Errorf("expected peer D to have score 0.0, got %f", aScore)
+	}
+
+}
+
 func TestScoreBehaviourPenalty(t *testing.T) {
 	params := &PeerScoreParams{
 		AppSpecificScore:       func(peer.ID) float64 { return 0 },