diff --git a/ratelimit/ratelimit.go b/ratelimit/ratelimit.go deleted file mode 100644 index 3a0bad70e..000000000 --- a/ratelimit/ratelimit.go +++ /dev/null @@ -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 -} diff --git a/ratelimit/ratelimit_test.go b/ratelimit/ratelimit_test.go deleted file mode 100644 index bcea8c46e..000000000 --- a/ratelimit/ratelimit_test.go +++ /dev/null @@ -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) - } -} diff --git a/wgengine/filter/filter.go b/wgengine/filter/filter.go index 4d5e66ded..03ae539a6 100644 --- a/wgengine/filter/filter.go +++ b/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)) } }