update limiter
This commit is contained in:
@ -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
|
||||
}
|
||||
|
Reference in New Issue
Block a user