39 lines
920 B
Go
39 lines
920 B
Go
|
package capnp
|
||
|
|
|
||
|
|
import "sync/atomic"
|
||
|
|
|
||
|
|
// A ReadLimiter tracks the number of bytes read from a message in order
|
||
|
|
// to avoid amplification attacks as detailed in
|
||
|
|
// https://capnproto.org/encoding.html#amplification-attack.
|
||
|
|
// It is safe to use from multiple goroutines.
|
||
|
|
type ReadLimiter struct {
|
||
|
|
limit uint64
|
||
|
|
}
|
||
|
|
|
||
|
|
// canRead reports whether the amount of bytes can be stored safely.
|
||
|
|
func (rl *ReadLimiter) canRead(sz Size) bool {
|
||
|
|
for {
|
||
|
|
curr := atomic.LoadUint64(&rl.limit)
|
||
|
|
ok := curr >= uint64(sz)
|
||
|
|
var new uint64
|
||
|
|
if ok {
|
||
|
|
new = curr - uint64(sz)
|
||
|
|
} else {
|
||
|
|
new = 0
|
||
|
|
}
|
||
|
|
if atomic.CompareAndSwapUint64(&rl.limit, curr, new) {
|
||
|
|
return ok
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Reset sets the number of bytes allowed to be read.
|
||
|
|
func (rl *ReadLimiter) Reset(limit uint64) {
|
||
|
|
atomic.StoreUint64(&rl.limit, limit)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Unread increases the limit by sz.
|
||
|
|
func (rl *ReadLimiter) Unread(sz Size) {
|
||
|
|
atomic.AddUint64(&rl.limit, uint64(sz))
|
||
|
|
}
|