add support for login throttling

irc/connection_limits/throttler.go

@@ -26,8 +26,45 @@ type ThrottlerConfig struct {
 
 // ThrottleDetails holds the connection-throttling details for a subnet/IP.
 type ThrottleDetails struct {
-	Start       time.Time
-	ClientCount int
+	Start time.Time
+	Count int
+}
+
+// GenericThrottle allows enforcing limits of the form
+// "at most X events per time window of duration Y"
+type GenericThrottle struct {
+	ThrottleDetails // variable state: what events have been seen
+	// these are constant after creation:
+	Duration time.Duration // window length to consider
+	Limit    int           // number of events allowed per window
+}
+
+// Touch checks whether an additional event is allowed:
+// it either denies it (by returning false) or allows it (by returning true)
+// and records it
+func (g *GenericThrottle) Touch() (throttled bool, remainingTime time.Duration) {
+	return g.touch(time.Now())
+}
+
+func (g *GenericThrottle) touch(now time.Time) (throttled bool, remainingTime time.Duration) {
+	if g.Limit == 0 {
+		return // limit of 0 disables throttling
+	}
+
+	elapsed := now.Sub(g.Start)
+	if elapsed > g.Duration {
+		// reset window, record the operation
+		g.Start = now
+		g.Count = 1
+		return false, 0
+	} else if g.Count >= g.Limit {
+		// we are throttled
+		return true, g.Start.Add(g.Duration).Sub(now)
+	} else {
+		// we are not throttled, record the operation
+		g.Count += 1
+		return false, 0
+	}
 }
 
 // Throttler manages automated client connection throttling.
@@ -102,21 +139,21 @@ func (ct *Throttler) AddClient(addr net.IP) error {
 	ct.maskAddr(addr)
 	addrString := addr.String()
 
-	details, exists := ct.population[addrString]
-	if !exists || details.Start.Add(ct.duration).Before(time.Now()) {
-		details = ThrottleDetails{
-			Start: time.Now(),
-		}
+	details := ct.population[addrString] // retrieve mutable throttle state from the map
+	// add in constant state to process the limiting operation
+	g := GenericThrottle{
+		ThrottleDetails: details,
+		Duration:        ct.duration,
+		Limit:           ct.subnetLimit,
 	}
+	throttled, _ := g.Touch()                     // actually check the limit
+	ct.population[addrString] = g.ThrottleDetails // store modified mutable state
 
-	if details.ClientCount+1 > ct.subnetLimit {
+	if throttled {
 		return errTooManyClients
+	} else {
+		return nil
 	}
-
-	details.ClientCount++
-	ct.population[addrString] = details
-
-	return nil
 }
 
 func (ct *Throttler) BanDuration() time.Duration {

irc/connection_limits/throttler_test.go

@@ -0,0 +1,86 @@
+// Copyright (c) 2018 Shivaram Lingamneni
+// released under the MIT license
+
+package connection_limits
+
+import (
+	"net"
+	"reflect"
+	"testing"
+	"time"
+)
+
+func assertEqual(supplied, expected interface{}, t *testing.T) {
+	if !reflect.DeepEqual(supplied, expected) {
+		t.Errorf("expected %v but got %v", expected, supplied)
+	}
+}
+
+func TestGenericThrottle(t *testing.T) {
+	minute, _ := time.ParseDuration("1m")
+	second, _ := time.ParseDuration("1s")
+	zero, _ := time.ParseDuration("0s")
+
+	throttler := GenericThrottle{
+		Duration: minute,
+		Limit:    2,
+	}
+
+	now := time.Now()
+	throttled, remaining := throttler.touch(now)
+	assertEqual(throttled, false, t)
+	assertEqual(remaining, zero, t)
+
+	now = now.Add(second)
+	throttled, remaining = throttler.touch(now)
+	assertEqual(throttled, false, t)
+	assertEqual(remaining, zero, t)
+
+	now = now.Add(second)
+	throttled, remaining = throttler.touch(now)
+	assertEqual(throttled, true, t)
+	assertEqual(remaining, 58*second, t)
+
+	now = now.Add(minute)
+	throttled, remaining = throttler.touch(now)
+	assertEqual(throttled, false, t)
+	assertEqual(remaining, zero, t)
+}
+
+func TestGenericThrottleDisabled(t *testing.T) {
+	minute, _ := time.ParseDuration("1m")
+	throttler := GenericThrottle{
+		Duration: minute,
+		Limit:    0,
+	}
+
+	for i := 0; i < 1024; i += 1 {
+		throttled, _ := throttler.Touch()
+		if throttled {
+			t.Error("disabled throttler should not throttle")
+		}
+	}
+}
+
+func TestConnectionThrottle(t *testing.T) {
+	minute, _ := time.ParseDuration("1m")
+	maxConnections := 3
+	config := ThrottlerConfig{
+		Enabled:            true,
+		CidrLenIPv4:        32,
+		CidrLenIPv6:        64,
+		ConnectionsPerCidr: maxConnections,
+		Duration:           minute,
+	}
+	throttler := NewThrottler()
+	throttler.ApplyConfig(config)
+
+	addr := net.ParseIP("8.8.8.8")
+
+	for i := 0; i < maxConnections; i += 1 {
+		err := throttler.AddClient(addr)
+		assertEqual(err, nil, t)
+	}
+	err := throttler.AddClient(addr)
+	assertEqual(err, errTooManyClients, t)
+}