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

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