forked from ergochat/ergo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
connection_limits.go
111 lines (91 loc) · 2.9 KB
/
connection_limits.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
// Copyright (c) 2016- Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license
package irc
import (
"errors"
"fmt"
"net"
)
var (
errTooManyClients = errors.New("Too many clients in subnet")
)
// ConnectionLimits manages the automated client connection limits.
type ConnectionLimits struct {
ipv4Mask net.IPMask
ipv6Mask net.IPMask
// subnetLimit is the maximum number of clients per subnet
subnetLimit int
// population holds IP -> count of clients connected from there
population map[string]int
// 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 (cl *ConnectionLimits) maskAddr(addr net.IP) net.IP {
if addr.To4() == nil {
// IPv6 addr
addr = addr.Mask(cl.ipv6Mask)
} else {
// IPv4 addr
addr = addr.Mask(cl.ipv4Mask)
}
return addr
}
// AddClient adds a client to our population if possible. If we can't, throws an error instead.
// 'force' is used to add already-existing clients (i.e. ones that are already on the network).
func (cl *ConnectionLimits) AddClient(addr net.IP, force bool) error {
// check exempted lists
// we don't track populations for exempted addresses or nets - this is by design
if cl.exemptedIPs[addr.String()] {
return nil
}
for _, ex := range cl.exemptedNets {
if ex.Contains(addr) {
return nil
}
}
// check population
cl.maskAddr(addr)
addrString := addr.String()
if cl.population[addrString]+1 > cl.subnetLimit && !force {
return errTooManyClients
}
cl.population[addrString] = cl.population[addrString] + 1
return nil
}
// RemoveClient removes the given address from our population
func (cl *ConnectionLimits) RemoveClient(addr net.IP) {
addrString := addr.String()
cl.population[addrString] = cl.population[addrString] - 1
// safety limiter
if cl.population[addrString] < 0 {
cl.population[addrString] = 0
}
}
// NewConnectionLimits returns a new connection limit handler.
func NewConnectionLimits(config ConnectionLimitsConfig) (*ConnectionLimits, error) {
var cl ConnectionLimits
cl.population = make(map[string]int)
cl.exemptedIPs = make(map[string]bool)
cl.ipv4Mask = net.CIDRMask(config.CidrLenIPv4, 32)
cl.ipv6Mask = net.CIDRMask(config.CidrLenIPv6, 128)
// subnetLimit is explicitly NOT capped at a minimum of one.
// this is so that CL config can be used to allow ONLY clients from exempted IPs/nets
cl.subnetLimit = config.IPsPerCidr
// assemble exempted nets
for _, cidr := range config.Exempted {
ipaddr := net.ParseIP(cidr)
_, netaddr, err := net.ParseCIDR(cidr)
if ipaddr == nil && err != nil {
return nil, fmt.Errorf("Could not parse exempted IP/network [%s]", cidr)
}
if ipaddr != nil {
cl.exemptedIPs[ipaddr.String()] = true
} else {
cl.exemptedNets = append(cl.exemptedNets, *netaddr)
}
}
return &cl, nil
}