ALL

Solo para leer y verificar. No puedes copiar, usar, modificar ni redistribuir este código. Aviso completo en la página principal.

tunnel/openai_prefixes_test.go

SHA-256: 2bfece3d472d348feede9cfbb3d75eaa3c8d6cedb1c7d676f476a771326e41a3
323 líneas · 12265 bytes · Ver en GitHub

1package main
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "log/slog"
8 "net/http"
9 "net/http/httptest"
10 "net/netip"
11 "os"
12 "path/filepath"
13 "strings"
14 "sync"
15 "sync/atomic"
16 "testing"
17 "time"
18)
19
20func testPrefixLimits() prefixLimits {
21 l := policyData.NodeLimits
22 return prefixLimits{maxBytes: l.OpenAIPrefixBodyMaxBytes, maxEntries: l.OpenAIPrefixMax, minV4: l.OpenAIPrefixMinV4Bits, minV6: l.OpenAIPrefixMinV6Bits}
23}
24
25func listOf(prefixes ...string) string {
26 var items []string
27 for _, p := range prefixes {
28 key := "ipv4Prefix"
29 if strings.Contains(p, ":") {
30 key = "ipv6Prefix"
31 }
32 items = append(items, fmt.Sprintf(`{%q:%q}`, key, p))
33 }
34 return `{"creationTime":"2026-10-01T00:00:00.000000","prefixes":[` + strings.Join(items, ",") + `]}`
35}
36
37func TestPrefixes_AGoodListIsReadAndLooksUpIPv4AndIPv6(t *testing.T) {
38 got, err := parsePrefixes([]byte(listOf("203.0.113.0/28", "198.51.100.64/26", "2001:db8:abcd::/48", "192.0.2.5/24")), testPrefixLimits())
39 if err != nil {
40 t.Fatal(err)
41 }
42 set := newPrefixSet(got)
43 for addr, want := range map[string]bool{
44 "203.0.113.7": true, "203.0.113.16": false, "198.51.100.100": true, "198.51.100.200": false, "2001:db8:abcd:1::9": true, "2001:db8:abce::1": false,
45 "192.0.2.99": true, // 192.0.2.5/24 is the /24
46 "::ffff:203.0.113.7": true, "198.51.100.8": false,
47 } {
48 if set.contains(netip.MustParseAddr(addr)) != want {
49 t.Errorf("%s: want %v", addr, want)
50 }
51 }
52 if set.len() != 4 {
53 t.Fatalf("%d prefixes", set.len())
54 }
55}
56
57func TestPrefixes_ACorruptEmptyHugeOrTooWideListIsRefused(t *testing.T) {
58 lim := testPrefixLimits()
59 manyItems := make([]string, lim.maxEntries+1)
60 for i := range manyItems {
61 manyItems[i] = fmt.Sprintf("10.%d.%d.0/24", i/256, i%256)
62 }
63 for name, c := range map[string]struct {
64 body string
65 want error
66 }{
67 "empty body": {"", errPrefixCorrupt},
68 "not json": {"<html>", errPrefixCorrupt},
69 "a list not an object": {`[{"ipv4Prefix":"203.0.113.0/28"}]`, errPrefixCorrupt},
70 "no prefixes key": {`{"x":[]}`, errPrefixCorrupt},
71 "an empty list": {`{"prefixes":[]}`, errPrefixEmpty},
72 "an entry that is a string": {`{"prefixes":["203.0.113.0/28"]}`, errPrefixCorrupt},
73 "an unknown key": {`{"prefixes":[{"cidr":"203.0.113.0/28"}]}`, errPrefixCorrupt},
74 "two keys in one entry": {`{"prefixes":[{"ipv4Prefix":"203.0.113.0/28","ipv6Prefix":"2001:db8::/48"}]}`, errPrefixCorrupt},
75 "a number as the prefix": {`{"prefixes":[{"ipv4Prefix":5}]}`, errPrefixCorrupt},
76 "not a prefix": {`{"prefixes":[{"ipv4Prefix":"203.0.113.0"}]}`, errPrefixBad},
77 "a v6 prefix under a v4 key": {`{"prefixes":[{"ipv4Prefix":"2001:db8::/48"}]}`, errPrefixBad},
78 "a v4 prefix under a v6 key": {`{"prefixes":[{"ipv6Prefix":"203.0.113.0/28"}]}`, errPrefixBad},
79 "everything": {listOf("0.0.0.0/0"), errPrefixBad},
80 "half the internet": {listOf("128.0.0.0/1"), errPrefixBad},
81 "a wide v4 range": {listOf("10.0.0.0/8"), errPrefixBad},
82 "a wide v6 range": {listOf("2000::/3"), errPrefixBad},
83 "a mapped address": {`{"prefixes":[{"ipv6Prefix":"::ffff:203.0.113.0/120"}]}`, errPrefixBad},
84 "two good prefixes": {listOf("203.0.113.0/28", "198.51.100.0/26") + "", nil},
85 "too many entries": {listOf(manyItems...), errPrefixTooMany},
86 "a body that is too big": {`{"prefixes":[],"pad":"` + strings.Repeat("x", lim.maxBytes) + `"}`, errPrefixTooLarge},
87 "a repeated key": {`{"prefixes":[{"ipv4Prefix":"203.0.113.0/28","ipv4Prefix":"198.51.100.0/26"}]}`, errPrefixCorrupt},
88 } {
89 _, err := parsePrefixes([]byte(c.body), lim)
90 if c.want == nil {
91 if err != nil {
92 t.Errorf("%s: %v", name, err)
93 }
94 continue
95 }
96 if !errors.Is(err, c.want) {
97 t.Errorf("%s: got %v, want %v", name, err, c.want)
98 }
99 }
100 // one bad entry refuses the whole list
101 if _, err := parsePrefixes([]byte(`{"prefixes":[{"ipv4Prefix":"203.0.113.0/28"},{"ipv4Prefix":"0.0.0.0/0"}]}`), lim); err == nil {
102 t.Error("a list with one too-wide prefix was taken")
103 }
104}
105
106type prefixHarness struct {
107 o *openaiPrefixes
108 fc *fakeClock
109 results sync.Map // outcome word -> *atomic.Int64
110 logs *syncBuf
111 file string
112}
113
114func (h *prefixHarness) count(word string) int64 {
115 if v, ok := h.results.Load(word); ok {
116 return v.(*atomic.Int64).Load()
117 }
118 return 0
119}
120
121func newPrefixHarness(t *testing.T, fetch func(context.Context) ([]byte, error)) *prefixHarness {
122 t.Helper()
123 h := &prefixHarness{fc: newFakeClock(), logs: &syncBuf{}, file: filepath.Join(t.TempDir(), "prefixes.json")}
124 cfg := DefaultConfig()
125 cfg.OpenAIPrefixesCache = h.file
126 cfg.PrefixFetch = fetch
127 cfg.Logger = slog.New(slog.NewJSONHandler(h.logs, nil))
128 h.o = newOpenAIPrefixes(&cfg, newClock(h.fc.Now, h.fc.AfterFunc), func(word string) {
129 v, _ := h.results.LoadOrStore(word, &atomic.Int64{})
130 v.(*atomic.Int64).Add(1)
131 })
132 return h
133}
134
135func TestPrefixes_ARefreshThatFailsKeepsTheLastGoodListAndAFileSurvivesARestart(t *testing.T) {
136 var mu sync.Mutex
137 body, fail := listOf("203.0.113.0/28"), error(nil)
138 h := newPrefixHarness(t, func(context.Context) ([]byte, error) {
139 mu.Lock()
140 defer mu.Unlock()
141 return []byte(body), fail
142 })
143 in := netip.MustParseAddr("203.0.113.5")
144 if h.o.contains(in) {
145 t.Fatal("a list that nobody read is in use")
146 }
147 if got := h.o.refresh(context.Background()); got != prefixOK || !h.o.contains(in) || h.o.size() != 1 {
148 t.Fatalf("first refresh: %s, size %d", got, h.o.size())
149 }
150 good := h.o.cur.Load()
151 for name, c := range map[string]struct {
152 body string
153 err error
154 want string
155 }{
156 "the publisher is down": {"", errors.New("connection refused"), prefixFetch},
157 "an empty list": {`{"prefixes":[]}`, nil, prefixEmpty},
158 "a corrupt body": {"{{{", nil, prefixCorrupt},
159 "a list of the internet": {listOf("0.0.0.0/0"), nil, prefixBadPrefix},
160 "a body past the limit": {"", errPrefixTooLarge, prefixTooLarge},
161 "a list with too many": {listOf(strings.Split(strings.Repeat("10.1.1.0/24,", 2049), ",")[:2049]...), nil, prefixTooMany},
162 "nothing but whitespace": {" ", nil, prefixCorrupt},
163 "a list of another shape": {`{"prefixes":[{"cidr":"203.0.113.0/28"}]}`, nil, prefixCorrupt},
164 } {
165 mu.Lock()
166 body, fail = c.body, c.err
167 mu.Unlock()
168 if got := h.o.refresh(context.Background()); got != c.want {
169 t.Errorf("%s: outcome %q, want %q", name, got, c.want)
170 }
171 if h.o.cur.Load() != good || !h.o.contains(in) {
172 t.Fatalf("%s: the list in use changed", name)
173 }
174 }
175 // a new good list replaces it, and the file holds it
176 mu.Lock()
177 body, fail = listOf("198.51.100.0/26", "2001:db8::/32"), nil
178 mu.Unlock()
179 if h.o.refresh(context.Background()) != prefixOK || h.o.contains(in) || !h.o.contains(netip.MustParseAddr("198.51.100.9")) {
180 t.Fatal("the new list was not taken")
181 }
182 // a restart with the publisher down: the file is the list
183 restarted := newPrefixHarness(t, func(context.Context) ([]byte, error) { return nil, errors.New("down") })
184 raw, err := os.ReadFile(h.file)
185 if err != nil {
186 t.Fatal(err)
187 }
188 if err := os.WriteFile(restarted.file, raw, 0o600); err != nil {
189 t.Fatal(err)
190 }
191 ctx, cancel := context.WithCancel(context.Background())
192 done := make(chan struct{})
193 go func() { restarted.o.Run(ctx); close(done) }()
194 eventually(t, "the list from the file", 3*time.Second, func() bool { return restarted.o.size() == 2 })
195 if !restarted.o.contains(netip.MustParseAddr("2001:db8:1::1")) {
196 t.Fatal("the IPv6 prefix of the file is not in use")
197 }
198 eventually(t, "the failed refresh was counted", 3*time.Second, func() bool { return restarted.count(prefixFetch) >= 1 })
199 if restarted.o.size() != 2 {
200 t.Fatal("a failed refresh dropped the list of the file")
201 }
202 cancel()
203 <-done
204 if info, err := os.Stat(h.file); err != nil || info.Mode().Perm() != 0o644 {
205 t.Fatalf("the file: %v %v", info, err)
206 }
207 // a file that is corrupt is ignored
208 if err := os.WriteFile(restarted.file, []byte("garbage"), 0o600); err != nil {
209 t.Fatal(err)
210 }
211 again := newPrefixHarness(t, nil)
212 again.file = restarted.file
213 again.o.file = restarted.file
214 again.o.Run(context.Background())
215 if again.o.size() != 0 || !strings.Contains(again.logs.String(), "openai_prefixes_file_refused") {
216 t.Fatalf("size %d logs %s", again.o.size(), again.logs.String())
217 }
218}
219
220func TestPrefixes_RunRefreshesDailyAndRetriesAFailureSoonerNotAtOnce(t *testing.T) {
221 var calls atomic.Int64
222 var fail atomic.Bool
223 fail.Store(true)
224 h := newPrefixHarness(t, func(context.Context) ([]byte, error) {
225 calls.Add(1)
226 if fail.Load() {
227 return nil, errors.New("down")
228 }
229 return []byte(listOf("203.0.113.0/28")), nil
230 })
231 ctx, cancel := context.WithCancel(context.Background())
232 defer cancel()
233 go h.o.Run(ctx)
234 eventually(t, "the first attempt", 3*time.Second, func() bool { return calls.Load() == 1 })
235 // 4 minutes later nothing was retried (the retry waits minutes, not seconds: the publisher is not hammered)
236 h.fc.Advance(4 * time.Minute)
237 time.Sleep(50 * time.Millisecond)
238 if calls.Load() != 1 {
239 t.Fatalf("%d calls after 4 minutes", calls.Load())
240 }
241 fail.Store(false)
242 h.fc.Advance(2 * time.Minute)
243 eventually(t, "the retry and the good list", 3*time.Second, func() bool { return h.o.size() == 1 })
244 n := calls.Load()
245 // then the next refresh is a day later (give or take the jitter that keeps nodes from calling in step)
246 h.fc.Advance(20 * time.Hour)
247 time.Sleep(50 * time.Millisecond)
248 if calls.Load() != n {
249 t.Fatalf("refreshed again before the day was over (%d calls, was %d)", calls.Load(), n)
250 }
251 h.fc.Advance(8 * time.Hour)
252 eventually(t, "the daily refresh", 3*time.Second, func() bool { return calls.Load() == n+1 })
253 if age := h.o.age(); age < 0 || age > time.Minute {
254 t.Fatalf("age %s", age)
255 }
256}
257
258func TestPrefixes_TheHTTPFetchReadsOneUrlWithoutRedirectsAndWithinTheLimit(t *testing.T) {
259 var hits atomic.Int64
260 mux := http.NewServeMux()
261 mux.HandleFunc("/list.json", func(w http.ResponseWriter, r *http.Request) {
262 hits.Add(1)
263 if r.Header.Get("cookie") != "" || r.Header.Get("authorization") != "" {
264 t.Error("the fetch sent credentials")
265 }
266 _, _ = w.Write([]byte(listOf("203.0.113.0/28")))
267 })
268 mux.HandleFunc("/moved", func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, "/list.json", http.StatusFound) })
269 mux.HandleFunc("/big", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write(make([]byte, 3000)) })
270 mux.HandleFunc("/missing", func(w http.ResponseWriter, r *http.Request) { http.NotFound(w, r) })
271 srv := httptest.NewServer(mux)
272 defer srv.Close()
273 ctx := context.Background()
274 if body, err := httpFetchList(ctx, srv.URL+"/list.json", 4096); err != nil || !strings.Contains(string(body), "203.0.113.0/28") {
275 t.Fatalf("%q %v", body, err)
276 }
277 if _, err := httpFetchList(ctx, srv.URL+"/moved", 4096); err == nil {
278 t.Fatal("a redirect was followed")
279 }
280 if _, err := httpFetchList(ctx, srv.URL+"/big", 1000); !errors.Is(err, errPrefixTooLarge) {
281 t.Fatalf("%v", err)
282 }
283 if _, err := httpFetchList(ctx, srv.URL+"/missing", 4096); err == nil {
284 t.Fatal("a 404 is not a list")
285 }
286}
287
288func TestNode_ASourceInTheOpenAIListIsClassedOpenAIAndGetsLargerPerSourceLimits(t *testing.T) {
289 release := make(chan struct{})
290 e := newVIPEnv(t, func(e *testEnv, c *Config) {
291 c.PerIPConcurrent, c.PerIPNewPerS, c.OpenAIPerIPFactor = 2, 1000, 4
292 c.PrefixFetch = func(context.Context) ([]byte, error) { return []byte(listOf("192.0.2.0/24")), nil }
293 })
294 eventually(t, "the list", 3*time.Second, func() bool { return e.srv.openai.size() == 1 })
295 d := e.newDevice()
296 dc := e.connectOK(d)
297 ds := dc.serve(holdHandler(d, release))
298 defer close(release)
299 held := func(ip string) int {
300 n := 0
301 for i := 0; i < 20; i++ {
302 if _, err := e.clientFrom(ip, d); err == nil {
303 n++
304 }
305 }
306 return n
307 }
308 if n := held("198.51.100.1"); n != 2 {
309 t.Fatalf("an ordinary source held %d, cap 2", n)
310 }
311 if n := held("192.0.2.1"); n != 8 {
312 t.Fatalf("an OpenAI address held %d, want 2 x 4", n)
313 }
314 seen := map[string]bool{}
315 for _, o := range ds.openList() {
316 seen[o.str("peer_ip")+" "+o.str("class")] = true
317 }
318 if !seen["192.0.2.1 openai"] || !seen["198.51.100.1 anon"] {
319 t.Fatalf("OPENs: %v", seen)
320 }
321 e.checkTranscript()
322 _ = fmt.Sprint
323}