Moved ratelimiter to internal package

4 files changed, 56 insertions, 21 deletions

diff --git a/ratelimiter.go b/internal/ratelimiter/ratelimiter.go
index 6e5f005..f9fc673 100644
--- a/ratelimiter.go
+++ b/internal/ratelimiter/ratelimiter.go
@@ -1,4 +1,4 @@
-package main
+package ratelimiter
 
 /* Copyright (C) 2015-2017 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved. */
 
@@ -26,21 +26,48 @@ type RatelimiterEntry struct {
 }
 
 type Ratelimiter struct {
-	mutex              sync.RWMutex
-	lastGarbageCollect time.Time
-	tableIPv4          map[[net.IPv4len]byte]*RatelimiterEntry
-	tableIPv6          map[[net.IPv6len]byte]*RatelimiterEntry
+	mutex     sync.RWMutex
+	stop      chan struct{}
+	tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry
+	tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry
+}
+
+func (rate *Ratelimiter) Close() {
+	rate.mutex.Lock()
+	defer rate.mutex.Unlock()
+
+	if rate.stop != nil {
+		close(rate.stop)
+	}
 }
 
 func (rate *Ratelimiter) Init() {
 	rate.mutex.Lock()
 	defer rate.mutex.Unlock()
+
+	if rate.stop != nil {
+		close(rate.stop)
+	}
+
+	rate.stop = make(chan struct{})
 	rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry)
 	rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry)
-	rate.lastGarbageCollect = time.Now()
+
+	go func() {
+		timer := time.NewTimer(time.Second)
+		for {
+			select {
+			case <-rate.stop:
+				return
+			case <-timer.C:
+				rate.garbageCollectEntries()
+				timer.Reset(time.Second)
+			}
+		}
+	}()
 }
 
-func (rate *Ratelimiter) GarbageCollectEntries() {
+func (rate *Ratelimiter) garbageCollectEntries() {
 	rate.mutex.Lock()
 
 	// remove unused IPv4 entries
@@ -66,19 +93,6 @@ func (rate *Ratelimiter) GarbageCollectEntries() {
 	rate.mutex.Unlock()
 }
 
-func (rate *Ratelimiter) RoutineGarbageCollector(stop Signal) {
-	timer := time.NewTimer(time.Second)
-	for {
-		select {
-		case <-stop.Wait():
-			return
-		case <-timer.C:
-			rate.GarbageCollectEntries()
-			timer.Reset(time.Second)
-		}
-	}
-}
-
 func (rate *Ratelimiter) Allow(ip net.IP) bool {
 	var entry *RatelimiterEntry
 	var KeyIPv4 [net.IPv4len]byte
diff --git a/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go
index 13b6a23..a6f618b 100644
--- a/ratelimiter_test.go
+++ b/internal/ratelimiter/ratelimiter_test.go
@@ -1,4 +1,4 @@
-package main
+package ratelimiter
 
 import (
 	"net"
diff --git a/internal/tai64n/tai64.go b/internal/tai64n/tai64n.go
index da5257c..da5257c 100644
--- a/internal/tai64n/tai64.go
+++ b/internal/tai64n/tai64n.go
diff --git a/internal/tai64n/tai64n_test.go b/internal/tai64n/tai64n_test.go
new file mode 100644
index 0000000..389b65c
--- /dev/null
+++ b/internal/tai64n/tai64n_test.go
@@ -0,0 +1,21 @@
+package tai64n
+
+import (
+	"testing"
+	"time"
+)
+
+/* Testing the essential property of the timestamp
+ * as used by WireGuard.
+ */
+func TestMonotonic(t *testing.T) {
+	old := Now()
+	for i := 0; i < 10000; i++ {
+		time.Sleep(time.Nanosecond)
+		next := Now()
+		if !next.After(old) {
+			t.Error("TAI64N, not monotonically increasing on nano-second scale")
+		}
+		old = next
+	}
+}