add rate limiter
This commit is contained in:
ginuerzh
61
limiter/generator.go
Normal file
61
limiter/generator.go
Normal file
@ -0,0 +1,61 @@
|
||||
package limiter
|
||||
|
||||
import (
|
||||
"github.com/go-gost/core/limiter"
|
||||
)
|
||||
|
||||
type RateLimitGenerator interface {
|
||||
In() limiter.Limiter
|
||||
Out() limiter.Limiter
|
||||
}
|
||||
|
||||
type rateLimitGenerator struct {
|
||||
in int
|
||||
out int
|
||||
}
|
||||
|
||||
func NewRateLimitGenerator(in, out int) RateLimitGenerator {
|
||||
return &rateLimitGenerator{
|
||||
in: in,
|
||||
out: out,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *rateLimitGenerator) In() limiter.Limiter {
|
||||
if p == nil || p.in <= 0 {
|
||||
return nil
|
||||
}
|
||||
return NewLimiter(p.in)
|
||||
}
|
||||
|
||||
func (p *rateLimitGenerator) Out() limiter.Limiter {
|
||||
if p == nil || p.out <= 0 {
|
||||
return nil
|
||||
}
|
||||
return NewLimiter(p.out)
|
||||
}
|
||||
|
||||
type rateLimitSingleGenerator struct {
|
||||
in limiter.Limiter
|
||||
out limiter.Limiter
|
||||
}
|
||||
|
||||
func NewRateLimitSingleGenerator(in, out int) RateLimitGenerator {
|
||||
p := &rateLimitSingleGenerator{}
|
||||
if in > 0 {
|
||||
p.in = NewLimiter(in)
|
||||
}
|
||||
if out > 0 {
|
||||
p.out = NewLimiter(out)
|
||||
}
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *rateLimitSingleGenerator) In() limiter.Limiter {
|
||||
return p.in
|
||||
}
|
||||
|
||||
func (p *rateLimitSingleGenerator) Out() limiter.Limiter {
|
||||
return p.out
|
||||
}
|
||||
30
limiter/limiter.go
Normal file
30
limiter/limiter.go
Normal file
@ -0,0 +1,30 @@
|
||||
package limiter
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/go-gost/core/limiter"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
type llimiter struct {
|
||||
limiter *rate.Limiter
|
||||
}
|
||||
|
||||
func NewLimiter(r int) limiter.Limiter {
|
||||
return &llimiter{
|
||||
limiter: rate.NewLimiter(rate.Limit(r), r),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *llimiter) Wait(ctx context.Context, n int) int {
|
||||
if l.limiter.Burst() < n {
|
||||
n = l.limiter.Burst()
|
||||
}
|
||||
l.limiter.WaitN(ctx, n)
|
||||
return n
|
||||
}
|
||||
|
||||
func (l *llimiter) Limit() int {
|
||||
return l.limiter.Burst()
|
||||
}
|
||||
439
limiter/rate.go
439
limiter/rate.go
@ -1,83 +1,412 @@
|
||||
package limiter
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/alecthomas/units"
|
||||
"github.com/go-gost/core/limiter"
|
||||
"golang.org/x/time/rate"
|
||||
"github.com/go-gost/core/logger"
|
||||
"github.com/go-gost/x/internal/loader"
|
||||
"github.com/yl2chen/cidranger"
|
||||
)
|
||||
|
||||
type llimiter struct {
|
||||
limiter *rate.Limiter
|
||||
}
|
||||
const (
|
||||
GlobalLimitKey = "$"
|
||||
ConnLimitKey = "$$"
|
||||
)
|
||||
|
||||
func Limiter(r int) limiter.Limiter {
|
||||
return &llimiter{
|
||||
limiter: rate.NewLimiter(rate.Limit(r), r),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *llimiter) Limit(b int) int {
|
||||
if l.limiter.Burst() < b {
|
||||
b = l.limiter.Burst()
|
||||
}
|
||||
l.limiter.WaitN(context.Background(), b)
|
||||
return b
|
||||
}
|
||||
|
||||
type Generator interface {
|
||||
Generate() limiter.Limiter
|
||||
}
|
||||
|
||||
type limiterGenerator struct {
|
||||
limit int
|
||||
}
|
||||
|
||||
func NewGenerator(r int) Generator {
|
||||
return &limiterGenerator{limit: r}
|
||||
}
|
||||
|
||||
// Generate creates a new Limiter.
|
||||
func (g *limiterGenerator) Generate() limiter.Limiter {
|
||||
return Limiter(g.limit)
|
||||
}
|
||||
|
||||
type multiLimiter struct {
|
||||
type limiterGroup struct {
|
||||
limiters []limiter.Limiter
|
||||
}
|
||||
|
||||
func MultiLimiter(limiters ...limiter.Limiter) limiter.Limiter {
|
||||
return &multiLimiter{
|
||||
limiters: limiters,
|
||||
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()
|
||||
}
|
||||
|
||||
type options struct {
|
||||
limits []string
|
||||
fileLoader loader.Loader
|
||||
redisLoader loader.Loader
|
||||
period time.Duration
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
type Option func(opts *options)
|
||||
|
||||
func LimitsOption(limits ...string) Option {
|
||||
return func(opts *options) {
|
||||
opts.limits = limits
|
||||
}
|
||||
}
|
||||
|
||||
func (l *multiLimiter) Limit(b int) int {
|
||||
for i := range l.limiters {
|
||||
b = l.limiters[i].Limit(b)
|
||||
func ReloadPeriodOption(period time.Duration) Option {
|
||||
return func(opts *options) {
|
||||
opts.period = period
|
||||
}
|
||||
}
|
||||
|
||||
func FileLoaderOption(fileLoader loader.Loader) Option {
|
||||
return func(opts *options) {
|
||||
opts.fileLoader = fileLoader
|
||||
}
|
||||
}
|
||||
|
||||
func RedisLoaderOption(redisLoader loader.Loader) Option {
|
||||
return func(opts *options) {
|
||||
opts.redisLoader = redisLoader
|
||||
}
|
||||
}
|
||||
|
||||
func LoggerOption(logger logger.Logger) Option {
|
||||
return func(opts *options) {
|
||||
opts.logger = logger
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
type rateLimiter struct {
|
||||
input limiter.Limiter
|
||||
output limiter.Limiter
|
||||
ipLimits map[string]RateLimitGenerator
|
||||
cidrLimits cidranger.Ranger
|
||||
inLimits map[string]limiter.Limiter
|
||||
outLimits map[string]limiter.Limiter
|
||||
mu sync.RWMutex
|
||||
cancelFunc context.CancelFunc
|
||||
options options
|
||||
}
|
||||
|
||||
func RateLimiter(input, output limiter.Limiter) limiter.RateLimiter {
|
||||
if input == nil || output == nil {
|
||||
return nil
|
||||
func NewRateLimiter(opts ...Option) limiter.RateLimiter {
|
||||
var options options
|
||||
for _, opt := range opts {
|
||||
opt(&options)
|
||||
}
|
||||
return &rateLimiter{
|
||||
input: input,
|
||||
output: output,
|
||||
|
||||
ctx, cancel := context.WithCancel(context.TODO())
|
||||
lim := &rateLimiter{
|
||||
ipLimits: make(map[string]RateLimitGenerator),
|
||||
cidrLimits: cidranger.NewPCTrieRanger(),
|
||||
inLimits: make(map[string]limiter.Limiter),
|
||||
outLimits: make(map[string]limiter.Limiter),
|
||||
options: options,
|
||||
cancelFunc: cancel,
|
||||
}
|
||||
|
||||
if err := lim.reload(ctx); err != nil {
|
||||
options.logger.Warnf("reload: %v", err)
|
||||
}
|
||||
if lim.options.period > 0 {
|
||||
go lim.periodReload(ctx)
|
||||
}
|
||||
return lim
|
||||
}
|
||||
|
||||
func (l *rateLimiter) In(key string) limiter.Limiter {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
if lim, ok := l.inLimits[key]; ok {
|
||||
return lim
|
||||
}
|
||||
|
||||
var lims []limiter.Limiter
|
||||
|
||||
if ip := net.ParseIP(key); ip != nil {
|
||||
found := false
|
||||
if p := l.ipLimits[key]; p != nil {
|
||||
if lim := p.In(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
|
||||
if v, _ := p[0].(*cidrLimitEntry); v != nil {
|
||||
if lim := v.limit.In(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if p := l.ipLimits[ConnLimitKey]; p != nil {
|
||||
if lim := p.In(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
}
|
||||
}
|
||||
if p := l.ipLimits[GlobalLimitKey]; p != nil {
|
||||
if lim := p.In(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
}
|
||||
}
|
||||
|
||||
var lim limiter.Limiter
|
||||
if len(lims) > 0 {
|
||||
lim = newLimiterGroup(lims...)
|
||||
}
|
||||
l.inLimits[key] = lim
|
||||
|
||||
if lim != nil && l.options.logger != nil {
|
||||
l.options.logger.Debugf("input limit for %s: %d", key, lim.Limit())
|
||||
}
|
||||
|
||||
return lim
|
||||
}
|
||||
|
||||
func (l *rateLimiter) Out(key string) limiter.Limiter {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
if lim, ok := l.outLimits[key]; ok {
|
||||
return lim
|
||||
}
|
||||
|
||||
var lims []limiter.Limiter
|
||||
|
||||
if ip := net.ParseIP(key); ip != nil {
|
||||
found := false
|
||||
if p := l.ipLimits[key]; p != nil {
|
||||
if lim := p.Out(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
|
||||
if v, _ := p[0].(*cidrLimitEntry); v != nil {
|
||||
if lim := v.limit.Out(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if p := l.ipLimits[ConnLimitKey]; p != nil {
|
||||
if lim := p.Out(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
}
|
||||
}
|
||||
if p := l.ipLimits[GlobalLimitKey]; p != nil {
|
||||
if lim := p.Out(); lim != nil {
|
||||
lims = append(lims, lim)
|
||||
}
|
||||
}
|
||||
|
||||
var lim limiter.Limiter
|
||||
if len(lims) > 0 {
|
||||
lim = newLimiterGroup(lims...)
|
||||
}
|
||||
l.outLimits[key] = lim
|
||||
|
||||
if lim != nil && l.options.logger != nil {
|
||||
l.options.logger.Debugf("output limit for %s: %d", key, lim.Limit())
|
||||
}
|
||||
|
||||
return lim
|
||||
}
|
||||
|
||||
func (l *rateLimiter) periodReload(ctx context.Context) error {
|
||||
period := l.options.period
|
||||
if period < time.Second {
|
||||
period = time.Second
|
||||
}
|
||||
ticker := time.NewTicker(period)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
if err := l.reload(ctx); err != nil {
|
||||
l.options.logger.Warnf("reload: %v", err)
|
||||
// return err
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (l *rateLimiter) Input() limiter.Limiter {
|
||||
return l.input
|
||||
func (l *rateLimiter) reload(ctx context.Context) error {
|
||||
v, err := l.load(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
lines := append(l.options.limits, v...)
|
||||
|
||||
ipLimits := make(map[string]RateLimitGenerator)
|
||||
cidrLimits := cidranger.NewPCTrieRanger()
|
||||
|
||||
for _, s := range lines {
|
||||
key, in, out := l.parseLimit(s)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
switch key {
|
||||
case GlobalLimitKey:
|
||||
ipLimits[key] = NewRateLimitSingleGenerator(in, out)
|
||||
case ConnLimitKey:
|
||||
ipLimits[key] = NewRateLimitGenerator(in, out)
|
||||
default:
|
||||
if ip := net.ParseIP(key); ip != nil {
|
||||
ipLimits[key] = NewRateLimitGenerator(in, out)
|
||||
break
|
||||
}
|
||||
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
|
||||
cidrLimits.Insert(&cidrLimitEntry{
|
||||
ipNet: *ipNet,
|
||||
limit: NewRateLimitGenerator(in, out),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
l.ipLimits = ipLimits
|
||||
l.cidrLimits = cidrLimits
|
||||
l.inLimits = make(map[string]limiter.Limiter)
|
||||
l.outLimits = make(map[string]limiter.Limiter)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *rateLimiter) Output() limiter.Limiter {
|
||||
return l.output
|
||||
func (l *rateLimiter) load(ctx context.Context) (patterns []string, err error) {
|
||||
if l.options.fileLoader != nil {
|
||||
if lister, ok := l.options.fileLoader.(loader.Lister); ok {
|
||||
list, er := lister.List(ctx)
|
||||
if er != nil {
|
||||
l.options.logger.Warnf("file loader: %v", er)
|
||||
}
|
||||
for _, s := range list {
|
||||
if line := l.parseLine(s); line != "" {
|
||||
patterns = append(patterns, line)
|
||||
}
|
||||
}
|
||||
} 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...)
|
||||
}
|
||||
}
|
||||
}
|
||||
if l.options.redisLoader != nil {
|
||||
if lister, ok := l.options.redisLoader.(loader.Lister); ok {
|
||||
list, er := lister.List(ctx)
|
||||
if er != nil {
|
||||
l.options.logger.Warnf("redis loader: %v", er)
|
||||
}
|
||||
patterns = append(patterns, list...)
|
||||
} 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...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
l.options.logger.Debugf("load items %d", len(patterns))
|
||||
return
|
||||
}
|
||||
|
||||
func (l *rateLimiter) parsePatterns(r io.Reader) (patterns []string, err error) {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(r)
|
||||
for scanner.Scan() {
|
||||
if line := l.parseLine(scanner.Text()); line != "" {
|
||||
patterns = append(patterns, line)
|
||||
}
|
||||
}
|
||||
|
||||
err = scanner.Err()
|
||||
return
|
||||
}
|
||||
|
||||
func (l *rateLimiter) parseLine(s string) string {
|
||||
if n := strings.IndexByte(s, '#'); n >= 0 {
|
||||
s = s[:n]
|
||||
}
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
func (l *rateLimiter) parseLimit(s string) (key string, in, out int) {
|
||||
s = strings.Replace(s, "\t", " ", -1)
|
||||
s = strings.TrimSpace(s)
|
||||
var ss []string
|
||||
for _, v := range strings.Split(s, " ") {
|
||||
if v != "" {
|
||||
ss = append(ss, v)
|
||||
}
|
||||
}
|
||||
if len(ss) < 2 {
|
||||
return
|
||||
}
|
||||
|
||||
key = ss[0]
|
||||
if v, _ := units.ParseBase2Bytes(ss[1]); v > 0 {
|
||||
in = int(v)
|
||||
}
|
||||
if len(ss) > 2 {
|
||||
if v, _ := units.ParseBase2Bytes(ss[2]); v > 0 {
|
||||
out = int(v)
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (l *rateLimiter) Close() error {
|
||||
l.cancelFunc()
|
||||
if l.options.fileLoader != nil {
|
||||
l.options.fileLoader.Close()
|
||||
}
|
||||
if l.options.redisLoader != nil {
|
||||
l.options.redisLoader.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type cidrLimitEntry struct {
|
||||
ipNet net.IPNet
|
||||
limit RateLimitGenerator
|
||||
}
|
||||
|
||||
func (p *cidrLimitEntry) Network() net.IPNet {
|
||||
return p.ipNet
|
||||
}
|
||||
|
||||
340
limiter/wrapper/conn.go
Normal file
340
limiter/wrapper/conn.go
Normal file
@ -0,0 +1,340 @@
|
||||
package wrapper
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"syscall"
|
||||
|
||||
"github.com/go-gost/core/limiter"
|
||||
xnet "github.com/go-gost/x/internal/net"
|
||||
"github.com/go-gost/x/internal/net/udp"
|
||||
)
|
||||
|
||||
var (
|
||||
errUnsupport = errors.New("unsupported operation")
|
||||
)
|
||||
|
||||
// serverConn is a server side Conn with metrics supported.
|
||||
type serverConn struct {
|
||||
net.Conn
|
||||
rbuf bytes.Buffer
|
||||
raddr string
|
||||
rlimiter limiter.RateLimiter
|
||||
}
|
||||
|
||||
func WrapConn(rlimiter limiter.RateLimiter, c net.Conn) net.Conn {
|
||||
if rlimiter == nil {
|
||||
return c
|
||||
}
|
||||
host, _, _ := net.SplitHostPort(c.RemoteAddr().String())
|
||||
return &serverConn{
|
||||
Conn: c,
|
||||
rlimiter: rlimiter,
|
||||
raddr: host,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *serverConn) Read(b []byte) (n int, err error) {
|
||||
if c.rlimiter == nil ||
|
||||
c.rlimiter.In(c.raddr) == nil {
|
||||
return c.Conn.Read(b)
|
||||
}
|
||||
|
||||
limiter := c.rlimiter.In(c.raddr)
|
||||
|
||||
if c.rbuf.Len() > 0 {
|
||||
burst := len(b)
|
||||
if c.rbuf.Len() < burst {
|
||||
burst = c.rbuf.Len()
|
||||
}
|
||||
lim := limiter.Wait(context.Background(), burst)
|
||||
return c.rbuf.Read(b[:lim])
|
||||
}
|
||||
|
||||
nn, err := c.Conn.Read(b)
|
||||
if err != nil {
|
||||
return nn, err
|
||||
}
|
||||
|
||||
n = limiter.Wait(context.Background(), nn)
|
||||
if n < nn {
|
||||
if _, err = c.rbuf.Write(b[n:nn]); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (c *serverConn) Write(b []byte) (n int, err error) {
|
||||
if c.rlimiter == nil ||
|
||||
c.rlimiter.Out(c.raddr) == nil {
|
||||
return c.Conn.Write(b)
|
||||
}
|
||||
|
||||
limiter := c.rlimiter.Out(c.raddr)
|
||||
nn := 0
|
||||
for len(b) > 0 {
|
||||
nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))])
|
||||
n += nn
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
b = b[nn:]
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
|
||||
if sc, ok := c.Conn.(syscall.Conn); ok {
|
||||
rc, err = sc.SyscallConn()
|
||||
return
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
type packetConn struct {
|
||||
net.PacketConn
|
||||
rlimiter limiter.RateLimiter
|
||||
}
|
||||
|
||||
func WrapPacketConn(rlimiter limiter.RateLimiter, pc net.PacketConn) net.PacketConn {
|
||||
if rlimiter == nil {
|
||||
return pc
|
||||
}
|
||||
return &packetConn{
|
||||
PacketConn: pc,
|
||||
rlimiter: rlimiter,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
|
||||
for {
|
||||
n, addr, err = c.PacketConn.ReadFrom(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
|
||||
if c.rlimiter == nil || c.rlimiter.In(host) == nil {
|
||||
return
|
||||
}
|
||||
|
||||
limiter := c.rlimiter.In(host)
|
||||
// discard when exceed the limit size.
|
||||
if limiter.Wait(context.Background(), n) < n {
|
||||
continue
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
|
||||
if c.rlimiter != nil {
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
// discard when exceed the limit size.
|
||||
if limiter := c.rlimiter.Out(host); limiter != nil &&
|
||||
limiter.Wait(context.Background(), len(p)) < len(p) {
|
||||
n = len(p)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
return c.PacketConn.WriteTo(p, addr)
|
||||
}
|
||||
|
||||
type udpConn struct {
|
||||
net.PacketConn
|
||||
rlimiter limiter.RateLimiter
|
||||
}
|
||||
|
||||
func WrapUDPConn(rlimiter limiter.RateLimiter, pc net.PacketConn) udp.Conn {
|
||||
return &udpConn{
|
||||
PacketConn: pc,
|
||||
rlimiter: rlimiter,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *udpConn) RemoteAddr() net.Addr {
|
||||
if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
|
||||
return nc.RemoteAddr()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *udpConn) SetReadBuffer(n int) error {
|
||||
if nc, ok := c.PacketConn.(xnet.SetBuffer); ok {
|
||||
return nc.SetReadBuffer(n)
|
||||
}
|
||||
return errUnsupport
|
||||
}
|
||||
|
||||
func (c *udpConn) SetWriteBuffer(n int) error {
|
||||
if nc, ok := c.PacketConn.(xnet.SetBuffer); ok {
|
||||
return nc.SetWriteBuffer(n)
|
||||
}
|
||||
return errUnsupport
|
||||
}
|
||||
|
||||
func (c *udpConn) Read(b []byte) (n int, err error) {
|
||||
if nc, ok := c.PacketConn.(io.Reader); ok {
|
||||
n, err = nc.Read(b)
|
||||
return
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
|
||||
for {
|
||||
n, addr, err = c.PacketConn.ReadFrom(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
|
||||
if c.rlimiter == nil || c.rlimiter.In(host) == nil {
|
||||
return
|
||||
}
|
||||
limiter := c.rlimiter.In(host)
|
||||
// discard when exceed the limit size.
|
||||
if limiter.Wait(context.Background(), n) < n {
|
||||
continue
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
|
||||
if nc, ok := c.PacketConn.(udp.ReadUDP); ok {
|
||||
for {
|
||||
n, addr, err = nc.ReadFromUDP(b)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
|
||||
if c.rlimiter == nil || c.rlimiter.In(host) == nil {
|
||||
return
|
||||
}
|
||||
limiter := c.rlimiter.In(host)
|
||||
// discard when exceed the limit size.
|
||||
if limiter.Wait(context.Background(), n) < n {
|
||||
continue
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
|
||||
if nc, ok := c.PacketConn.(udp.ReadUDP); ok {
|
||||
for {
|
||||
n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
|
||||
if c.rlimiter == nil || c.rlimiter.In(host) == nil {
|
||||
return
|
||||
}
|
||||
limiter := c.rlimiter.In(host)
|
||||
// discard when exceed the limit size.
|
||||
if limiter.Wait(context.Background(), n) < n {
|
||||
continue
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) Write(b []byte) (n int, err error) {
|
||||
if nc, ok := c.PacketConn.(io.Writer); ok {
|
||||
n, err = nc.Write(b)
|
||||
return
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
|
||||
if c.rlimiter != nil {
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
// discard when exceed the limit size.
|
||||
if limiter := c.rlimiter.Out(host); limiter != nil &&
|
||||
limiter.Wait(context.Background(), len(p)) < len(p) {
|
||||
n = len(p)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
n, err = c.PacketConn.WriteTo(p, addr)
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
|
||||
if c.rlimiter != nil {
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
// discard when exceed the limit size.
|
||||
if limiter := c.rlimiter.Out(host); limiter != nil &&
|
||||
limiter.Wait(context.Background(), len(b)) < len(b) {
|
||||
n = len(b)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
|
||||
n, err = nc.WriteToUDP(b, addr)
|
||||
return
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
|
||||
if c.rlimiter != nil {
|
||||
host, _, _ := net.SplitHostPort(addr.String())
|
||||
// discard when exceed the limit size.
|
||||
if limiter := c.rlimiter.Out(host); limiter != nil &&
|
||||
limiter.Wait(context.Background(), len(b)) < len(b) {
|
||||
n = len(b)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
|
||||
n, oobn, err = nc.WriteMsgUDP(b, oob, addr)
|
||||
return
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) {
|
||||
if nc, ok := c.PacketConn.(xnet.SyscallConn); ok {
|
||||
return nc.SyscallConn()
|
||||
}
|
||||
err = errUnsupport
|
||||
return
|
||||
}
|
||||
|
||||
func (c *udpConn) SetDSCP(n int) error {
|
||||
if nc, ok := c.PacketConn.(xnet.SetDSCP); ok {
|
||||
return nc.SetDSCP(n)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
32
limiter/wrapper/listener.go
Normal file
32
limiter/wrapper/listener.go
Normal file
@ -0,0 +1,32 @@
|
||||
package wrapper
|
||||
|
||||
import (
|
||||
"net"
|
||||
|
||||
"github.com/go-gost/core/limiter"
|
||||
)
|
||||
|
||||
type listener struct {
|
||||
net.Listener
|
||||
rlimiter limiter.RateLimiter
|
||||
}
|
||||
|
||||
func WrapListener(rlimiter limiter.RateLimiter, ln net.Listener) net.Listener {
|
||||
if rlimiter == nil {
|
||||
return ln
|
||||
}
|
||||
|
||||
return &listener{
|
||||
rlimiter: rlimiter,
|
||||
Listener: ln,
|
||||
}
|
||||
}
|
||||
|
||||
func (ln *listener) Accept() (net.Conn, error) {
|
||||
c, err := ln.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return WrapConn(ln.rlimiter, c), nil
|
||||
}
|
||||
Reference in New Issue
Block a user