The Jailer in My Codebase: In-Memory Jail and Why It Became a Library

  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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178

package main

import (
	"sync"
	"time"
)

// Jailer is the only interface any caller talks to.
// The rules (what counts as a try, how long a principal stays
// jailed) live in Policy, not here, so a test swaps the rules
// without touching the store.
type Jailer interface {
	// Try records one try. Returns the principal's current
	// in-window count and a jailed flag. A jailed principal
	// reports jailed on every call until the jail expires.
	Try(principal string) (count int, jailed bool)
	// Clear releases a principal (e.g. after a successful
	// login): the counter resets and any active jail lifts.
	Clear(principal string)
	// IsJailed reports whether the principal is jailed right now.
	IsJailed(principal string) bool
}

// Policy is the rulebook: how many tries per window, how big
// the window is, how long the jail lasts, and how it escalates.
// It is a value type: a test builds a custom one in one literal.
type Policy struct {
	MaxTries  int
	Window    time.Duration
	BaseDelay time.Duration
	MaxDelay  time.Duration
	Now       func() time.Time // injectable clock; tests rewind it
}

func DefaultPolicy() Policy {
	return Policy{
		MaxTries:  5,
		Window:    time.Minute,
		BaseDelay: 30 * time.Second,
		MaxDelay:  time.Hour,
		Now:       time.Now,
	}
}

// entry is per-principal state: one counter, one window start,
// one jail timestamp, one escalation streak.
type entry struct {
	count       int
	windowStart time.Time
	jailUntil   time.Time
	streak      int
}

// MemoryJailer is the in-memory store.
// ponytail: single-process only — state is lost on restart, and a
// multi-replica deploy jails each replica independently. Swap in a
// Redis impl behind the same interface when you need a shared jail;
// the interface does not change.
type MemoryJailer struct {
	mu      sync.Mutex
	p       Policy
	entries map[string]*entry
	stop    chan struct{}
}

// New builds a jailer and starts the background sweeper.
func New(p Policy) *MemoryJailer {
	if p.Now == nil {
		p.Now = time.Now
	}
	j := &MemoryJailer{
		p:       p,
		entries: make(map[string]*entry),
		stop:    make(chan struct{}),
	}
	go j.sweep()
	return j
}

// Stop shuts the sweeper down. Call it on process shutdown;
// without it the sweeper leaks a goroutine per jailer.
func (j *MemoryJailer) Stop() { close(j.stop) }

// Try implements Jailer. One try = one locked, atomic increment —
// never a read-modify-write across two calls.
func (j *MemoryJailer) Try(principal string) (int, bool) {
	j.mu.Lock()
	defer j.mu.Unlock()
	now := j.p.Now()
	e := j.entry(principal, now)

	// A jailed principal *stays* jailed: the jail is state,
	// not a return value of this call. A failed try does not
	// reset the counter.
	if now.Before(e.jailUntil) {
		return e.count, true
	}

	e.count++
	if e.count > j.p.MaxTries {
		e.streak++
		e.jailUntil = now.Add(j.delay(e.streak))
		e.count = 0 // next window starts clean
		return e.count, true
	}
	return e.count, false
}

// delay is the exponential backoff: base * 2^(streak-1), capped.
// The escalation is a policy — a function of the streak — not a
// property of the principal.
func (j *MemoryJailer) delay(streak int) time.Duration {
	d := j.p.BaseDelay
	for i := 1; i < streak && d < j.p.MaxDelay; i++ {
		d *= 2
	}
	if d > j.p.MaxDelay {
		d = j.p.MaxDelay
	}
	return d
}

// entry returns the principal's entry, resetting the counter when
// the window has elapsed.
func (j *MemoryJailer) entry(principal string, now time.Time) *entry {
	e, ok := j.entries[principal]
	if !ok {
		e = &entry{windowStart: now}
		j.entries[principal] = e
		return e
	}
	// Tumbling window: when it elapses, the counter restarts.
	// ponytail: tumbling, not sliding — a sliding window needs a
	// timestamp per try, which buys nothing for login attempts.
	if now.Sub(e.windowStart) > j.p.Window {
		e.count = 0
		e.windowStart = now
	}
	return e
}

