Replace our ratelimiter with standard rate package (#359)

* Replace our ratelimiter with standard rate package

Signed-off-by: Wendi Yu <wendi.yu@yahoo.ca>

ratelimit/ratelimit.go

@@ -1,81 +0,0 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-package ratelimit
-
-import (
-	"sync"
-	"time"
-
-	"tailscale.com/types/structs"
-)
-
-type Bucket struct {
-	_            structs.Incomparable
-	mu           sync.Mutex
-	FillInterval time.Duration
-	Burst        int
-	v            int
-	quitCh       chan struct{}
-	started      bool
-	closed       bool
-}
-
-func (b *Bucket) startLocked() {
-	b.v = b.Burst
-	b.quitCh = make(chan struct{})
-	b.started = true
-
-	t := time.NewTicker(b.FillInterval)
-	go func() {
-		for {
-			select {
-			case <-b.quitCh:
-				return
-			case <-t.C:
-				b.tick()
-			}
-		}
-	}()
-}
-
-func (b *Bucket) tick() {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	if b.v < b.Burst {
-		b.v++
-	}
-}
-
-func (b *Bucket) Close() {
-	b.mu.Lock()
-	if !b.started {
-		b.closed = true
-		b.mu.Unlock()
-		return
-	}
-	if b.closed {
-		b.mu.Unlock()
-		return
-	}
-	b.closed = true
-	b.mu.Unlock()
-
-	b.quitCh <- struct{}{}
-}
-
-func (b *Bucket) TryGet() int {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	if !b.started {
-		b.startLocked()
-	}
-	if b.v > 0 {
-		b.v--
-		return b.v + 1
-	}
-	return 0
-}

ratelimit/ratelimit_test.go

@@ -1,28 +0,0 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-package ratelimit
-
-import (
-	"testing"
-	"time"
-)
-
-func TestBucket(t *testing.T) {
-	b := Bucket{
-		FillInterval: time.Second,
-		Burst:        3,
-	}
-	expect := []int{3, 2, 1, 0, 0}
-	for i, want := range expect {
-		got := b.TryGet()
-		if want != got {
-			t.Errorf("#%d want=%d got=%d\n", i, want, got)
-		}
-	}
-	b.tick()
-	if want, got := 1, b.TryGet(); want != got {
-		t.Errorf("after tick: want=%d got=%d\n", want, got)
-	}
-}

wgengine/filter/filter.go

@@ -12,7 +12,7 @@ import (
 	"time"
 
 	"github.com/golang/groupcache/lru"
-	"tailscale.com/ratelimit"
+	"golang.org/x/time/rate"
 	"tailscale.com/wgengine/packet"
 )
 
@@ -105,11 +105,10 @@ func New(matches Matches, shareStateWith *Filter) *Filter {
 }
 
 func maybeHexdump(flag RunFlags, b []byte) string {
-	if flag != 0 {
-		return packet.Hexdump(b) + "\n"
-	} else {
+	if flag == 0 {
 		return ""
 	}
+	return packet.Hexdump(b) + "\n"
 }
 
 // TODO(apenwarr): use a bigger bucket for specifically TCP SYN accept logging?
@@ -117,17 +116,11 @@ func maybeHexdump(flag RunFlags, b []byte) string {
 //   we have to be cautious about flooding the logs vs letting people use
 //   flood protection to hide their traffic. We could use a rate limiter in
 //   the actual *filter* for SYN accepts, perhaps.
-var acceptBucket = ratelimit.Bucket{
-	Burst:        3,
-	FillInterval: 10 * time.Second,
-}
-var dropBucket = ratelimit.Bucket{
-	Burst:        10,
-	FillInterval: 5 * time.Second,
-}
+var acceptBucket = rate.NewLimiter(rate.Every(10*time.Second), 3)
+var dropBucket = rate.NewLimiter(rate.Every(5*time.Second), 10)
 
 func logRateLimit(runflags RunFlags, b []byte, q *packet.QDecode, r Response, why string) {
-	if r == Drop && (runflags&LogDrops) != 0 && dropBucket.TryGet() > 0 {
+	if r == Drop && (runflags&LogDrops) != 0 && dropBucket.Allow() {
 		var qs string
 		if q == nil {
 			qs = fmt.Sprintf("(%d bytes)", len(b))
@@ -135,7 +128,7 @@ func logRateLimit(runflags RunFlags, b []byte, q *packet.QDecode, r Response, wh
 			qs = q.String()
 		}
 		log.Printf("Drop: %v %v %s\n%s", qs, len(b), why, maybeHexdump(runflags&HexdumpDrops, b))
-	} else if r == Accept && (runflags&LogAccepts) != 0 && acceptBucket.TryGet() > 0 {
+	} else if r == Accept && (runflags&LogAccepts) != 0 && acceptBucket.Allow() {
 		log.Printf("Accept: %v %v %s\n%s", q, len(b), why, maybeHexdump(runflags&HexdumpAccepts, b))
 	}
 }