114 lines
3.0 KiB
Go
114 lines
3.0 KiB
Go
package services
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// The rate limit, which is a cost control and not a security one.
|
|
|
|
func TestABurstOfQuestionsGetsThrough(t *testing.T) {
|
|
// Somebody working through the prompt chips on a page fires four in a row
|
|
// and must not be stopped mid-thought.
|
|
at := time.Now()
|
|
l := newAskLimiter(func() time.Time { return at })
|
|
|
|
for i := range askBurst {
|
|
if err := l.allow(904); err != nil {
|
|
t.Fatalf("question %d of a burst was refused: %v", i+1, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTheSeventhQuestionInARowWaits(t *testing.T) {
|
|
at := time.Now()
|
|
l := newAskLimiter(func() time.Time { return at })
|
|
for range askBurst {
|
|
_ = l.allow(904)
|
|
}
|
|
|
|
err := l.allow(904)
|
|
var tooFast ErrTooFast
|
|
if !errors.As(err, &tooFast) {
|
|
t.Fatalf("an unbounded burst was allowed: %v", err)
|
|
}
|
|
if tooFast.RetryIn <= 0 {
|
|
t.Fatal("the refusal does not say how long to wait")
|
|
}
|
|
// The message is read by a shopkeeper, so it has to say something useful.
|
|
if !contains(err.Error(), "seconds") {
|
|
t.Fatalf("unhelpful message: %q", err)
|
|
}
|
|
}
|
|
|
|
func TestTheAllowanceComesBackWithTime(t *testing.T) {
|
|
// Continuous refill, not a fixed window: a person is never told to wait out
|
|
// a window that started before they arrived.
|
|
at := time.Now()
|
|
l := newAskLimiter(func() time.Time { return at })
|
|
for range askBurst {
|
|
_ = l.allow(904)
|
|
}
|
|
if err := l.allow(904); err == nil {
|
|
t.Fatal("expected to be out of allowance")
|
|
}
|
|
|
|
at = at.Add(askRefill)
|
|
if err := l.allow(904); err != nil {
|
|
t.Fatalf("one refill period bought nothing: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestOnePersonCannotLockOutTheirColleagues(t *testing.T) {
|
|
// Per user rather than per tenant, so an impatient person in a shop cannot
|
|
// turn a cost control into an outage for everybody else.
|
|
at := time.Now()
|
|
l := newAskLimiter(func() time.Time { return at })
|
|
for range askBurst + 3 {
|
|
_ = l.allow(904)
|
|
}
|
|
|
|
if err := l.allow(905); err != nil {
|
|
t.Fatalf("a colleague was refused because somebody else was busy: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestAnUnidentifiedCallerIsNotSilentlyThrottled(t *testing.T) {
|
|
// Every such request is already refused for having no session. Rate
|
|
// limiting it too would hide the real reason behind a confusing one.
|
|
l := newAskLimiter(nil)
|
|
for range askBurst * 3 {
|
|
if err := l.allow(0); err != nil {
|
|
t.Fatalf("a caller with no user id was throttled: %v", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestQuietCallersAreForgotten(t *testing.T) {
|
|
// So the map does not grow with every account that ever asked anything.
|
|
at := time.Now()
|
|
l := newAskLimiter(func() time.Time { return at })
|
|
|
|
for user := 1; user <= 300; user++ {
|
|
_ = l.allow(user)
|
|
}
|
|
before := len(l.buckets)
|
|
|
|
at = at.Add(askIdle + time.Minute)
|
|
_ = l.allow(99999) // a new caller is what triggers the sweep
|
|
|
|
if len(l.buckets) >= before {
|
|
t.Fatalf("nothing was swept: %d then %d", before, len(l.buckets))
|
|
}
|
|
}
|
|
|
|
func contains(haystack, needle string) bool {
|
|
for i := 0; i+len(needle) <= len(haystack); i++ {
|
|
if haystack[i:i+len(needle)] == needle {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|