// Clear implements Jailer.
func (j *MemoryJailer) Clear(principal string) {
	j.mu.Lock()
	defer j.mu.Unlock()
	delete(j.entries, principal)
}

// IsJailed implements Jailer.
func (j *MemoryJailer) IsJailed(principal string) bool {
	j.mu.Lock()
	defer j.mu.Unlock()
	e, ok := j.entries[principal]
	return ok && j.p.Now().Before(e.jailUntil)
}

// sweep evicts entries whose window and jail have both expired.
// Without it the map grows one entry per principal that ever tried
// (and failed) and never came back.
func (j *MemoryJailer) sweep() {
	t := time.NewTicker(j.p.Window)
	defer t.Stop()
	for {
		select {
		case <-j.stop:
			return
		case <-t.C:
			j.mu.Lock()
			now := j.p.Now()
			for p, e := range j.entries {
				if now.Sub(e.windowStart) > j.p.Window && now.After(e.jailUntil) {
					delete(j.entries, p)
				}
			}
			j.mu.Unlock()
		}
	}
}

 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

package main

import (
	"testing"
	"time"
)

// testPolicy builds a jail policy against a fake clock the tests
// move by hand: 3 tries per 1ms window, 10ms base jail, 100ms cap.
func testPolicy(now *time.Time) Policy {
	return Policy{
		MaxTries:  3,
		Window:    1 * time.Millisecond,
		BaseDelay: 10 * time.Millisecond,
		MaxDelay:  100 * time.Millisecond,
		Now:       func() time.Time { return *now },
	}
}

// The 4th try (after 3) jails the principal, and the jail sticks
// until the delay elapses.
func TestJailAfterMaxTries(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	j := New(testPolicy(&now))
	defer j.Stop()

	for i := 0; i < 3; i++ {
		if _, jailed := j.Try("alice"); jailed {
			t.Fatalf("try %d: jailed before exceeding MaxTries", i+1)
		}
	}
	if _, jailed := j.Try("alice"); !jailed {
		t.Fatal("4th try: expected jailed, got clean")
	}
	if !j.IsJailed("alice") {
		t.Fatal("IsJailed: expected true immediately after jail")
	}

	now = now.Add(10 * time.Millisecond) // exactly one BaseDelay
	if j.IsJailed("alice") {
		t.Fatal("IsJailed: expected free after BaseDelay")
	}
}

// The second jail runs for 2x the first (exponential backoff),
// and the streak cap holds at MaxDelay.
func TestExponentialBackoff(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	j := New(testPolicy(&now))
	defer j.Stop()

	for i := 0; i < 3; i++ {
		j.Try("bob")
	}
	j.Try("bob") // first jail: 10ms
	if !j.IsJailed("bob") {
		t.Fatal("first jail: expected jailed")
	}
	now = now.Add(10 * time.Millisecond) // first jail expires

	for i := 0; i < 3; i++ {
		j.Try("bob")
	}
	j.Try("bob") // second jail: 20ms
	now = now.Add(10 * time.Millisecond)
	if !j.IsJailed("bob") {
		t.Fatal("second jail: still jailed at 1x base")
	}
	now = now.Add(10 * time.Millisecond) // 2x base elapsed
	if j.IsJailed("bob") {
		t.Fatal("second jail: free at 2x base")
	}
}

// Clear is a hard release: counter and jail both die.
func TestClear(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	j := New(testPolicy(&now))
	defer j.Stop()

	for i := 0; i < 4; i++ {
		j.Try("carol")
	}
	if !j.IsJailed("carol") {
		t.Fatal("expected jailed after 4 tries")
	}
	j.Clear("carol")
	if j.IsJailed("carol") {
		t.Fatal("Clear: expected not jailed")
	}
}