@ -2,16 +2,23 @@ package connlimit
import (
"errors"
"fmt"
"net"
"net/http"
"sync"
"sync/atomic"
"time"
)
var (
// ErrPerClientIPLimitReached is returned if accepting a new conn would exceed
// the per-client-ip limit set.
ErrPerClientIPLimitReached = errors . New ( "client connection limit reached" )
tooManyConnsMsg = "Your IP is issuing too many concurrent connections, please rate limit your calls\n"
tooManyRequestsResponse = [ ] byte ( fmt . Sprintf ( "HTTP/1.1 429 Too Many Requests\r\n" +
"Content-Type: text/plain\r\n" +
"Content-Length: %d\r\n" +
"Connection: close\r\n\r\n%s" , len ( tooManyConnsMsg ) , tooManyConnsMsg ) )
)
// Limiter implements a simple limiter that tracks the number of connections
@ -173,7 +180,7 @@ func (l *Limiter) SetConfig(c Config) {
l . cfg . Store ( c )
}
// HTTPConnStateFunc returns a func that can be passed as the ConnState field of
// HTTPConnStateFunc WithErrorHandler returns a func that can be passed as the ConnState field of
// an http.Server. This intercepts new HTTP connections to the server and
// applies the limiting to new connections.
//
@ -181,13 +188,15 @@ func (l *Limiter) SetConfig(c Config) {
// in the limiter as if it was closed. Servers that use Hijacking must implement
// their own calls if they need to continue limiting the number of concurrent
// hijacked connections.
func ( l * Limiter ) HTTPConnStateFunc ( ) func ( net . Conn , http . ConnState ) {
// errorHandler MUST close the connection itself
func ( l * Limiter ) HTTPConnStateFuncWithErrorHandler ( errorHandler func ( error , net . Conn ) ) func ( net . Conn , http . ConnState ) {
return func ( conn net . Conn , state http . ConnState ) {
switch state {
case http . StateNew :
_ , err := l . Accept ( conn )
if err != nil {
conn . Close ( )
errorHandler ( err , conn )
}
case http . StateHijacked :
l . freeConn ( conn )
@ -199,3 +208,26 @@ func (l *Limiter) HTTPConnStateFunc() func(net.Conn, http.ConnState) {
}
}
}
// HTTPConnStateFunc is here for ascending compatibility reasons.
func ( l * Limiter ) HTTPConnStateFunc ( ) func ( net . Conn , http . ConnState ) {
return l . HTTPConnStateFuncWithErrorHandler ( func ( err error , conn net . Conn ) {
conn . Close ( )
} )
}
// HTTPConnStateFuncWithDefault429Handler return an HTTP 429 if too many connections occur.
// BEWARE that returning HTTP 429 is done on critical path, you might choose to use
// HTTPConnStateFuncWithErrorHandler if you want to use a non-blocking strategy.
func ( l * Limiter ) HTTPConnStateFuncWithDefault429Handler ( writeDeadlineMaxDelay time . Duration ) func ( net . Conn , http . ConnState ) {
return l . HTTPConnStateFuncWithErrorHandler ( func ( err error , conn net . Conn ) {
if err == ErrPerClientIPLimitReached {
// We don't care about slow players
if writeDeadlineMaxDelay > 0 {
conn . SetDeadline ( time . Now ( ) . Add ( writeDeadlineMaxDelay ) )
}
conn . Write ( tooManyRequestsResponse )
}
conn . Close ( )
} )
}
Pierre Souchay
4 years ago
committed by