From 05f55dc0afdd62522c3c0aa1caee0a933280ef00 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Thu, 4 Oct 2018 14:29:05 +0200
Subject: [PATCH] p2p/netutil: GC also when adding statements and contacts

This prevents indefinite memory growth if Predict* isn't called.
---
 p2p/netutil/iptrack.go      | 28 ++++++++++++++++++++--------
 p2p/netutil/iptrack_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/p2p/netutil/iptrack.go b/p2p/netutil/iptrack.go
index ddfc91f961..b9cbd5e1ca 100644
--- a/p2p/netutil/iptrack.go
+++ b/p2p/netutil/iptrack.go
@@ -25,12 +25,14 @@ import (
 // IPTracker predicts the external endpoint, i.e. IP address and port, of the local host
 // based on statements made by other hosts.
 type IPTracker struct {
-	window        time.Duration
-	contactWindow time.Duration
-	minStatements int
-	clock         mclock.Clock
-	statements    map[string]ipStatement
-	contact       map[string]mclock.AbsTime
+	window          time.Duration
+	contactWindow   time.Duration
+	minStatements   int
+	clock           mclock.Clock
+	statements      map[string]ipStatement
+	contact         map[string]mclock.AbsTime
+	lastStatementGC mclock.AbsTime
+	lastContactGC   mclock.AbsTime
 }
 
 type ipStatement struct {
@@ -90,16 +92,25 @@ func (it *IPTracker) PredictEndpoint() string {
 
 // AddStatement records that a certain host thinks our external endpoint is the one given.
 func (it *IPTracker) AddStatement(host, endpoint string) {
-	it.statements[host] = ipStatement{endpoint, it.clock.Now()}
+	now := it.clock.Now()
+	it.statements[host] = ipStatement{endpoint, now}
+	if time.Duration(now-it.lastStatementGC) >= it.window {
+		it.gcStatements(now)
+	}
 }
 
 // AddContact records that a packet containing our endpoint information has been sent to a
 // certain host.
 func (it *IPTracker) AddContact(host string) {
-	it.contact[host] = it.clock.Now()
+	now := it.clock.Now()
+	it.contact[host] = now
+	if time.Duration(now-it.lastContactGC) >= it.contactWindow {
+		it.gcContact(now)
+	}
 }
 
 func (it *IPTracker) gcStatements(now mclock.AbsTime) {
+	it.lastStatementGC = now
 	cutoff := now.Add(-it.window)
 	for host, s := range it.statements {
 		if s.time < cutoff {
@@ -109,6 +120,7 @@ func (it *IPTracker) gcStatements(now mclock.AbsTime) {
 }
 
 func (it *IPTracker) gcContact(now mclock.AbsTime) {
+	it.lastContactGC = now
 	cutoff := now.Add(-it.contactWindow)
 	for host, ct := range it.contact {
 		if ct < cutoff {
diff --git a/p2p/netutil/iptrack_test.go b/p2p/netutil/iptrack_test.go
index 115f407bf8..a9a2998a65 100644
--- a/p2p/netutil/iptrack_test.go
+++ b/p2p/netutil/iptrack_test.go
@@ -18,6 +18,7 @@ package netutil
 
 import (
 	"fmt"
+	mrand "math/rand"
 	"testing"
 	"time"
 
@@ -107,3 +108,31 @@ func runIPTrackerTest(t *testing.T, evs []iptrackTestEvent) {
 		}
 	}
 }
+
+// This checks that old statements and contacts are GCed even if Predict* isn't called.
+func TestIPTrackerForceGC(t *testing.T) {
+	var (
+		clock  mclock.Simulated
+		window = 10 * time.Second
+		rate   = 50 * time.Millisecond
+		max    = int(window/rate) + 1
+		it     = NewIPTracker(window, window, 3)
+	)
+	it.clock = &clock
+
+	for i := 0; i < 5*max; i++ {
+		e1 := make([]byte, 4)
+		e2 := make([]byte, 4)
+		mrand.Read(e1)
+		mrand.Read(e2)
+		it.AddStatement(string(e1), string(e2))
+		it.AddContact(string(e1))
+		clock.Run(rate)
+	}
+	if len(it.contact) > 2*max {
+		t.Errorf("contacts not GCed, have %d", len(it.contact))
+	}
+	if len(it.statements) > 2*max {
+		t.Errorf("statements not GCed, have %d", len(it.statements))
+	}
+}