go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix traffic limiter

Browse Source
This commit is contained in:
ginuerzh
2022-12-19 19:33:29 +08:00
parent 1cb719f694
commit 15feb7599e
6 changed files with 438 additions and 246 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
go.mod
View File

@ -18,6 +18,7 @@ require (
github.com/gorilla/websocket v1.5.0
github.com/lucas-clemente/quic-go v0.30.0
github.com/miekg/dns v1.1.50
github.com/patrickmn/go-cache v2.1.0+incompatible
github.com/pion/dtls/v2 v2.1.5
github.com/pires/go-proxyproto v0.6.2
github.com/prometheus/client_golang v1.12.1

2
go.sum
View File

@ -280,6 +280,8 @@ github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
github.com/onsi/ginkgo/v2 v2.2.0 h1:3ZNA3L1c5FYDFTTxbFeVGGD8jYvjYauHD30YgLxVsNI=
github.com/onsi/ginkgo/v2 v2.2.0/go.mod h1:MEH45j8TBi6u9BMogfbp0stKC5cdGjumZj5Y7AG4VIk=
github.com/onsi/gomega v1.20.1 h1:PA/3qinGoukvymdIDV8pii6tiZgC8kbmJO6Z5+b002Q=
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
github.com/pelletier/go-toml v1.9.4 h1:tjENF6MfZAg8e4ZmZTeWaWiT2vXtsoO6+iuOjFhECwM=
github.com/pelletier/go-toml v1.9.4/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
github.com/pion/dtls/v2 v2.1.5 h1:jlh2vtIyUBShchoTDqpCCqiYCyRFJ/lvf/gQ8TALs+c=

40
limiter/traffic/generator.go
View File

@ -4,58 +4,28 @@ import (
limiter "github.com/go-gost/core/limiter/traffic"
)
type TrafficLimitGenerator interface {
In() limiter.Limiter
Out() limiter.Limiter
}
type trafficLimitGenerator struct {
type limitGenerator struct {
in int
out int
}
func NewTrafficLimitGenerator(in, out int) TrafficLimitGenerator {
return &trafficLimitGenerator{
func newLimitGenerator(in, out int) *limitGenerator {
return &limitGenerator{
in: in,
out: out,
}
}
func (p *trafficLimitGenerator) In() limiter.Limiter {
func (p *limitGenerator) In() limiter.Limiter {
if p == nil || p.in <= 0 {
return nil
}
return NewLimiter(p.in)
}
func (p *trafficLimitGenerator) Out() limiter.Limiter {
func (p *limitGenerator) Out() limiter.Limiter {
if p == nil || p.out <= 0 {
return nil
}
return NewLimiter(p.out)
}
type trafficLimitSingleGenerator struct {
in limiter.Limiter
out limiter.Limiter
}
func NewTrafficLimitSingleGenerator(in, out int) TrafficLimitGenerator {
p := &trafficLimitSingleGenerator{}
if in > 0 {
p.in = NewLimiter(in)
}
if out > 0 {
p.out = NewLimiter(out)
}
return p
}
func (p *trafficLimitSingleGenerator) In() limiter.Limiter {
return p.in
}
func (p *trafficLimitSingleGenerator) Out() limiter.Limiter {
return p.out
}

48
limiter/traffic/limiter.go
View File

@ -2,6 +2,9 @@ package traffic
import (
"context"
"fmt"
"sort"
"strconv"
limiter "github.com/go-gost/core/limiter/traffic"
"golang.org/x/time/rate"
@ -26,5 +29,48 @@ func (l *llimiter) Wait(ctx context.Context, n int) int {
}
func (l *llimiter) Limit() int {
return l.limiter.Burst()
return int(l.limiter.Limit())
}
func (l *llimiter) Set(n int) {
l.limiter.SetLimit(rate.Limit(n))
l.limiter.SetBurst(n)
}
func (l *llimiter) String() string {
return strconv.Itoa(int(l.limiter.Limit()))
}
type limiterGroup struct {
limiters []limiter.Limiter
}
func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
sort.Slice(limiters, func(i, j int) bool {
return limiters[i].Limit() < limiters[j].Limit()
})
return &limiterGroup{limiters: limiters}
}
func (l *limiterGroup) Wait(ctx context.Context, n int) int {
for i := range l.limiters {
if v := l.limiters[i].Wait(ctx, n); v < n {
n = v
}
}
return n
}
func (l *limiterGroup) Limit() int {
if len(l.limiters) == 0 {
return 0
}
return l.limiters[0].Limit()
}
func (l *limiterGroup) Set(n int) {}
func (l *limiterGroup) String() string {
return fmt.Sprintf("%v", l.limiters)
}

407
limiter/traffic/traffic.go
View File

@ -5,7 +5,6 @@ import (
"context"
"io"
"net"
"sort"
"strings"
"sync"
"time"
@ -14,6 +13,7 @@ import (
limiter "github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/logger"
"github.com/go-gost/x/internal/loader"
"github.com/patrickmn/go-cache"
"github.com/yl2chen/cidranger"
)
@ -22,33 +22,10 @@ const (
ConnLimitKey = "$$"
)
type limiterGroup struct {
limiters []limiter.Limiter
}
func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
sort.Slice(limiters, func(i, j int) bool {
return limiters[i].Limit() < limiters[j].Limit()
})
return &limiterGroup{limiters: limiters}
}
func (l *limiterGroup) Wait(ctx context.Context, n int) int {
for i := range l.limiters {
if v := l.limiters[i].Wait(ctx, n); v < n {
n = v
}
}
return n
}
func (l *limiterGroup) Limit() int {
if len(l.limiters) == 0 {
return 0
}
return l.limiters[0].Limit()
}
const (
defaultExpiration = 15 * time.Second
cleanupInterval = 30 * time.Second
)
type options struct {
limits []string
@ -97,12 +74,19 @@ func LoggerOption(logger logger.Logger) Option {
}
}
type limitValue struct {
in int
out int
}
type trafficLimiter struct {
limits map[string]TrafficLimitGenerator
cidrLimits cidranger.Ranger
inLimits map[string]limiter.Limiter
outLimits map[string]limiter.Limiter
mu sync.Mutex
generators sync.Map
cidrGenerators cidranger.Ranger
connInLimits *cache.Cache
connOutLimits *cache.Cache
inLimits *cache.Cache
outLimits *cache.Cache
mu sync.RWMutex
cancelFunc context.CancelFunc
options options
}
@ -115,10 +99,11 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
ctx, cancel := context.WithCancel(context.TODO())
lim := &trafficLimiter{
limits: make(map[string]TrafficLimitGenerator),
cidrLimits: cidranger.NewPCTrieRanger(),
inLimits: make(map[string]limiter.Limiter),
outLimits: make(map[string]limiter.Limiter),
cidrGenerators: cidranger.NewPCTrieRanger(),
connInLimits: cache.New(defaultExpiration, cleanupInterval),
connOutLimits: cache.New(defaultExpiration, cleanupInterval),
inLimits: cache.New(defaultExpiration, cleanupInterval),
outLimits: cache.New(defaultExpiration, cleanupInterval),
options: options,
cancelFunc: cancel,
}
@ -132,47 +117,57 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
return lim
}
// In obtains a traffic input limiter based on key.
// The key should be client connection address.
func (l *trafficLimiter) In(key string) limiter.Limiter {
l.mu.Lock()
defer l.mu.Unlock()
var lims []limiter.Limiter
if p := l.limits[ConnLimitKey]; p != nil {
if lim := p.In(); lim != nil {
lims = append(lims, lim)
// service level limiter
if lim, ok := l.inLimits.Get(GlobalLimitKey); ok && lim != nil {
lims = append(lims, lim.(limiter.Limiter))
}
// connection level limiter
if lim, ok := l.connInLimits.Get(key); ok {
if lim != nil {
// cached connection level limiter
lims = append(lims, lim.(limiter.Limiter))
// reset expiration
l.connInLimits.Set(key, lim, defaultExpiration)
}
if p := l.limits[GlobalLimitKey]; p != nil {
if lim := p.In(); lim != nil {
} else {
// generate a new connection level limiter and cache it
if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil {
lim := v.(*limitGenerator).In()
if lim != nil {
lims = append(lims, lim)
l.connInLimits.Set(key, lim, defaultExpiration)
}
}
}
// IP limiter
if lim, ok := l.inLimits[key]; ok {
host, _, _ := net.SplitHostPort(key)
// IP level limiter
if lim, ok := l.inLimits.Get(host); ok {
// cached IP limiter
if lim != nil {
lims = append(lims, lim)
lims = append(lims, lim.(limiter.Limiter))
}
} else {
if ip := net.ParseIP(key); ip != nil {
if p := l.limits[key]; p != nil {
if lim = p.In(); lim != nil {
lims = append(lims, lim)
}
}
if lim == nil {
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
l.mu.RLock()
ranger := l.cidrGenerators
l.mu.RUnlock()
// CIDR level limiter
if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 {
if v, _ := p[0].(*cidrLimitEntry); v != nil {
if lim = v.limit.In(); lim != nil {
if lim := v.generator.In(); lim != nil {
lims = append(lims, lim)
l.inLimits.Set(host, lim, cache.NoExpiration)
}
}
}
}
l.inLimits[key] = lim
}
}
var lim limiter.Limiter
if len(lims) > 0 {
@ -180,53 +175,63 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
}
if lim != nil && l.options.logger != nil {
l.options.logger.Debugf("input limit for %s: %d", key, lim.Limit())
l.options.logger.Debugf("input limit for %s: %s", key, lim)
}
return lim
}
// Out obtains a traffic output limiter based on key.
// The key should be client connection address.
func (l *trafficLimiter) Out(key string) limiter.Limiter {
l.mu.Lock()
defer l.mu.Unlock()
var lims []limiter.Limiter
if p := l.limits[ConnLimitKey]; p != nil {
if lim := p.Out(); lim != nil {
lims = append(lims, lim)
// service level limiter
if lim, ok := l.outLimits.Get(GlobalLimitKey); ok && lim != nil {
lims = append(lims, lim.(limiter.Limiter))
}
// connection level limiter
if lim, ok := l.connOutLimits.Get(key); ok {
if lim != nil {
// cached connection level limiter
lims = append(lims, lim.(limiter.Limiter))
// reset expiration
l.connOutLimits.Set(key, lim, defaultExpiration)
}
if p := l.limits[GlobalLimitKey]; p != nil {
if lim := p.Out(); lim != nil {
} else {
// generate a new connection level limiter
if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil {
lim := v.(*limitGenerator).Out()
if lim != nil {
lims = append(lims, lim)
l.connOutLimits.Set(key, lim, defaultExpiration)
}
}
}
// IP limiter
if lim, ok := l.outLimits[key]; ok {
host, _, _ := net.SplitHostPort(key)
// IP level limiter
if lim, ok := l.outLimits.Get(host); ok {
if lim != nil {
lims = append(lims, lim)
// cached IP level limiter
lims = append(lims, lim.(limiter.Limiter))
}
} else {
if ip := net.ParseIP(key); ip != nil {
if p := l.limits[key]; p != nil {
if lim = p.Out(); lim != nil {
lims = append(lims, lim)
}
}
if lim == nil {
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
l.mu.RLock()
ranger := l.cidrGenerators
l.mu.RUnlock()
// CIDR level limiter
if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 {
if v, _ := p[0].(*cidrLimitEntry); v != nil {
if lim = v.limit.Out(); lim != nil {
if lim := v.generator.Out(); lim != nil {
lims = append(lims, lim)
l.outLimits.Set(host, lim, cache.NoExpiration)
}
}
}
}
l.outLimits[key] = lim
}
}
var lim limiter.Limiter
if len(lims) > 0 {
@ -234,7 +239,7 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
}
if lim != nil && l.options.logger != nil {
l.options.logger.Debugf("output limit for %s: %d", key, lim.Limit())
l.options.logger.Debugf("output limit for %s: %s", key, lim)
}
return lim
@ -262,36 +267,160 @@ func (l *trafficLimiter) periodReload(ctx context.Context) error {
}
func (l *trafficLimiter) reload(ctx context.Context) error {
v, err := l.load(ctx)
values, err := l.load(ctx)
if err != nil {
return err
}
lines := append(l.options.limits, v...)
// service level limiter, never expired
{
value := values[GlobalLimitKey]
if v, _ := l.inLimits.Get(GlobalLimitKey); v != nil {
lim := v.(limiter.Limiter)
if value.in <= 0 {
l.inLimits.Delete(GlobalLimitKey)
} else {
lim.Set(value.in)
}
} else {
if value.in > 0 {
l.inLimits.Set(GlobalLimitKey, NewLimiter(value.in), cache.NoExpiration)
}
}
limits := make(map[string]TrafficLimitGenerator)
cidrLimits := cidranger.NewPCTrieRanger()
if v, _ := l.outLimits.Get(GlobalLimitKey); v != nil {
lim := v.(limiter.Limiter)
if value.out <= 0 {
l.outLimits.Delete(GlobalLimitKey)
} else {
lim.Set(value.out)
}
} else {
if value.out > 0 {
l.outLimits.Set(GlobalLimitKey, NewLimiter(value.out), cache.NoExpiration)
}
}
delete(values, GlobalLimitKey)
}
for _, s := range lines {
key, in, out := l.parseLimit(s)
if key == "" {
// connection level limiters
{
value := values[ConnLimitKey]
var in, out int
if v, _ := l.generators.Load(ConnLimitKey); v != nil {
in, out = v.(*limitGenerator).in, v.(*limitGenerator).out
}
l.generators.Store(ConnLimitKey, newLimitGenerator(value.in, value.out))
if value.in <= 0 {
l.connInLimits.Flush()
} else {
if in != value.in {
for _, item := range l.connInLimits.Items() {
if v := item.Object; v != nil {
v.(limiter.Limiter).Set(in)
}
}
}
}
if value.out <= 0 {
l.connOutLimits.Flush()
} else {
if out != value.out {
for _, item := range l.connOutLimits.Items() {
if v := item.Object; v != nil {
v.(limiter.Limiter).Set(out)
}
}
}
}
delete(values, ConnLimitKey)
}
cidrGenerators := cidranger.NewPCTrieRanger()
// IP/CIDR level limiters
{
// snapshot of the current limiters
inLimits := l.inLimits.Items()
outLimits := l.outLimits.Items()
delete(inLimits, GlobalLimitKey)
delete(outLimits, GlobalLimitKey)
for key, value := range values {
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
cidrGenerators.Insert(&cidrLimitEntry{
ipNet: *ipNet,
generator: newLimitGenerator(value.in, value.out),
})
continue
}
switch key {
case GlobalLimitKey:
limits[key] = NewTrafficLimitSingleGenerator(in, out)
case ConnLimitKey:
limits[key] = NewTrafficLimitGenerator(in, out)
default:
if ip := net.ParseIP(key); ip != nil {
limits[key] = NewTrafficLimitSingleGenerator(in, out)
break
if v, _ := l.inLimits.Get(key); v != nil {
lim := v.(limiter.Limiter)
if value.in <= 0 {
l.inLimits.Delete(key)
} else {
lim.Set(value.in)
}
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
cidrLimits.Insert(&cidrLimitEntry{
ipNet: *ipNet,
limit: NewTrafficLimitGenerator(in, out),
})
delete(inLimits, key)
} else {
if value.in > 0 {
l.inLimits.Set(key, NewLimiter(value.in), cache.NoExpiration)
}
}
if v, _ := l.outLimits.Get(key); v != nil {
lim := v.(limiter.Limiter)
if value.out <= 0 {
l.outLimits.Delete(key)
} else {
lim.Set(value.out)
}
delete(outLimits, key)
} else {
if value.out > 0 {
l.outLimits.Set(key, NewLimiter(value.out), cache.NoExpiration)
}
}
}
// check the CIDR for remain limiters, clean the unmatched ones.
for k, v := range inLimits {
if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 {
if le, _ := p[0].(*cidrLimitEntry); le != nil {
in := le.generator.in
if in <= 0 {
l.inLimits.Delete(k)
continue
}
lim := v.Object.(limiter.Limiter)
if lim.Limit() != in {
lim.Set(in)
}
}
} else {
l.inLimits.Delete(k)
}
}
for k, v := range outLimits {
if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 {
if le, _ := p[0].(*cidrLimitEntry); le != nil {
out := le.generator.out
if out <= 0 {
l.outLimits.Delete(k)
continue
}
lim := v.Object.(limiter.Limiter)
if lim.Limit() != out {
lim.Set(out)
}
delete(outLimits, k)
}
} else {
l.outLimits.Delete(k)
}
}
}
@ -299,15 +428,22 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
l.mu.Lock()
defer l.mu.Unlock()
l.limits = limits
l.cidrLimits = cidrLimits
l.inLimits = make(map[string]limiter.Limiter)
l.outLimits = make(map[string]limiter.Limiter)
l.cidrGenerators = cidrGenerators
return nil
}
func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error) {
func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue, err error) {
values = make(map[string]limitValue)
for _, v := range l.options.limits {
key, in, out := l.parseLimit(v)
if key == "" {
continue
}
values[key] = limitValue{in: in, out: out}
}
if l.options.fileLoader != nil {
if lister, ok := l.options.fileLoader.(loader.Lister); ok {
list, er := lister.List(ctx)
@ -315,17 +451,24 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error
l.options.logger.Warnf("file loader: %v", er)
}
for _, s := range list {
if line := l.parseLine(s); line != "" {
patterns = append(patterns, line)
key, in, out := l.parseLimit(l.parseLine(s))
if key == "" {
continue
}
values[key] = limitValue{in: in, out: out}
}
} else {
r, er := l.options.fileLoader.Load(ctx)
if er != nil {
l.options.logger.Warnf("file loader: %v", er)
}
if v, _ := l.parsePatterns(r); v != nil {
patterns = append(patterns, v...)
patterns, _ := l.parsePatterns(r)
for _, s := range patterns {
key, in, out := l.parseLimit(l.parseLine(s))
if key == "" {
continue
}
values[key] = limitValue{in: in, out: out}
}
}
}
@ -335,14 +478,25 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error
if er != nil {
l.options.logger.Warnf("redis loader: %v", er)
}
patterns = append(patterns, list...)
for _, s := range list {
key, in, out := l.parseLimit(l.parseLine(s))
if key == "" {
continue
}
values[key] = limitValue{in: in, out: out}
}
} else {
r, er := l.options.redisLoader.Load(ctx)
if er != nil {
l.options.logger.Warnf("redis loader: %v", er)
}
if v, _ := l.parsePatterns(r); v != nil {
patterns = append(patterns, v...)
patterns, _ := l.parsePatterns(r)
for _, s := range patterns {
key, in, out := l.parseLimit(l.parseLine(s))
if key == "" {
continue
}
values[key] = limitValue{in: in, out: out}
}
}
}
@ -351,12 +505,17 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error
if er != nil {
l.options.logger.Warnf("http loader: %v", er)
}
if v, _ := l.parsePatterns(r); v != nil {
patterns = append(patterns, v...)
patterns, _ := l.parsePatterns(r)
for _, s := range patterns {
key, in, out := l.parseLimit(l.parseLine(s))
if key == "" {
continue
}
values[key] = limitValue{in: in, out: out}
}
}
l.options.logger.Debugf("load items %d", len(patterns))
l.options.logger.Debugf("load items %d", len(values))
return
}
@ -386,6 +545,10 @@ func (l *trafficLimiter) parseLine(s string) string {
func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) {
s = strings.Replace(s, "\t", " ", -1)
s = strings.TrimSpace(s)
if s == "" {
return
}
var ss []string
for _, v := range strings.Split(s, " ") {
if v != "" {
@ -422,7 +585,7 @@ func (l *trafficLimiter) Close() error {
type cidrLimitEntry struct {
ipNet net.IPNet
limit TrafficLimitGenerator
generator *limitGenerator
}
func (p *cidrLimitEntry) Network() net.IPNet {

158
limiter/traffic/wrapper/conn.go
View File

@ -6,40 +6,63 @@ import (
"errors"
"io"
"net"
"sync"
"syscall"
"time"
limiter "github.com/go-gost/core/limiter/traffic"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/internal/net/udp"
"github.com/patrickmn/go-cache"
)
var (
errUnsupport = errors.New("unsupported operation")
)
// serverConn is a server side Conn with metrics supported.
// serverConn is a server side Conn with traffic limiter supported.
type serverConn struct {
net.Conn
rbuf bytes.Buffer
limiter limiter.TrafficLimiter
limiterIn limiter.Limiter
expIn int64
limiterOut limiter.Limiter
expOut int64
}
func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn {
if rlimiter == nil {
func WrapConn(limiter limiter.TrafficLimiter, c net.Conn) net.Conn {
if limiter == nil {
return c
}
host, _, _ := net.SplitHostPort(c.RemoteAddr().String())
return &serverConn{
Conn: c,
limiterIn: rlimiter.In(host),
limiterOut: rlimiter.Out(host),
limiter: limiter,
}
}
func (c *serverConn) getInLimiter(addr net.Addr) limiter.Limiter {
now := time.Now().UnixNano()
// cache the limiter for 1s
if c.limiter != nil && time.Duration(now-c.expIn) > time.Second {
c.limiterIn = c.limiter.In(addr.String())
c.expIn = now
}
return c.limiterIn
}
func (c *serverConn) getOutLimiter(addr net.Addr) limiter.Limiter {
now := time.Now().UnixNano()
// cache the limiter for 1s
if c.limiter != nil && time.Duration(now-c.expOut) > time.Second {
c.limiterOut = c.limiter.Out(addr.String())
c.expOut = now
}
return c.limiterOut
}
func (c *serverConn) Read(b []byte) (n int, err error) {
if c.limiterIn == nil {
limiter := c.getInLimiter(c.RemoteAddr())
if limiter == nil {
return c.Conn.Read(b)
}
@ -48,7 +71,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
if c.rbuf.Len() < burst {
burst = c.rbuf.Len()
}
lim := c.limiterIn.Wait(context.Background(), burst)
lim := limiter.Wait(context.Background(), burst)
return c.rbuf.Read(b[:lim])
}
@ -57,7 +80,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
return nn, err
}
n = c.limiterIn.Wait(context.Background(), nn)
n = limiter.Wait(context.Background(), nn)
if n < nn {
if _, err = c.rbuf.Write(b[n:nn]); err != nil {
return 0, err
@ -68,13 +91,14 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
}
func (c *serverConn) Write(b []byte) (n int, err error) {
if c.limiterOut == nil {
limiter := c.getOutLimiter(c.RemoteAddr())
if limiter == nil {
return c.Conn.Write(b)
}
nn := 0
for len(b) > 0 {
nn, err = c.Conn.Write(b[:c.limiterOut.Wait(context.Background(), len(b))])
nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))])
n += nn
if err != nil {
return
@ -97,10 +121,8 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
type packetConn struct {
net.PacketConn
limiter limiter.TrafficLimiter
inLimits map[string]limiter.Limiter
inMux sync.RWMutex
outLimits map[string]limiter.Limiter
outMux sync.RWMutex
inLimits *cache.Cache
outLimits *cache.Cache
}
func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn {
@ -110,8 +132,8 @@ func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketCon
return &packetConn{
PacketConn: pc,
limiter: lim,
inLimits: make(map[string]limiter.Limiter),
outLimits: make(map[string]limiter.Limiter),
inLimits: cache.New(time.Second, 10*time.Second),
outLimits: cache.New(time.Second, 10*time.Second),
}
}
@ -120,24 +142,21 @@ func (c *packetConn) getInLimiter(addr net.Addr) limiter.Limiter {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.inMux.RLock()
defer c.inMux.RUnlock()
lim, ok := c.inLimits[addr.String()]
return lim, ok
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.inLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.In(host)
c.inMux.Lock()
defer c.inMux.Unlock()
c.inLimits[addr.String()] = lim
lim = c.limiter.In(addr.String())
c.inLimits.Set(addr.String(), lim, 0)
return lim
}
@ -147,24 +166,21 @@ func (c *packetConn) getOutLimiter(addr net.Addr) limiter.Limiter {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.outMux.RLock()
defer c.outMux.RUnlock()
lim, ok := c.outLimits[addr.String()]
return lim, ok
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.outLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.Out(host)
c.outMux.Lock()
defer c.outMux.Unlock()
c.outLimits[addr.String()] = lim
lim = c.limiter.Out(addr.String())
c.outLimits.Set(addr.String(), lim, 0)
return lim
}
@ -204,16 +220,16 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
type udpConn struct {
net.PacketConn
limiter limiter.TrafficLimiter
inLimits map[string]limiter.Limiter
inMux sync.RWMutex
outLimits map[string]limiter.Limiter
outMux sync.RWMutex
inLimits *cache.Cache
outLimits *cache.Cache
}
func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn {
return &udpConn{
PacketConn: pc,
limiter: limiter,
inLimits: cache.New(time.Second, 10*time.Second),
outLimits: cache.New(time.Second, 10*time.Second),
}
}
@ -222,24 +238,21 @@ func (c *udpConn) getInLimiter(addr net.Addr) limiter.Limiter {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.inMux.RLock()
defer c.inMux.RUnlock()
lim, ok := c.inLimits[addr.String()]
return lim, ok
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.inLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.In(host)
c.inMux.Lock()
defer c.inMux.Unlock()
c.inLimits[addr.String()] = lim
lim = c.limiter.In(addr.String())
c.inLimits.Set(addr.String(), lim, 0)
return lim
}
@ -249,24 +262,21 @@ func (c *udpConn) getOutLimiter(addr net.Addr) limiter.Limiter {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.outMux.RLock()
defer c.outMux.RUnlock()
lim, ok := c.outLimits[addr.String()]
return lim, ok
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.outLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.Out(host)
c.outMux.Lock()
defer c.outMux.Unlock()
c.outLimits[addr.String()] = lim
lim = c.limiter.Out(addr.String())
c.outLimits.Set(addr.String(), lim, 0)
return lim
}