mirror of
https://github.com/tailscale/tailscale.git
synced 2025-01-05 14:57:49 +00:00
Replace our ratelimiter with standard rate package (#359)
* Replace our ratelimiter with standard rate package Signed-off-by: Wendi Yu <wendi.yu@yahoo.ca>
This commit is contained in:
parent
b01db109f5
commit
499c8fcbb3
@ -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
|
||||
}
|
@ -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)
|
||||
}
|
||||
}
|
@ -12,7 +12,7 @@
|
||||
"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))
|
||||
}
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user