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

update limiter

Browse Source
This commit is contained in:
ginuerzh
2022-09-14 20:00:35 +08:00
parent 91c12882f5
commit 01d7dc77c6
34 changed files with 1171 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
handler/dns/handler.go
View File

@ -130,6 +130,10 @@ func (h *dnsHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
b := bufpool.Get(h.md.bufferSize)
defer bufpool.Put(b)
@ -152,6 +156,18 @@ func (h *dnsHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
return nil
}
func (h *dnsHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}
func (h *dnsHandler) exchange(ctx context.Context, msg []byte, log logger.Logger) ([]byte, error) {
mq := dns.Msg{}
if err := mq.Unpack(msg); err != nil {

16
handler/forward/local/handler.go
View File

@ -77,6 +77,10 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
target := h.group.Next(ctx)
if target == nil {
err := errors.New("target not available")
@ -119,3 +123,15 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return nil
}
func (h *forwardHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/forward/remote/handler.go
View File

@ -71,6 +71,10 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
target := h.group.Next(ctx)
if target == nil {
err := errors.New("target not available")
@ -113,3 +117,15 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return nil
}
func (h *forwardHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/http/handler.go
View File

@ -75,6 +75,10 @@ func (h *httpHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
req, err := http.ReadRequest(bufio.NewReader(conn))
if err != nil {
log.Error(err)
@ -337,3 +341,15 @@ func (h *httpHandler) authenticate(conn net.Conn, req *http.Request, resp *http.
resp.Write(conn)
return
}
func (h *httpHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/http2/handler.go
View File

@ -75,6 +75,10 @@ func (h *http2Handler) Handle(ctx context.Context, conn net.Conn, opts ...handle
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
v, ok := conn.(md.Metadatable)
if !ok || v == nil {
err := errors.New("wrong connection type")
@ -345,3 +349,15 @@ func (h *http2Handler) writeResponse(w http.ResponseWriter, resp *http.Response)
_, err := io.Copy(flushWriter{w}, resp.Body)
return err
}
func (h *http2Handler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/redirect/tcp/handler.go
View File

@ -74,6 +74,10 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
var dstAddr net.Addr
if h.md.tproxy {
@ -269,3 +273,15 @@ func (h *redirectHandler) getServerName(ctx context.Context, r io.Reader) (host
return
}
func (h *redirectHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/redirect/udp/handler.go
View File

@ -63,6 +63,10 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
dstAddr := conn.LocalAddr()
log = log.WithFields(map[string]any{
@ -92,3 +96,15 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han
return nil
}
func (h *redirectHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/relay/handler.go
View File

@ -75,6 +75,10 @@ func (h *relayHandler) Handle(ctx context.Context, conn net.Conn, opts ...handle
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
if h.md.readTimeout > 0 {
conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))
}
@ -145,3 +149,15 @@ func (h *relayHandler) Handle(ctx context.Context, conn net.Conn, opts ...handle
}
return ErrUnknownCmd
}
func (h *relayHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/sni/handler.go
View File

@ -76,6 +76,10 @@ func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
var hdr [dissector.RecordHeaderLen]byte
if _, err := io.ReadFull(conn, hdr[:]); err != nil {
log.Error(err)
@ -251,3 +255,15 @@ func (h *sniHandler) decodeServerName(s string) (string, error) {
}
return string(v), nil
}
func (h *sniHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/socks/v4/handler.go
View File

@ -72,6 +72,10 @@ func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
if h.md.readTimeout > 0 {
conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))
}
@ -150,3 +154,15 @@ func (h *socks4Handler) handleBind(ctx context.Context, conn net.Conn, req *goso
// TODO: bind
return ErrUnimplemented
}
func (h *socks4Handler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/socks/v5/handler.go
View File

@ -78,6 +78,10 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
if h.md.readTimeout > 0 {
conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))
}
@ -113,3 +117,15 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
return err
}
}
func (h *socks5Handler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/ss/handler.go
View File

@ -76,6 +76,10 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
if h.cipher != nil {
conn = ss.ShadowConn(h.cipher.StreamConn(conn), nil)
}
@ -117,3 +121,15 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H
return nil
}
func (h *ssHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/ss/udp/handler.go
View File

@ -77,6 +77,10 @@ func (h *ssuHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
pc, ok := conn.(net.PacketConn)
if ok {
if h.cipher != nil {
@ -186,3 +190,15 @@ func (h *ssuHandler) relayPacket(pc1, pc2 net.PacketConn, log logger.Logger) (er
return <-errc
}
func (h *ssuHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}

16
handler/sshd/handler.go
View File

@ -66,6 +66,10 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
"local": conn.LocalAddr().String(),
})
if !h.checkRateLimit(conn.RemoteAddr()) {
return nil
}
switch cc := conn.(type) {
case *sshd_util.DirectForwardConn:
return h.handleDirectForward(ctx, cc, log)
@ -217,6 +221,18 @@ func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_uti
return nil
}
func (h *forwardHandler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}
func getHostPortFromAddr(addr net.Addr) (host string, port int, err error) {
host, portString, err := net.SplitHostPort(addr.String())
if err != nil {