For reading and verification only. You may not copy, use, modify or redistribute this code. Full notice on the main page.
tunnel/limits.go
| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "hash/maphash" |
| 5 | "net/netip" |
| 6 | "sync" |
| 7 | "time" |
| 8 | ) |
| 9 | |
| 10 | // Limits of the public listener (the design plan, section 5): who may hold how many connections, how fast new ones may arrive, and who is dropped |
| 11 | // for trying names that are not a device. Everything is kept only in memory, for the life of the process: a source's address is a map key here and is never |
| 12 | // part of a log line, a metric label or a file (a log line that must tell two sources apart uses Server.short, the per-boot salted short hash). |
| 13 | // |
| 14 | // A source is an IPv4 address, or the /64 (policy node_limits.ipv6_prefix_bits) of an IPv6 one, so that rotating inside a prefix an attacker owns does not |
| 15 | // make new sources. An address of the OpenAI prefix list gets limits openai_per_ip_factor times larger (one egress address serves many users), is never |
| 16 | // banned, and is still bounded: the list is a quota class, not an allowlist (anyone can make ChatGPT call a name through those addresses). |
| 17 | // |
| 18 | // A device is a source too, and a busy one: for every client stream it serves it opens one data connection to the node, from one address. While an address |
| 19 | // holds a device's authenticated control connection its limits are device_per_ip_factor times larger, and the connections that prove who they are (the |
| 20 | // control connection once it authenticates, a data connection once it attaches) give back the token they took from the bucket: a device is not |
| 21 | // charged for serving its clients, only for what it tries before it has proven itself. |
| 22 | |
| 23 | // Why a connection was refused before any byte of it was read, as the metrics name it. |
| 24 | const ( |
| 25 | refuseBanned = "ip_banned" |
| 26 | refuseRate = "ip_rate" |
| 27 | refuseConcurrent = "ip_concurrent" |
| 28 | ) |
| 29 | |
| 30 | // ipEntry is what the node remembers of one source. |
| 31 | type ipEntry struct { |
| 32 | conns int // connections it holds now (pre-hello, and its client streams) |
| 33 | tokens float64 // new-connection bucket |
| 34 | refilled time.Time // when tokens were last computed |
| 35 | banUntil time.Time |
| 36 | strikes []uint64 // the distinct names (or junk hellos) tried in the current window, as per-boot hashes: never the names |
| 37 | strikeAt time.Time // when the current window began |
| 38 | junk uint64 |
| 39 | } |
| 40 | |
| 41 | // ipLimits is the table of sources. It is safe for concurrent use. |
| 42 | type ipLimits struct { |
| 43 | clk clock |
| 44 | |
| 45 | concurrent int |
| 46 | rate float64 // new connections per second, and the burst |
| 47 | factor int // the multiple of the limits an address of the OpenAI list has |
| 48 | devFactor int // and an address that holds an authenticated control connection |
| 49 | v6bits int |
| 50 | max int |
| 51 | strikeMax int |
| 52 | window time.Duration |
| 53 | ban time.Duration |
| 54 | |
| 55 | seed maphash.Seed |
| 56 | |
| 57 | mu sync.Mutex |
| 58 | entries map[netip.Addr]*ipEntry |
| 59 | overflow ipEntry // the shared entry of every source that does not fit in the table |
| 60 | devices map[netip.Addr]int // the sources that hold an authenticated control connection, and how many |
| 61 | past map[netip.Addr]time.Time // the sources that held one and let go of it, and when: a device that reconnects is not an anonymous source (pastTTL) |
| 62 | pastTTL time.Duration |
| 63 | lastSweep time.Time |
| 64 | |
| 65 | overflowed, bans uint64 // counters for the metrics |
| 66 | } |
| 67 | |
| 68 | func newIPLimits(cfg *Config, clk clock, seed maphash.Seed) *ipLimits { |
| 69 | return &ipLimits{ |
| 70 | clk: clk, concurrent: cfg.PerIPConcurrent, rate: float64(cfg.PerIPNewPerS), factor: cfg.OpenAIPerIPFactor, devFactor: cfg.DevicePerIPFactor, v6bits: cfg.IPv6PrefixBits, |
| 71 | max: cfg.IPTableMax, strikeMax: cfg.UnknownSNIPerIP, window: cfg.UnknownSNIWindow, ban: cfg.UnknownSNIBan, seed: seed, |
| 72 | entries: map[netip.Addr]*ipEntry{}, devices: map[netip.Addr]int{}, past: map[netip.Addr]time.Time{}, pastTTL: cfg.HintTTL, |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | // keyOf is the identity limits are kept under. |
| 77 | func (l *ipLimits) keyOf(a netip.Addr) netip.Addr { |
| 78 | a = a.WithZone("").Unmap() |
| 79 | if a.Is6() { |
| 80 | if p, err := a.Prefix(l.v6bits); err == nil { |
| 81 | return p.Addr() |
| 82 | } |
| 83 | } |
| 84 | return a |
| 85 | } |
| 86 | |
| 87 | // entryLocked finds or makes the entry of key. A full table is swept of the idle ones first (at most once a second); when it is still full the source |
| 88 | // shares the overflow entry, which has the limits of a crowd (the table's size divided by 1024 times a source's), never none. |
| 89 | func (l *ipLimits) entryLocked(key netip.Addr, now time.Time) (*ipEntry, int) { |
| 90 | if e := l.entries[key]; e != nil { |
| 91 | return e, 1 |
| 92 | } |
| 93 | if len(l.entries) >= l.max && now.Sub(l.lastSweep) >= time.Second { |
| 94 | l.sweepLocked(now) |
| 95 | } |
| 96 | if len(l.entries) >= l.max { |
| 97 | l.overflowed++ |
| 98 | return &l.overflow, max(1, l.max/1024) |
| 99 | } |
| 100 | e := &ipEntry{tokens: l.rate, refilled: now} |
| 101 | l.entries[key] = e |
| 102 | return e, 1 |
| 103 | } |
| 104 | |
| 105 | // sweepLocked forgets the sources that hold nothing, are not banned, have a full bucket and no strikes in the window. |
| 106 | func (l *ipLimits) sweepLocked(now time.Time) { |
| 107 | l.lastSweep = now |
| 108 | for key, e := range l.entries { |
| 109 | if e.conns == 0 && !now.Before(e.banUntil) && now.Sub(e.refilled) >= time.Second && now.Sub(e.strikeAt) >= l.window { |
| 110 | delete(l.entries, key) |
| 111 | } |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | // Sweep is the periodic cleanup (a background task calls it), so an idle table shrinks without a new source having to arrive. |
| 116 | func (l *ipLimits) Sweep() { |
| 117 | l.mu.Lock() |
| 118 | defer l.mu.Unlock() |
| 119 | now := l.clk.Now() |
| 120 | l.sweepLocked(now) |
| 121 | for key, at := range l.past { |
| 122 | if now.Sub(at) >= l.pastTTL { |
| 123 | delete(l.past, key) |
| 124 | } |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | // multLocked is the multiple of a plain source's limits that key has: the crowd's (scale) for the overflow entry, times the larger of the OpenAI list's |
| 129 | // and a device's. |
| 130 | func (l *ipLimits) multLocked(key netip.Addr, scale int, openai bool) int { |
| 131 | mult := scale |
| 132 | switch { |
| 133 | case openai: |
| 134 | mult *= l.factor |
| 135 | case l.devices[key] > 0: |
| 136 | mult *= l.devFactor |
| 137 | } |
| 138 | return mult |
| 139 | } |
| 140 | |
| 141 | // markDevice records that a device holds an authenticated control connection from addr, until the function it returns is called (once). |
| 142 | func (l *ipLimits) markDevice(addr netip.Addr) (unmark func()) { |
| 143 | key := l.keyOf(addr) |
| 144 | l.mu.Lock() |
| 145 | if l.devices[key]++; l.devices[key] == 1 { |
| 146 | // The bucket of this address was made when the control connection was admitted, with a plain source's burst: it now holds the device's. |
| 147 | e, scale := l.entries[key], 1 |
| 148 | if e == nil { |
| 149 | e, scale = &l.overflow, max(1, l.max/1024) |
| 150 | } |
| 151 | e.tokens = max(e.tokens, l.rate*float64(l.multLocked(key, scale, false))) |
| 152 | } |
| 153 | l.mu.Unlock() |
| 154 | var once sync.Once |
| 155 | return func() { |
| 156 | once.Do(func() { |
| 157 | l.mu.Lock() |
| 158 | defer l.mu.Unlock() |
| 159 | if l.devices[key]--; l.devices[key] <= 0 { |
| 160 | delete(l.devices, key) |
| 161 | if len(l.past) < l.max { // a table that is full forgets the newest, never grows: the address is simply not remembered as a device's |
| 162 | l.past[key] = l.clk.Now() |
| 163 | } |
| 164 | } |
| 165 | }) |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | // knownDevice reports whether addr holds an authenticated control connection now or let go of one less than hint_ttl_s ago: a device that reconnects after its |
| 170 | // connection dropped is not an anonymous source (it is unauthenticated until its PROOF, and it is what an attack on the unauthenticated budget must not lock out). |
| 171 | func (l *ipLimits) knownDevice(addr netip.Addr) bool { |
| 172 | key := l.keyOf(addr) |
| 173 | l.mu.Lock() |
| 174 | defer l.mu.Unlock() |
| 175 | if l.devices[key] > 0 { |
| 176 | return true |
| 177 | } |
| 178 | at, ok := l.past[key] |
| 179 | return ok && l.clk.Now().Sub(at) < l.pastTTL |
| 180 | } |
| 181 | |
| 182 | // isDevice reports whether addr holds an authenticated control connection now. |
| 183 | func (l *ipLimits) isDevice(addr netip.Addr) bool { |
| 184 | key := l.keyOf(addr) |
| 185 | l.mu.Lock() |
| 186 | defer l.mu.Unlock() |
| 187 | return l.devices[key] > 0 |
| 188 | } |
| 189 | |
| 190 | // refund gives back the new-connection token that admit took for a connection that has since proved who it is (a device's control connection after its |
| 191 | // proof, a data connection after its ATTACH): those are not what the per-source rate is there to bound. |
| 192 | func (l *ipLimits) refund(addr netip.Addr, openai bool) { |
| 193 | key := l.keyOf(addr) |
| 194 | l.mu.Lock() |
| 195 | defer l.mu.Unlock() |
| 196 | e, scale := l.entries[key], 1 |
| 197 | if e == nil { |
| 198 | e, scale = &l.overflow, max(1, l.max/1024) |
| 199 | } |
| 200 | e.tokens = min(l.rate*float64(l.multLocked(key, scale, openai)), e.tokens+1) |
| 201 | } |
| 202 | |
| 203 | // admit decides whether a new connection from addr may be served. It returns the function that gives its slot back (once), or the reason it was |
| 204 | // refused. The new-connection bucket is charged for every accepted connection, the concurrent count until release is called. |
| 205 | func (l *ipLimits) admit(addr netip.Addr, openai bool) (release func(), refused string) { |
| 206 | now := l.clk.Now() |
| 207 | key := l.keyOf(addr) |
| 208 | l.mu.Lock() |
| 209 | defer l.mu.Unlock() |
| 210 | e, scale := l.entryLocked(key, now) |
| 211 | mult := l.multLocked(key, scale, openai) |
| 212 | if !openai && now.Before(e.banUntil) { |
| 213 | return nil, refuseBanned |
| 214 | } |
| 215 | burst := l.rate * float64(mult) // the bucket holds a second's worth of new connections and refills at that rate |
| 216 | e.tokens = min(burst, e.tokens+now.Sub(e.refilled).Seconds()*burst) |
| 217 | e.refilled = now |
| 218 | if e.tokens < 1 { |
| 219 | return nil, refuseRate |
| 220 | } |
| 221 | if e.conns >= l.concurrent*mult { |
| 222 | return nil, refuseConcurrent |
| 223 | } |
| 224 | e.tokens-- |
| 225 | e.conns++ |
| 226 | var once sync.Once |
| 227 | return func() { |
| 228 | once.Do(func() { |
| 229 | l.mu.Lock() |
| 230 | e.conns-- |
| 231 | l.mu.Unlock() |
| 232 | }) |
| 233 | }, "" |
| 234 | } |
| 235 | |
| 236 | // strike records that addr tried name, which is not a device that is online (name "" for a ClientHello that is no TLS or names nothing: each one is a |
| 237 | // strike of its own). The same name again within the window costs nothing, so ChatGPT retrying an offline device is not a scan; more than the |
| 238 | // budget of distinct names bans the source for the ban period, which is reported (true) once. An address of the OpenAI list is never struck, nor are the |
| 239 | // sources that do not fit in the table (they share one entry that must never be banned). |
| 240 | func (l *ipLimits) strike(addr netip.Addr, name string, openai bool) (banned bool) { |
| 241 | if openai { |
| 242 | return false |
| 243 | } |
| 244 | now := l.clk.Now() |
| 245 | key := l.keyOf(addr) |
| 246 | l.mu.Lock() |
| 247 | defer l.mu.Unlock() |
| 248 | e, _ := l.entryLocked(key, now) |
| 249 | if e == &l.overflow { |
| 250 | // The overflow entry is a crowd, not a source: a scanner inside it must not ban every other new source for the ban period. Its limits (concurrent |
| 251 | // and rate) are the crowd's whole defence. |
| 252 | return false |
| 253 | } |
| 254 | if now.Sub(e.strikeAt) >= l.window { |
| 255 | e.strikes, e.strikeAt = e.strikes[:0], now |
| 256 | } |
| 257 | var h uint64 |
| 258 | if name != "" { |
| 259 | h = maphash.String(l.seed, name) |
| 260 | for _, seen := range e.strikes { |
| 261 | if seen == h { |
| 262 | return false |
| 263 | } |
| 264 | } |
| 265 | } else { |
| 266 | e.junk++ |
| 267 | h = e.junk // a counter cannot collide with a name's hash in practice, and a repeat of it would only cost a strike |
| 268 | } |
| 269 | e.strikes = append(e.strikes, h) |
| 270 | if len(e.strikes) <= l.strikeMax { |
| 271 | return false |
| 272 | } |
| 273 | e.strikes, e.banUntil = e.strikes[:0], now.Add(l.ban) |
| 274 | l.bans++ |
| 275 | return true |
| 276 | } |
| 277 | |
| 278 | // snapshot is what the metrics report: the size of the table, the sources banned now, the bans and the overflows so far. |
| 279 | func (l *ipLimits) snapshot() (size, banned int, bans, overflowed uint64) { |
| 280 | now := l.clk.Now() |
| 281 | l.mu.Lock() |
| 282 | defer l.mu.Unlock() |
| 283 | for _, e := range l.entries { |
| 284 | if now.Before(e.banUntil) { |
| 285 | banned++ |
| 286 | } |
| 287 | } |
| 288 | return len(l.entries), banned, l.bans, l.overflowed |
| 289 | } |
| 290 | |
| 291 | // bucket is a token bucket: rate tokens a second up to burst. |
| 292 | type bucket struct { |
| 293 | tokens float64 |
| 294 | refilled time.Time |
| 295 | } |
| 296 | |
| 297 | // take spends one token at now; it reports false when there is none. |
| 298 | func (b *bucket) take(now time.Time, rate, burst float64) bool { |
| 299 | if b.refilled.IsZero() { |
| 300 | b.tokens = burst |
| 301 | } else { |
| 302 | b.tokens = min(burst, b.tokens+now.Sub(b.refilled).Seconds()*rate) |
| 303 | } |
| 304 | b.refilled = now |
| 305 | if b.tokens < 1 { |
| 306 | return false |
| 307 | } |
| 308 | b.tokens-- |
| 309 | return true |
| 310 | } |
| 311 | |
| 312 | // logSampler lets a log line that a flood can repeat out at most once a second per reason, with how many were left out: an attack must not fill the disk |
| 313 | // through the log (nor hide the other lines in it). |
| 314 | type logSampler struct { |
| 315 | clk clock |
| 316 | mu sync.Mutex |
| 317 | at map[string]time.Time |
| 318 | cut map[string]int |
| 319 | } |
| 320 | |
| 321 | func newLogSampler(clk clock) *logSampler { |
| 322 | return &logSampler{clk: clk, at: map[string]time.Time{}, cut: map[string]int{}} |
| 323 | } |
| 324 | |
| 325 | // allow reports whether a line for reason may be written now, and how many of its kind were left out since the last one. |
| 326 | func (l *logSampler) allow(reason string) (ok bool, suppressed int) { |
| 327 | now := l.clk.Now() |
| 328 | l.mu.Lock() |
| 329 | defer l.mu.Unlock() |
| 330 | if last, seen := l.at[reason]; seen && now.Sub(last) < time.Second { |
| 331 | l.cut[reason]++ |
| 332 | return false, 0 |
| 333 | } |
| 334 | l.at[reason] = now |
| 335 | suppressed = l.cut[reason] |
| 336 | delete(l.cut, reason) |
| 337 | return true, suppressed |
| 338 | } |