add rate limiter
This commit is contained in:
parent
e23da0f319
commit
4c2131ca6d
219
admission/wrapper/conn.go
Normal file
219
admission/wrapper/conn.go
Normal file
@ -0,0 +1,219 @@
|
|||||||
|
package wrapper
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"syscall"
|
||||||
|
|
||||||
|
"github.com/go-gost/core/admission"
|
||||||
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
"github.com/go-gost/x/internal/net/udp"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
errUnsupport = errors.New("unsupported operation")
|
||||||
|
)
|
||||||
|
|
||||||
|
type serverConn struct {
|
||||||
|
net.Conn
|
||||||
|
admission admission.Admission
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapConn(admission admission.Admission, c net.Conn) net.Conn {
|
||||||
|
if admission == nil {
|
||||||
|
return c
|
||||||
|
}
|
||||||
|
return &serverConn{
|
||||||
|
Conn: c,
|
||||||
|
admission: admission,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *serverConn) Read(b []byte) (n int, err error) {
|
||||||
|
if c.admission != nil &&
|
||||||
|
!c.admission.Admit(c.RemoteAddr().String()) {
|
||||||
|
err = io.EOF
|
||||||
|
return
|
||||||
|
}
|
||||||
|
return c.Conn.Read(b)
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
admission admission.Admission
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapPacketConn(admission admission.Admission, pc net.PacketConn) net.PacketConn {
|
||||||
|
if admission == nil {
|
||||||
|
return pc
|
||||||
|
}
|
||||||
|
return &packetConn{
|
||||||
|
PacketConn: pc,
|
||||||
|
admission: admission,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
if c.admission != nil &&
|
||||||
|
!c.admission.Admit(addr.String()) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type udpConn struct {
|
||||||
|
net.PacketConn
|
||||||
|
admission admission.Admission
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapUDPConn(admission admission.Admission, pc net.PacketConn) udp.Conn {
|
||||||
|
return &udpConn{
|
||||||
|
PacketConn: pc,
|
||||||
|
admission: admission,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
if c.admission != nil &&
|
||||||
|
!c.admission.Admit(addr.String()) {
|
||||||
|
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
|
||||||
|
}
|
||||||
|
if c.admission != nil &&
|
||||||
|
!c.admission.Admit(addr.String()) {
|
||||||
|
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
|
||||||
|
}
|
||||||
|
if c.admission != nil &&
|
||||||
|
!c.admission.Admit(addr.String()) {
|
||||||
|
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) {
|
||||||
|
n, err = c.PacketConn.WriteTo(p, addr)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
|
||||||
|
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 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
|
||||||
|
}
|
37
admission/wrapper/listener.go
Normal file
37
admission/wrapper/listener.go
Normal file
@ -0,0 +1,37 @@
|
|||||||
|
package wrapper
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net"
|
||||||
|
|
||||||
|
"github.com/go-gost/core/admission"
|
||||||
|
)
|
||||||
|
|
||||||
|
type listener struct {
|
||||||
|
net.Listener
|
||||||
|
admission admission.Admission
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapListener(admission admission.Admission, ln net.Listener) net.Listener {
|
||||||
|
if admission == nil {
|
||||||
|
return ln
|
||||||
|
}
|
||||||
|
return &listener{
|
||||||
|
Listener: ln,
|
||||||
|
admission: admission,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (ln *listener) Accept() (net.Conn, error) {
|
||||||
|
for {
|
||||||
|
c, err := ln.Listener.Accept()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if ln.admission != nil &&
|
||||||
|
!ln.admission.Admit(c.RemoteAddr().String()) {
|
||||||
|
c.Close()
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
return c, err
|
||||||
|
}
|
||||||
|
}
|
@ -186,18 +186,19 @@ type RecorderObject struct {
|
|||||||
|
|
||||||
type LimiterConfig struct {
|
type LimiterConfig struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
RateLimit *RateLimitConfig `yaml:"rate" json:"rate"`
|
Rate *RateLimiterConfig `yaml:"rate" json:"rate"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type RateLimitConfig struct {
|
type RateLimiterConfig struct {
|
||||||
Input string `yaml:",omitempty" json:"input,omitempty"`
|
Limits []string `yaml:",omitempty" json:"limits,omitempty"`
|
||||||
Output string `yaml:",omitempty" json:"output,omitempty"`
|
Reload time.Duration `yaml:",omitempty" json:"reload,omitempty"`
|
||||||
Conn *LimitConfig `yaml:",omitempty" json:"conn,omitempty"`
|
File *FileLoader `yaml:",omitempty" json:"file,omitempty"`
|
||||||
|
Redis *RedisLoader `yaml:",omitempty" json:"redis,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type LimitConfig struct {
|
type LimitConfig struct {
|
||||||
Input string `yaml:",omitempty" json:"input,omitempty"`
|
In string `yaml:",omitempty" json:"in,omitempty"`
|
||||||
Output string `yaml:",omitempty" json:"output,omitempty"`
|
Out string `yaml:",omitempty" json:"out,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type ListenerConfig struct {
|
type ListenerConfig struct {
|
||||||
|
@ -4,7 +4,6 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"net/url"
|
"net/url"
|
||||||
|
|
||||||
"github.com/alecthomas/units"
|
|
||||||
"github.com/go-gost/core/admission"
|
"github.com/go-gost/core/admission"
|
||||||
"github.com/go-gost/core/auth"
|
"github.com/go-gost/core/auth"
|
||||||
"github.com/go-gost/core/bypass"
|
"github.com/go-gost/core/bypass"
|
||||||
@ -323,33 +322,41 @@ func defaultChainSelector() selector.Selector[chain.Chainer] {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func ParseRateLimiter(cfg *config.LimiterConfig) (lim limiter.RateLimiter) {
|
func ParseRateLimiter(cfg *config.LimiterConfig) (lim limiter.RateLimiter) {
|
||||||
if cfg == nil || cfg.RateLimit == nil {
|
if cfg == nil || cfg.Rate == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
var rlimiters []limiter.Limiter
|
var opts []xlimiter.Option
|
||||||
var wlimiters []limiter.Limiter
|
|
||||||
if cfg.RateLimit.Conn != nil {
|
|
||||||
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Conn.Input); v > 0 {
|
|
||||||
rlimiters = append(rlimiters, xlimiter.Limiter(int(v)))
|
|
||||||
}
|
|
||||||
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Conn.Output); v > 0 {
|
|
||||||
wlimiters = append(wlimiters, xlimiter.Limiter(int(v)))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Input); v > 0 {
|
|
||||||
rlimiters = append(rlimiters, xlimiter.Limiter(int(v)))
|
|
||||||
}
|
|
||||||
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Output); v > 0 {
|
|
||||||
wlimiters = append(wlimiters, xlimiter.Limiter(int(v)))
|
|
||||||
}
|
|
||||||
|
|
||||||
var input, output limiter.Limiter
|
if cfg.Rate.File != nil && cfg.Rate.File.Path != "" {
|
||||||
if len(rlimiters) > 0 {
|
opts = append(opts, xlimiter.FileLoaderOption(loader.FileLoader(cfg.Rate.File.Path)))
|
||||||
input = xlimiter.MultiLimiter(rlimiters...)
|
|
||||||
}
|
}
|
||||||
if len(wlimiters) > 0 {
|
if cfg.Rate.Redis != nil && cfg.Rate.Redis.Addr != "" {
|
||||||
output = xlimiter.MultiLimiter(wlimiters...)
|
switch cfg.Rate.Redis.Type {
|
||||||
|
case "list": // redis list
|
||||||
|
opts = append(opts, xlimiter.RedisLoaderOption(loader.RedisListLoader(
|
||||||
|
cfg.Rate.Redis.Addr,
|
||||||
|
loader.DBRedisLoaderOption(cfg.Rate.Redis.DB),
|
||||||
|
loader.PasswordRedisLoaderOption(cfg.Rate.Redis.Password),
|
||||||
|
loader.KeyRedisLoaderOption(cfg.Rate.Redis.Key),
|
||||||
|
)))
|
||||||
|
default: // redis set
|
||||||
|
opts = append(opts, xlimiter.RedisLoaderOption(loader.RedisSetLoader(
|
||||||
|
cfg.Rate.Redis.Addr,
|
||||||
|
loader.DBRedisLoaderOption(cfg.Rate.Redis.DB),
|
||||||
|
loader.PasswordRedisLoaderOption(cfg.Rate.Redis.Password),
|
||||||
|
loader.KeyRedisLoaderOption(cfg.Rate.Redis.Key),
|
||||||
|
)))
|
||||||
}
|
}
|
||||||
return xlimiter.RateLimiter(input, output)
|
}
|
||||||
|
opts = append(opts,
|
||||||
|
xlimiter.LimitsOption(cfg.Rate.Limits...),
|
||||||
|
xlimiter.ReloadPeriodOption(cfg.Rate.Reload),
|
||||||
|
xlimiter.LoggerOption(logger.Default().WithFields(map[string]any{
|
||||||
|
"kind": "limiter",
|
||||||
|
"hosts": cfg.Name,
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
|
||||||
|
return xlimiter.NewRateLimiter(opts...)
|
||||||
}
|
}
|
||||||
|
@ -1,5 +1,28 @@
|
|||||||
package net
|
package net
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net"
|
||||||
|
"syscall"
|
||||||
|
)
|
||||||
|
|
||||||
|
type SetBuffer interface {
|
||||||
|
SetReadBuffer(bytes int) error
|
||||||
|
SetWriteBuffer(bytes int) error
|
||||||
|
}
|
||||||
|
|
||||||
|
type SyscallConn interface {
|
||||||
|
SyscallConn() (syscall.RawConn, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
type RemoteAddr interface {
|
||||||
|
RemoteAddr() net.Addr
|
||||||
|
}
|
||||||
|
|
||||||
|
// tcpraw.TCPConn
|
||||||
|
type SetDSCP interface {
|
||||||
|
SetDSCP(int) error
|
||||||
|
}
|
||||||
|
|
||||||
func IsIPv4(address string) bool {
|
func IsIPv4(address string) bool {
|
||||||
return address != "" && address[0] != ':' && address[0] != '['
|
return address != "" && address[0] != ':' && address[0] != '['
|
||||||
}
|
}
|
||||||
|
29
internal/net/udp/conn.go
Normal file
29
internal/net/udp/conn.go
Normal file
@ -0,0 +1,29 @@
|
|||||||
|
package udp
|
||||||
|
|
||||||
|
import (
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
|
||||||
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Conn interface {
|
||||||
|
net.PacketConn
|
||||||
|
io.Reader
|
||||||
|
io.Writer
|
||||||
|
ReadUDP
|
||||||
|
WriteUDP
|
||||||
|
xnet.SetBuffer
|
||||||
|
xnet.SyscallConn
|
||||||
|
xnet.RemoteAddr
|
||||||
|
}
|
||||||
|
|
||||||
|
type ReadUDP interface {
|
||||||
|
ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
|
||||||
|
ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
|
||||||
|
}
|
||||||
|
|
||||||
|
type WriteUDP interface {
|
||||||
|
WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
|
||||||
|
WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
|
||||||
|
}
|
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()
|
||||||
|
}
|
441
limiter/rate.go
441
limiter/rate.go
@ -1,83 +1,412 @@
|
|||||||
package limiter
|
package limiter
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bufio"
|
||||||
"context"
|
"context"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/alecthomas/units"
|
||||||
"github.com/go-gost/core/limiter"
|
"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 {
|
const (
|
||||||
limiter *rate.Limiter
|
GlobalLimitKey = "$"
|
||||||
}
|
ConnLimitKey = "$$"
|
||||||
|
)
|
||||||
|
|
||||||
func Limiter(r int) limiter.Limiter {
|
type limiterGroup struct {
|
||||||
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 {
|
|
||||||
limiters []limiter.Limiter
|
limiters []limiter.Limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
func MultiLimiter(limiters ...limiter.Limiter) limiter.Limiter {
|
func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
|
||||||
return &multiLimiter{
|
sort.Slice(limiters, func(i, j int) bool {
|
||||||
limiters: limiters,
|
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 {
|
func ReloadPeriodOption(period time.Duration) Option {
|
||||||
for i := range l.limiters {
|
return func(opts *options) {
|
||||||
b = l.limiters[i].Limit(b)
|
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 {
|
type rateLimiter struct {
|
||||||
input limiter.Limiter
|
ipLimits map[string]RateLimitGenerator
|
||||||
output limiter.Limiter
|
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 {
|
func NewRateLimiter(opts ...Option) limiter.RateLimiter {
|
||||||
if input == nil || output == nil {
|
var options options
|
||||||
|
for _, opt := range opts {
|
||||||
|
opt(&options)
|
||||||
|
}
|
||||||
|
|
||||||
|
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) 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
|
return nil
|
||||||
}
|
|
||||||
return &rateLimiter{
|
|
||||||
input: input,
|
|
||||||
output: output,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *rateLimiter) Input() limiter.Limiter {
|
func (l *rateLimiter) load(ctx context.Context) (patterns []string, err error) {
|
||||||
return l.input
|
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) Output() limiter.Limiter {
|
func (l *rateLimiter) parsePatterns(r io.Reader) (patterns []string, err error) {
|
||||||
return l.output
|
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
|
||||||
|
}
|
@ -9,10 +9,13 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/miekg/dns"
|
"github.com/miekg/dns"
|
||||||
)
|
)
|
||||||
@ -114,6 +117,8 @@ func (l *dnsListener) Accept() (conn net.Conn, err error) {
|
|||||||
select {
|
select {
|
||||||
case conn = <-l.cqueue:
|
case conn = <-l.cqueue:
|
||||||
conn = metrics.WrapConn(l.options.Service, conn)
|
conn = metrics.WrapConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapConn(l.options.RateLimiter, conn)
|
||||||
case err, ok = <-l.errChan:
|
case err, ok = <-l.errChan:
|
||||||
if !ok {
|
if !ok {
|
||||||
err = listener.ErrClosed
|
err = listener.ErrClosed
|
||||||
|
@ -7,8 +7,10 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/xtaci/tcpraw"
|
"github.com/xtaci/tcpraw"
|
||||||
)
|
)
|
||||||
@ -50,6 +52,8 @@ func (l *ftcpListener) Init(md md.Metadata) (err error) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn = metrics.WrapPacketConn(l.options.Service, conn)
|
conn = metrics.WrapPacketConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapPacketConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapPacketConn(l.options.RateLimiter, conn)
|
||||||
|
|
||||||
l.ln = udp.NewListener(
|
l.ln = udp.NewListener(
|
||||||
conn,
|
conn,
|
||||||
|
@ -3,13 +3,14 @@ package grpc
|
|||||||
import (
|
import (
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
pb "github.com/go-gost/x/internal/util/grpc/proto"
|
pb "github.com/go-gost/x/internal/util/grpc/proto"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"google.golang.org/grpc"
|
"google.golang.org/grpc"
|
||||||
"google.golang.org/grpc/credentials"
|
"google.golang.org/grpc/credentials"
|
||||||
@ -55,6 +56,7 @@ func (l *grpcListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
var opts []grpc.ServerOption
|
var opts []grpc.ServerOption
|
||||||
if !l.md.insecure {
|
if !l.md.insecure {
|
||||||
|
@ -7,12 +7,13 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httputil"
|
"net/http/httputil"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"golang.org/x/net/http2"
|
"golang.org/x/net/http2"
|
||||||
"golang.org/x/net/http2/h2c"
|
"golang.org/x/net/http2/h2c"
|
||||||
@ -77,6 +78,7 @@ func (l *h2Listener) Init(md md.Metadata) (err error) {
|
|||||||
l.addr = ln.Addr()
|
l.addr = ln.Addr()
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
if l.h2c {
|
if l.h2c {
|
||||||
l.server.Handler = h2c.NewHandler(
|
l.server.Handler = h2c.NewHandler(
|
||||||
|
@ -5,13 +5,14 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
mdx "github.com/go-gost/x/metadata"
|
mdx "github.com/go-gost/x/metadata"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"golang.org/x/net/http2"
|
"golang.org/x/net/http2"
|
||||||
)
|
)
|
||||||
@ -66,6 +67,7 @@ func (l *http2Listener) Init(md md.Metadata) (err error) {
|
|||||||
l.addr = ln.Addr()
|
l.addr = ln.Addr()
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
ln = tls.NewListener(
|
ln = tls.NewListener(
|
||||||
ln,
|
ln,
|
||||||
|
@ -6,9 +6,11 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
pht_util "github.com/go-gost/x/internal/util/pht"
|
pht_util "github.com/go-gost/x/internal/util/pht"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/lucas-clemente/quic-go"
|
"github.com/lucas-clemente/quic-go"
|
||||||
)
|
)
|
||||||
@ -74,7 +76,11 @@ func (l *http3Listener) Accept() (conn net.Conn, err error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
return metrics.WrapConn(l.options.Service, conn), nil
|
|
||||||
|
conn = metrics.WrapConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapConn(l.options.RateLimiter, conn)
|
||||||
|
return conn, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l *http3Listener) Addr() net.Addr {
|
func (l *http3Listener) Addr() net.Addr {
|
||||||
|
@ -4,12 +4,13 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
icmp_pkg "github.com/go-gost/x/internal/util/icmp"
|
icmp_pkg "github.com/go-gost/x/internal/util/icmp"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/lucas-clemente/quic-go"
|
"github.com/lucas-clemente/quic-go"
|
||||||
"golang.org/x/net/icmp"
|
"golang.org/x/net/icmp"
|
||||||
@ -57,6 +58,7 @@ func (l *icmpListener) Init(md md.Metadata) (err error) {
|
|||||||
conn = icmp_pkg.ServerConn(conn)
|
conn = icmp_pkg.ServerConn(conn)
|
||||||
conn = metrics.WrapPacketConn(l.options.Service, conn)
|
conn = metrics.WrapPacketConn(l.options.Service, conn)
|
||||||
conn = admission.WrapPacketConn(l.options.Admission, conn)
|
conn = admission.WrapPacketConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapPacketConn(l.options.RateLimiter, conn)
|
||||||
|
|
||||||
config := &quic.Config{
|
config := &quic.Config{
|
||||||
KeepAlivePeriod: l.md.keepAlivePeriod,
|
KeepAlivePeriod: l.md.keepAlivePeriod,
|
||||||
|
@ -4,13 +4,14 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
kcp_util "github.com/go-gost/x/internal/util/kcp"
|
kcp_util "github.com/go-gost/x/internal/util/kcp"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/xtaci/kcp-go/v5"
|
"github.com/xtaci/kcp-go/v5"
|
||||||
"github.com/xtaci/smux"
|
"github.com/xtaci/smux"
|
||||||
@ -75,6 +76,7 @@ func (l *kcpListener) Init(md md.Metadata) (err error) {
|
|||||||
|
|
||||||
conn = metrics.WrapUDPConn(l.options.Service, conn)
|
conn = metrics.WrapUDPConn(l.options.Service, conn)
|
||||||
conn = admission.WrapUDPConn(l.options.Admission, conn)
|
conn = admission.WrapUDPConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapUDPConn(l.options.RateLimiter, conn)
|
||||||
|
|
||||||
ln, err := kcp.ServeConn(
|
ln, err := kcp.ServeConn(
|
||||||
kcp_util.BlockCrypt(config.Key, config.Crypt, kcp_util.DefaultSalt),
|
kcp_util.BlockCrypt(config.Key, config.Crypt, kcp_util.DefaultSalt),
|
||||||
|
@ -4,12 +4,13 @@ import (
|
|||||||
"crypto/tls"
|
"crypto/tls"
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/xtaci/smux"
|
"github.com/xtaci/smux"
|
||||||
)
|
)
|
||||||
@ -54,6 +55,7 @@ func (l *mtlsListener) Init(md md.Metadata) (err error) {
|
|||||||
|
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
l.Listener = tls.NewListener(ln, l.options.TLSConfig)
|
l.Listener = tls.NewListener(ln, l.options.TLSConfig)
|
||||||
|
|
||||||
l.cqueue = make(chan net.Conn, l.md.backlog)
|
l.cqueue = make(chan net.Conn, l.md.backlog)
|
||||||
|
@ -6,13 +6,14 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httputil"
|
"net/http/httputil"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
ws_util "github.com/go-gost/x/internal/util/ws"
|
ws_util "github.com/go-gost/x/internal/util/ws"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/gorilla/websocket"
|
"github.com/gorilla/websocket"
|
||||||
"github.com/xtaci/smux"
|
"github.com/xtaci/smux"
|
||||||
@ -96,6 +97,7 @@ func (l *mwsListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
if l.tlsEnabled {
|
if l.tlsEnabled {
|
||||||
ln = tls.NewListener(ln, l.options.TLSConfig)
|
ln = tls.NewListener(ln, l.options.TLSConfig)
|
||||||
|
@ -3,11 +3,13 @@ package http
|
|||||||
import (
|
import (
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
|
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
@ -49,6 +51,7 @@ func (l *obfsListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
l.Listener = ln
|
l.Listener = ln
|
||||||
return
|
return
|
||||||
|
@ -3,12 +3,13 @@ package tls
|
|||||||
import (
|
import (
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -49,6 +50,7 @@ func (l *obfsListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
l.Listener = ln
|
l.Listener = ln
|
||||||
return
|
return
|
||||||
|
@ -8,9 +8,11 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
pht_util "github.com/go-gost/x/internal/util/pht"
|
pht_util "github.com/go-gost/x/internal/util/pht"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -89,6 +91,8 @@ func (l *phtListener) Accept() (conn net.Conn, err error) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn = metrics.WrapConn(l.options.Service, conn)
|
conn = metrics.WrapConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapConn(l.options.RateLimiter, conn)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -7,9 +7,11 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
quic_util "github.com/go-gost/x/internal/util/quic"
|
quic_util "github.com/go-gost/x/internal/util/quic"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/lucas-clemente/quic-go"
|
"github.com/lucas-clemente/quic-go"
|
||||||
)
|
)
|
||||||
@ -99,6 +101,8 @@ func (l *quicListener) Accept() (conn net.Conn, err error) {
|
|||||||
select {
|
select {
|
||||||
case conn = <-l.cqueue:
|
case conn = <-l.cqueue:
|
||||||
conn = metrics.WrapConn(l.options.Service, conn)
|
conn = metrics.WrapConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapConn(l.options.RateLimiter, conn)
|
||||||
case err, ok = <-l.errChan:
|
case err, ok = <-l.errChan:
|
||||||
if !ok {
|
if !ok {
|
||||||
err = listener.ErrClosed
|
err = listener.ErrClosed
|
||||||
|
@ -7,8 +7,10 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -54,7 +56,10 @@ func (l *redirectListener) Init(md md.Metadata) (err error) {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
l.ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
l.ln = ln
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -6,7 +6,9 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -52,6 +54,8 @@ func (l *redirectListener) Accept() (conn net.Conn, err error) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn = metrics.WrapConn(l.options.Service, conn)
|
conn = metrics.WrapConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapConn(l.options.Admission, conn)
|
||||||
|
conn = limiter.WrapConn(l.options.RateLimiter, conn)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -8,8 +8,10 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -76,6 +78,8 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) {
|
|||||||
return nil, listener.NewAcceptError(err)
|
return nil, listener.NewAcceptError(err)
|
||||||
}
|
}
|
||||||
l.ln = metrics.WrapListener(l.options.Service, l.ln)
|
l.ln = metrics.WrapListener(l.options.Service, l.ln)
|
||||||
|
l.ln = admission.WrapListener(l.options.Admission, l.ln)
|
||||||
|
l.ln = limiter.WrapListener(l.options.RateLimiter, l.ln)
|
||||||
}
|
}
|
||||||
conn, err = l.ln.Accept()
|
conn, err = l.ln.Accept()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
@ -8,8 +8,10 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -88,7 +90,9 @@ func (l *rudpListener) Accept() (conn net.Conn, err error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if pc, ok := conn.(net.PacketConn); ok {
|
if pc, ok := conn.(net.PacketConn); ok {
|
||||||
conn = metrics.WrapUDPConn(l.options.Service, pc)
|
uc := metrics.WrapUDPConn(l.options.Service, pc)
|
||||||
|
uc = admission.WrapUDPConn(l.options.Admission, uc)
|
||||||
|
conn = limiter.WrapUDPConn(l.options.RateLimiter, uc)
|
||||||
}
|
}
|
||||||
|
|
||||||
return
|
return
|
||||||
|
@ -5,13 +5,14 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
ssh_util "github.com/go-gost/x/internal/util/ssh"
|
ssh_util "github.com/go-gost/x/internal/util/ssh"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"golang.org/x/crypto/ssh"
|
"golang.org/x/crypto/ssh"
|
||||||
)
|
)
|
||||||
@ -57,6 +58,7 @@ func (l *sshListener) Init(md md.Metadata) (err error) {
|
|||||||
|
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
l.Listener = ln
|
l.Listener = ln
|
||||||
|
|
||||||
config := &ssh.ServerConfig{
|
config := &ssh.ServerConfig{
|
||||||
|
@ -7,14 +7,15 @@ import (
|
|||||||
"strconv"
|
"strconv"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
ssh_util "github.com/go-gost/x/internal/util/ssh"
|
ssh_util "github.com/go-gost/x/internal/util/ssh"
|
||||||
sshd_util "github.com/go-gost/x/internal/util/sshd"
|
sshd_util "github.com/go-gost/x/internal/util/sshd"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"golang.org/x/crypto/ssh"
|
"golang.org/x/crypto/ssh"
|
||||||
)
|
)
|
||||||
@ -66,6 +67,7 @@ func (l *sshdListener) Init(md md.Metadata) (err error) {
|
|||||||
|
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
l.Listener = ln
|
l.Listener = ln
|
||||||
|
|
||||||
config := &ssh.ServerConfig{
|
config := &ssh.ServerConfig{
|
||||||
|
@ -6,9 +6,10 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
mdata "github.com/go-gost/core/metadata"
|
mdata "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
mdx "github.com/go-gost/x/metadata"
|
mdx "github.com/go-gost/x/metadata"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -79,6 +80,7 @@ func (l *tapListener) Init(md mdata.Metadata) (err error) {
|
|||||||
raddr: &net.IPAddr{IP: ip},
|
raddr: &net.IPAddr{IP: ip},
|
||||||
}
|
}
|
||||||
c = metrics.WrapConn(l.options.Service, c)
|
c = metrics.WrapConn(l.options.Service, c)
|
||||||
|
c = limiter.WrapConn(l.options.RateLimiter, c)
|
||||||
c = withMetadata(mdx.NewMetadata(map[string]any{
|
c = withMetadata(mdx.NewMetadata(map[string]any{
|
||||||
"config": l.md.config,
|
"config": l.md.config,
|
||||||
}), c)
|
}), c)
|
||||||
|
@ -3,12 +3,13 @@ package tcp
|
|||||||
import (
|
import (
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
limiter "github.com/go-gost/core/limiter/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -49,6 +50,7 @@ func (l *tcpListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
l.ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
l.ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
return
|
return
|
||||||
|
@ -4,12 +4,13 @@ import (
|
|||||||
"crypto/tls"
|
"crypto/tls"
|
||||||
"net"
|
"net"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -50,6 +51,7 @@ func (l *tlsListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
l.ln = tls.NewListener(ln, l.options.TLSConfig)
|
l.ln = tls.NewListener(ln, l.options.TLSConfig)
|
||||||
|
|
||||||
|
@ -8,9 +8,10 @@ import (
|
|||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
mdata "github.com/go-gost/core/metadata"
|
mdata "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
mdx "github.com/go-gost/x/metadata"
|
mdx "github.com/go-gost/x/metadata"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -88,6 +89,7 @@ func (l *tunListener) listenLoop() {
|
|||||||
cancel: cancel,
|
cancel: cancel,
|
||||||
}
|
}
|
||||||
c = metrics.WrapConn(l.options.Service, c)
|
c = metrics.WrapConn(l.options.Service, c)
|
||||||
|
c = limiter.WrapConn(l.options.RateLimiter, c)
|
||||||
c = withMetadata(mdx.NewMetadata(map[string]any{
|
c = withMetadata(mdx.NewMetadata(map[string]any{
|
||||||
"config": l.md.config,
|
"config": l.md.config,
|
||||||
}), c)
|
}), c)
|
||||||
|
@ -4,12 +4,13 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
|
|
||||||
"github.com/go-gost/core/common/net/udp"
|
"github.com/go-gost/core/common/net/udp"
|
||||||
limiter "github.com/go-gost/core/limiter/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -55,6 +56,7 @@ func (l *udpListener) Init(md md.Metadata) (err error) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn = metrics.WrapPacketConn(l.options.Service, conn)
|
conn = metrics.WrapPacketConn(l.options.Service, conn)
|
||||||
|
conn = admission.WrapPacketConn(l.options.Admission, conn)
|
||||||
conn = limiter.WrapPacketConn(l.options.RateLimiter, conn)
|
conn = limiter.WrapPacketConn(l.options.RateLimiter, conn)
|
||||||
|
|
||||||
l.ln = udp.NewListener(conn, &udp.ListenConfig{
|
l.ln = udp.NewListener(conn, &udp.ListenConfig{
|
||||||
|
@ -6,13 +6,14 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httputil"
|
"net/http/httputil"
|
||||||
|
|
||||||
admission "github.com/go-gost/core/admission/wrapper"
|
|
||||||
"github.com/go-gost/core/listener"
|
"github.com/go-gost/core/listener"
|
||||||
"github.com/go-gost/core/logger"
|
"github.com/go-gost/core/logger"
|
||||||
md "github.com/go-gost/core/metadata"
|
md "github.com/go-gost/core/metadata"
|
||||||
metrics "github.com/go-gost/core/metrics/wrapper"
|
admission "github.com/go-gost/x/admission/wrapper"
|
||||||
xnet "github.com/go-gost/x/internal/net"
|
xnet "github.com/go-gost/x/internal/net"
|
||||||
ws_util "github.com/go-gost/x/internal/util/ws"
|
ws_util "github.com/go-gost/x/internal/util/ws"
|
||||||
|
limiter "github.com/go-gost/x/limiter/wrapper"
|
||||||
|
metrics "github.com/go-gost/x/metrics/wrapper"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
"github.com/gorilla/websocket"
|
"github.com/gorilla/websocket"
|
||||||
)
|
)
|
||||||
@ -91,6 +92,7 @@ func (l *wsListener) Init(md md.Metadata) (err error) {
|
|||||||
}
|
}
|
||||||
ln = metrics.WrapListener(l.options.Service, ln)
|
ln = metrics.WrapListener(l.options.Service, ln)
|
||||||
ln = admission.WrapListener(l.options.Admission, ln)
|
ln = admission.WrapListener(l.options.Admission, ln)
|
||||||
|
ln = limiter.WrapListener(l.options.RateLimiter, ln)
|
||||||
|
|
||||||
if l.tlsEnabled {
|
if l.tlsEnabled {
|
||||||
ln = tls.NewListener(ln, l.options.TLSConfig)
|
ln = tls.NewListener(ln, l.options.TLSConfig)
|
||||||
|
272
metrics/wrapper/conn.go
Normal file
272
metrics/wrapper/conn.go
Normal file
@ -0,0 +1,272 @@
|
|||||||
|
package wrapper
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"syscall"
|
||||||
|
|
||||||
|
"github.com/go-gost/core/metrics"
|
||||||
|
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
|
||||||
|
service string
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapConn(service string, c net.Conn) net.Conn {
|
||||||
|
if !metrics.IsEnabled() {
|
||||||
|
return c
|
||||||
|
}
|
||||||
|
return &serverConn{
|
||||||
|
service: service,
|
||||||
|
Conn: c,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *serverConn) Read(b []byte) (n int, err error) {
|
||||||
|
n, err = c.Conn.Read(b)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferInputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *serverConn) Write(b []byte) (n int, err error) {
|
||||||
|
n, err = c.Conn.Write(b)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferOutputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
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
|
||||||
|
service string
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapPacketConn(service string, pc net.PacketConn) net.PacketConn {
|
||||||
|
if !metrics.IsEnabled() {
|
||||||
|
return pc
|
||||||
|
}
|
||||||
|
return &packetConn{
|
||||||
|
PacketConn: pc,
|
||||||
|
service: service,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
|
||||||
|
n, addr, err = c.PacketConn.ReadFrom(p)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferInputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
|
||||||
|
n, err = c.PacketConn.WriteTo(p, addr)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferOutputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
type udpConn struct {
|
||||||
|
net.PacketConn
|
||||||
|
service string
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapUDPConn(service string, pc net.PacketConn) udp.Conn {
|
||||||
|
return &udpConn{
|
||||||
|
PacketConn: pc,
|
||||||
|
service: service,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferInputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
err = errUnsupport
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
|
||||||
|
n, addr, err = c.PacketConn.ReadFrom(p)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferInputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
|
||||||
|
if nc, ok := c.PacketConn.(udp.ReadUDP); ok {
|
||||||
|
n, addr, err = nc.ReadFromUDP(b)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferInputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
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 {
|
||||||
|
n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferInputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferOutputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
err = errUnsupport
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
|
||||||
|
n, err = c.PacketConn.WriteTo(p, addr)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferOutputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
|
||||||
|
if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
|
||||||
|
n, err = nc.WriteToUDP(b, addr)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferOutputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
err = errUnsupport
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
|
||||||
|
if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
|
||||||
|
n, oobn, err = nc.WriteMsgUDP(b, oob, addr)
|
||||||
|
if counter := metrics.GetCounter(
|
||||||
|
metrics.MetricServiceTransferOutputBytesCounter,
|
||||||
|
metrics.Labels{
|
||||||
|
"service": c.service,
|
||||||
|
}); counter != nil {
|
||||||
|
counter.Add(float64(n))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
err = errUnsupport
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) {
|
||||||
|
if nc, ok := c.PacketConn.(syscall.Conn); 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
metrics/wrapper/listener.go
Normal file
32
metrics/wrapper/listener.go
Normal file
@ -0,0 +1,32 @@
|
|||||||
|
package wrapper
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net"
|
||||||
|
|
||||||
|
"github.com/go-gost/core/metrics"
|
||||||
|
)
|
||||||
|
|
||||||
|
type listener struct {
|
||||||
|
service string
|
||||||
|
net.Listener
|
||||||
|
}
|
||||||
|
|
||||||
|
func WrapListener(service string, ln net.Listener) net.Listener {
|
||||||
|
if !metrics.IsEnabled() {
|
||||||
|
return ln
|
||||||
|
}
|
||||||
|
|
||||||
|
return &listener{
|
||||||
|
service: service,
|
||||||
|
Listener: ln,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (ln *listener) Accept() (net.Conn, error) {
|
||||||
|
c, err := ln.Listener.Accept()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return WrapConn(ln.service, c), nil
|
||||||
|
}
|
@ -31,18 +31,18 @@ type rlimiterWrapper struct {
|
|||||||
r *rlimiterRegistry
|
r *rlimiterRegistry
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *rlimiterWrapper) Input() limiter.Limiter {
|
func (w *rlimiterWrapper) In(key string) limiter.Limiter {
|
||||||
v := w.r.get(w.name)
|
v := w.r.get(w.name)
|
||||||
if v == nil {
|
if v == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return v.Input()
|
return v.In(key)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *rlimiterWrapper) Output() limiter.Limiter {
|
func (w *rlimiterWrapper) Out(key string) limiter.Limiter {
|
||||||
v := w.r.get(w.name)
|
v := w.r.get(w.name)
|
||||||
if v == nil {
|
if v == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return v.Output()
|
return v.Out(key)
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user