mirror of
https://github.com/ergochat/ergo.git
synced 2024-11-23 12:29:25 +01:00
219 lines
5.4 KiB
Go
219 lines
5.4 KiB
Go
// Copyright (c) 2016-2017 Daniel Oaks <daniel@danieloaks.net>
|
|
// released under the MIT license
|
|
|
|
package connection_limits
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// ThrottlerConfig controls the automated connection throttling.
|
|
type ThrottlerConfig struct {
|
|
Enabled bool
|
|
CidrLenIPv4 int `yaml:"cidr-len-ipv4"`
|
|
CidrLenIPv6 int `yaml:"cidr-len-ipv6"`
|
|
ConnectionsPerCidr int `yaml:"max-connections"`
|
|
DurationString string `yaml:"duration"`
|
|
Duration time.Duration `yaml:"duration-time"`
|
|
BanDurationString string `yaml:"ban-duration"`
|
|
BanDuration time.Duration
|
|
BanMessage string `yaml:"ban-message"`
|
|
Exempted []string
|
|
}
|
|
|
|
// ThrottleDetails holds the connection-throttling details for a subnet/IP.
|
|
type ThrottleDetails struct {
|
|
Start time.Time
|
|
Count int
|
|
}
|
|
|
|
// GenericThrottle allows enforcing limits of the form
|
|
// "at most X events per time window of duration Y"
|
|
type GenericThrottle struct {
|
|
ThrottleDetails // variable state: what events have been seen
|
|
// these are constant after creation:
|
|
Duration time.Duration // window length to consider
|
|
Limit int // number of events allowed per window
|
|
}
|
|
|
|
// Touch checks whether an additional event is allowed:
|
|
// it either denies it (by returning false) or allows it (by returning true)
|
|
// and records it
|
|
func (g *GenericThrottle) Touch() (throttled bool, remainingTime time.Duration) {
|
|
return g.touch(time.Now())
|
|
}
|
|
|
|
func (g *GenericThrottle) touch(now time.Time) (throttled bool, remainingTime time.Duration) {
|
|
if g.Limit == 0 {
|
|
return // limit of 0 disables throttling
|
|
}
|
|
|
|
elapsed := now.Sub(g.Start)
|
|
if elapsed > g.Duration {
|
|
// reset window, record the operation
|
|
g.Start = now
|
|
g.Count = 1
|
|
return false, 0
|
|
} else if g.Count >= g.Limit {
|
|
// we are throttled
|
|
return true, g.Start.Add(g.Duration).Sub(now)
|
|
} else {
|
|
// we are not throttled, record the operation
|
|
g.Count += 1
|
|
return false, 0
|
|
}
|
|
}
|
|
|
|
// Throttler manages automated client connection throttling.
|
|
type Throttler struct {
|
|
sync.RWMutex
|
|
|
|
enabled bool
|
|
ipv4Mask net.IPMask
|
|
ipv6Mask net.IPMask
|
|
subnetLimit int
|
|
duration time.Duration
|
|
population map[string]ThrottleDetails
|
|
|
|
// used by the server to ban clients that go over this limit
|
|
banDuration time.Duration
|
|
banMessage string
|
|
|
|
// exemptedIPs holds IPs that are exempt from limits
|
|
exemptedIPs map[string]bool
|
|
// exemptedNets holds networks that are exempt from limits
|
|
exemptedNets []net.IPNet
|
|
}
|
|
|
|
// maskAddr masks the given IPv4/6 address with our cidr limit masks.
|
|
func (ct *Throttler) maskAddr(addr net.IP) net.IP {
|
|
if addr.To4() == nil {
|
|
// IPv6 addr
|
|
addr = addr.Mask(ct.ipv6Mask)
|
|
} else {
|
|
// IPv4 addr
|
|
addr = addr.Mask(ct.ipv4Mask)
|
|
}
|
|
|
|
return addr
|
|
}
|
|
|
|
// ResetFor removes any existing count for the given address.
|
|
func (ct *Throttler) ResetFor(addr net.IP) {
|
|
ct.Lock()
|
|
defer ct.Unlock()
|
|
|
|
if !ct.enabled {
|
|
return
|
|
}
|
|
|
|
// remove
|
|
ct.maskAddr(addr)
|
|
addrString := addr.String()
|
|
delete(ct.population, addrString)
|
|
}
|
|
|
|
// AddClient introduces a new client connection if possible. If we can't, throws an error instead.
|
|
func (ct *Throttler) AddClient(addr net.IP) error {
|
|
ct.Lock()
|
|
defer ct.Unlock()
|
|
|
|
if !ct.enabled {
|
|
return nil
|
|
}
|
|
|
|
// check exempted lists
|
|
if ct.exemptedIPs[addr.String()] {
|
|
return nil
|
|
}
|
|
for _, ex := range ct.exemptedNets {
|
|
if ex.Contains(addr) {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// check throttle
|
|
ct.maskAddr(addr)
|
|
addrString := addr.String()
|
|
|
|
details := ct.population[addrString] // retrieve mutable throttle state from the map
|
|
// add in constant state to process the limiting operation
|
|
g := GenericThrottle{
|
|
ThrottleDetails: details,
|
|
Duration: ct.duration,
|
|
Limit: ct.subnetLimit,
|
|
}
|
|
throttled, _ := g.Touch() // actually check the limit
|
|
ct.population[addrString] = g.ThrottleDetails // store modified mutable state
|
|
|
|
if throttled {
|
|
return errTooManyClients
|
|
} else {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (ct *Throttler) BanDuration() time.Duration {
|
|
ct.RLock()
|
|
defer ct.RUnlock()
|
|
|
|
return ct.banDuration
|
|
}
|
|
|
|
func (ct *Throttler) BanMessage() string {
|
|
ct.RLock()
|
|
defer ct.RUnlock()
|
|
|
|
return ct.banMessage
|
|
}
|
|
|
|
// NewThrottler returns a new client connection throttler.
|
|
// The throttler is functional, but disabled; it can be enabled via `ApplyConfig`.
|
|
func NewThrottler() *Throttler {
|
|
var ct Throttler
|
|
|
|
// initialize empty population; all other state is configurable
|
|
ct.population = make(map[string]ThrottleDetails)
|
|
|
|
return &ct
|
|
}
|
|
|
|
// ApplyConfig atomically applies a config update to a throttler
|
|
func (ct *Throttler) ApplyConfig(config ThrottlerConfig) error {
|
|
// assemble exempted nets
|
|
exemptedIPs := make(map[string]bool)
|
|
var exemptedNets []net.IPNet
|
|
for _, cidr := range config.Exempted {
|
|
ipaddr := net.ParseIP(cidr)
|
|
_, netaddr, err := net.ParseCIDR(cidr)
|
|
|
|
if ipaddr == nil && err != nil {
|
|
return fmt.Errorf("Could not parse exempted IP/network [%s]", cidr)
|
|
}
|
|
|
|
if ipaddr != nil {
|
|
exemptedIPs[ipaddr.String()] = true
|
|
} else {
|
|
exemptedNets = append(exemptedNets, *netaddr)
|
|
}
|
|
}
|
|
|
|
ct.Lock()
|
|
defer ct.Unlock()
|
|
|
|
ct.enabled = config.Enabled
|
|
ct.ipv4Mask = net.CIDRMask(config.CidrLenIPv4, 32)
|
|
ct.ipv6Mask = net.CIDRMask(config.CidrLenIPv6, 128)
|
|
ct.subnetLimit = config.ConnectionsPerCidr
|
|
ct.duration = config.Duration
|
|
ct.banDuration = config.BanDuration
|
|
ct.banMessage = config.BanMessage
|
|
ct.exemptedIPs = exemptedIPs
|
|
ct.exemptedNets = exemptedNets
|
|
|
|
return nil
|
|
}
|