fix masking bug

IP.Mask() returns a new IP value, rather than modifying its target in place
This commit is contained in:
Shivaram Lingamneni 2019-02-05 01:44:58 -05:00
parent 1c23af8767
commit eb8f0e50df
2 changed files with 10 additions and 29 deletions

@ -42,17 +42,14 @@ type Limiter struct {
exemptedNets []net.IPNet
}
// maskAddr masks the given IPv4/6 address with our cidr limit masks.
func (cl *Limiter) maskAddr(addr net.IP) net.IP {
if addr.To4() == nil {
// IPv6 addr
addr = addr.Mask(cl.ipv6Mask)
// addrToKey canonicalizes `addr` to a string key.
func addrToKey(addr net.IP, v4Mask net.IPMask, v6Mask net.IPMask) string {
if addr.To4() != nil {
addr = addr.Mask(v4Mask) // IP.Mask() handles the 4-in-6 mapping for us
} else {
// IPv4 addr
addr = addr.Mask(cl.ipv4Mask)
addr = addr.Mask(v6Mask)
}
return addr
return addr.String()
}
// AddClient adds a client to our population if possible. If we can't, throws an error instead.
@ -72,8 +69,7 @@ func (cl *Limiter) AddClient(addr net.IP, force bool) error {
}
// check population
cl.maskAddr(addr)
addrString := addr.String()
addrString := addrToKey(addr, cl.ipv4Mask, cl.ipv6Mask)
if cl.population[addrString]+1 > cl.subnetLimit && !force {
return errTooManyClients
@ -93,7 +89,7 @@ func (cl *Limiter) RemoveClient(addr net.IP) {
return
}
addrString := addr.String()
addrString := addrToKey(addr, cl.ipv4Mask, cl.ipv6Mask)
cl.population[addrString] = cl.population[addrString] - 1
// safety limiter

@ -88,19 +88,6 @@ type Throttler struct {
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()
@ -111,8 +98,7 @@ func (ct *Throttler) ResetFor(addr net.IP) {
}
// remove
ct.maskAddr(addr)
addrString := addr.String()
addrString := addrToKey(addr, ct.ipv4Mask, ct.ipv6Mask)
delete(ct.population, addrString)
}
@ -131,8 +117,7 @@ func (ct *Throttler) AddClient(addr net.IP) error {
}
// check throttle
ct.maskAddr(addr)
addrString := addr.String()
addrString := addrToKey(addr, ct.ipv4Mask, ct.ipv6Mask)
details := ct.population[addrString] // retrieve mutable throttle state from the map
// add in constant state to process the limiting operation