ALL

For reading and verification only. You may not copy, use, modify or redistribute this code. Full notice on the main page.

tunnel/openai_prefixes.go

SHA-256: 12809b4af8aa3e1c69b657a8be18553af43e3857a44e31d5e18f71704b66ce4f
336 lines · 10216 bytes · View on GitHub

1package main
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "log/slog"
10 "net/http"
11 "net/netip"
12 "os"
13 "path/filepath"
14 "slices"
15 "sync/atomic"
16 "time"
17)
18
19// The address ranges OpenAI publishes for its connector traffic (the design plan, section 5). The node reads them once a day into the "openai"
20// source class of an OPEN, which is a QUOTA and not an allowlist: ChatGPT can be made to call any name from those addresses, so the class never lets a
21// stream past a limit; it only keeps ChatGPT's own share of the node (policy node_limits.openai_reserved_pct) and per-source limits larger by a factor
22// (openai_per_ip_factor) for an address that serves many users. The published list is "best effort" (G1 measures the addresses ChatGPT really comes from).
23//
24// A list that is empty, corrupt, too big, or that holds a prefix too wide to be a connector's range (0.0.0.0/0 would make everyone "openai") replaces
25// nothing: the node keeps the last good one, which is also kept in a file (public data) so that a restart during an outage of the publisher does not
26// start with none.
27
28// The words of a refresh's outcome, in the log and in the metrics.
29const (
30 prefixOK = "ok"
31 prefixFetch = "fetch_failed"
32 prefixEmpty = "empty"
33 prefixCorrupt = "corrupt"
34 prefixTooLarge = "too_large"
35 prefixTooMany = "too_many"
36 prefixBadPrefix = "bad_prefix"
37)
38
39var (
40 errPrefixEmpty = errors.New(prefixEmpty)
41 errPrefixCorrupt = errors.New(prefixCorrupt)
42 errPrefixTooLarge = errors.New(prefixTooLarge)
43 errPrefixTooMany = errors.New(prefixTooMany)
44 errPrefixBad = errors.New(prefixBadPrefix)
45)
46
47// prefixSet is an immutable list of prefixes with a lookup that costs one map probe per distinct prefix length.
48type prefixSet struct {
49 set map[netip.Prefix]struct{}
50 lens []int
51}
52
53func newPrefixSet(prefixes []netip.Prefix) *prefixSet {
54 p := &prefixSet{set: make(map[netip.Prefix]struct{}, len(prefixes))}
55 for _, pf := range prefixes {
56 p.set[pf] = struct{}{}
57 if !slices.Contains(p.lens, pf.Bits()) {
58 p.lens = append(p.lens, pf.Bits())
59 }
60 }
61 return p
62}
63
64func (p *prefixSet) len() int { return len(p.set) }
65
66func (p *prefixSet) contains(a netip.Addr) bool {
67 a = a.WithZone("").Unmap()
68 for _, bits := range p.lens {
69 if pf, err := a.Prefix(bits); err == nil {
70 if _, ok := p.set[pf]; ok {
71 return true
72 }
73 }
74 }
75 return false
76}
77
78// prefixLimits are the sanity bounds a list is read under (policy node_limits).
79type prefixLimits struct {
80 maxBytes, maxEntries, minV4, minV6 int
81}
82
83// parsePrefixes reads a published list: {"prefixes": [{"ipv4Prefix": "a.b.c.d/n"}, {"ipv6Prefix": "..."}]} and nothing else of consequence. Every entry
84// must be one of those two shapes with a prefix of its own family, not narrower than the policy's minimum width (a connector range is a /28 or so, not a
85// continent); a single bad entry refuses the whole list, because a list that is partly wrong is not one to trust the rest of.
86func parsePrefixes(body []byte, lim prefixLimits) ([]netip.Prefix, error) {
87 if len(body) > lim.maxBytes {
88 return nil, errPrefixTooLarge
89 }
90 value, err := parseJSON(string(body))
91 if err != nil {
92 return nil, errPrefixCorrupt
93 }
94 doc, ok := value.(map[string]any)
95 if !ok {
96 return nil, errPrefixCorrupt
97 }
98 list, ok := doc["prefixes"].([]any)
99 if !ok {
100 return nil, errPrefixCorrupt
101 }
102 if len(list) == 0 {
103 return nil, errPrefixEmpty
104 }
105 if len(list) > lim.maxEntries {
106 return nil, errPrefixTooMany
107 }
108 out := make([]netip.Prefix, 0, len(list))
109 for _, item := range list {
110 entry, ok := item.(map[string]any)
111 if !ok || len(entry) != 1 {
112 return nil, errPrefixCorrupt
113 }
114 for key, raw := range entry {
115 text, ok := raw.(string)
116 if !ok || (key != "ipv4Prefix" && key != "ipv6Prefix") {
117 return nil, errPrefixCorrupt
118 }
119 pf, err := netip.ParsePrefix(text)
120 if err != nil || pf.Addr().Zone() != "" || pf.Addr().Is4In6() || pf.Addr().IsUnspecified() {
121 return nil, errPrefixBad
122 }
123 minBits := lim.minV6
124 if key == "ipv4Prefix" {
125 if !pf.Addr().Is4() {
126 return nil, errPrefixBad
127 }
128 minBits = lim.minV4
129 } else if !pf.Addr().Is6() {
130 return nil, errPrefixBad
131 }
132 if pf.Bits() < minBits {
133 return nil, errPrefixBad
134 }
135 out = append(out, pf.Masked())
136 }
137 }
138 return out, nil
139}
140
141// encodePrefixes writes a list in the format parsePrefixes reads (what the file of the last good list holds).
142func encodePrefixes(set *prefixSet) []byte {
143 all := make([]netip.Prefix, 0, len(set.set))
144 for pf := range set.set {
145 all = append(all, pf)
146 }
147 slices.SortFunc(all, func(a, b netip.Prefix) int { return a.Addr().Compare(b.Addr()) })
148 entries := make([]map[string]string, 0, len(all))
149 for _, pf := range all {
150 key := "ipv6Prefix"
151 if pf.Addr().Is4() {
152 key = "ipv4Prefix"
153 }
154 entries = append(entries, map[string]string{key: pf.String()})
155 }
156 raw, _ := json.Marshal(map[string]any{"prefixes": entries})
157 return raw
158}
159
160// openaiPrefixes is the node's current list and what keeps it fresh.
161type openaiPrefixes struct {
162 cur atomic.Pointer[prefixSet]
163 at atomic.Int64 // when the list in use was installed (clock nanoseconds since the epoch), 0 when none
164 clk clock
165 log *slog.Logger
166 lim prefixLimits
167 every time.Duration
168 file string
169 fetch func(ctx context.Context) ([]byte, error)
170 count func(result string) // the metrics' counter of outcomes
171}
172
173func newOpenAIPrefixes(cfg *Config, clk clock, count func(string)) *openaiPrefixes {
174 l := policyData.NodeLimits
175 o := &openaiPrefixes{
176 clk: clk, log: cfg.Logger, every: cfg.OpenAIPrefixRefresh, file: cfg.OpenAIPrefixesCache, count: count,
177 lim: prefixLimits{maxBytes: l.OpenAIPrefixBodyMaxBytes, maxEntries: l.OpenAIPrefixMax, minV4: l.OpenAIPrefixMinV4Bits, minV6: l.OpenAIPrefixMinV6Bits},
178 }
179 switch {
180 case cfg.PrefixFetch != nil:
181 o.fetch = cfg.PrefixFetch
182 case cfg.OpenAIPrefixesURL != "":
183 o.fetch = func(ctx context.Context) ([]byte, error) {
184 return httpFetchList(ctx, cfg.OpenAIPrefixesURL, int64(o.lim.maxBytes))
185 }
186 }
187 return o
188}
189
190// contains reports whether a is in the list in use (false when there is none).
191func (o *openaiPrefixes) contains(a netip.Addr) bool {
192 if set := o.cur.Load(); set != nil {
193 return set.contains(a)
194 }
195 return false
196}
197
198// size is how many prefixes are in use, and age how long ago the list was installed (-1 when there is none).
199func (o *openaiPrefixes) size() int {
200 if set := o.cur.Load(); set != nil {
201 return set.len()
202 }
203 return 0
204}
205
206func (o *openaiPrefixes) age() time.Duration {
207 at := o.at.Load()
208 if at == 0 {
209 return -1
210 }
211 return o.clk.Now().Sub(time.Unix(0, at))
212}
213
214// apply reads body and, when it is a good list, puts it in use and (unless it came from the file) keeps it in the file. A list that is not good
215// changes nothing, and the error says why in one of the outcome words.
216func (o *openaiPrefixes) apply(body []byte, fromFile bool) error {
217 prefixes, err := parsePrefixes(body, o.lim)
218 if err != nil {
219 return err
220 }
221 set := newPrefixSet(prefixes)
222 o.cur.Store(set)
223 o.at.Store(o.clk.Now().UnixNano())
224 if !fromFile && o.file != "" {
225 if werr := writeAtomic(o.file, encodePrefixes(set)); werr != nil {
226 o.log.Warn("openai_prefixes_not_kept", "reason", "write_failed")
227 }
228 }
229 return nil
230}
231
232// writeAtomic replaces path with data in one rename, so a crash never leaves half a list.
233func writeAtomic(path string, data []byte) error {
234 tmp, err := os.CreateTemp(filepath.Dir(path), ".prefixes-*")
235 if err != nil {
236 return err
237 }
238 _, werr := tmp.Write(data)
239 cerr := tmp.Close()
240 if werr != nil || cerr != nil {
241 _ = os.Remove(tmp.Name())
242 return errors.Join(werr, cerr)
243 }
244 if err := os.Chmod(tmp.Name(), 0o644); err != nil {
245 _ = os.Remove(tmp.Name())
246 return err
247 }
248 if err := os.Rename(tmp.Name(), path); err != nil {
249 _ = os.Remove(tmp.Name())
250 return err
251 }
252 return nil
253}
254
255// refresh fetches and applies once, and says what came of it.
256func (o *openaiPrefixes) refresh(ctx context.Context) string {
257 if o.fetch == nil {
258 return prefixFetch
259 }
260 fctx, cancel := context.WithTimeout(ctx, 30*time.Second)
261 defer cancel()
262 body, err := o.fetch(fctx)
263 result := prefixOK
264 switch {
265 case err != nil && errors.Is(err, errPrefixTooLarge):
266 result = prefixTooLarge
267 case err != nil:
268 result = prefixFetch
269 default:
270 if aerr := o.apply(body, false); aerr != nil {
271 result = aerr.Error()
272 }
273 }
274 o.count(result)
275 if result == prefixOK {
276 o.log.Info("openai_prefixes_refreshed", "prefixes", o.size())
277 } else {
278 o.log.Warn("openai_prefixes_refused", "reason", result, "keeping", o.size())
279 }
280 return result
281}
282
283// Run is the node's background task: the last good list from the file, then a refresh at start and every refresh period (a failed one is tried again
284// after minutes, not a day), until the node ends.
285func (o *openaiPrefixes) Run(ctx context.Context) {
286 if o.file != "" {
287 if body, err := os.ReadFile(o.file); err == nil {
288 if aerr := o.apply(body, true); aerr != nil {
289 o.log.Warn("openai_prefixes_file_refused", "reason", aerr.Error())
290 }
291 }
292 }
293 if o.fetch == nil {
294 return
295 }
296 failures := 0
297 for {
298 wait := o.every
299 if o.refresh(ctx) == prefixOK {
300 failures = 0
301 } else {
302 failures++
303 wait = min(5*time.Minute<<min(failures-1, 4), o.every)
304 }
305 if !o.clk.sleep(ctx, jitter(wait, 10)) {
306 return
307 }
308 }
309}
310
311// httpFetchList is the one request the node makes to the publisher: HTTPS, no redirect, no credentials, at most limit bytes.
312func httpFetchList(ctx context.Context, url string, limit int64) ([]byte, error) {
313 req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
314 if err != nil {
315 return nil, err
316 }
317 req.Header.Set("accept", "application/json")
318 req.Header.Set("user-agent", "all-tunnel/"+version)
319 client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
320 resp, err := client.Do(req)
321 if err != nil {
322 return nil, err
323 }
324 defer resp.Body.Close()
325 if resp.StatusCode != http.StatusOK {
326 return nil, fmt.Errorf("status %d", resp.StatusCode)
327 }
328 body, err := io.ReadAll(io.LimitReader(resp.Body, limit+1))
329 if err != nil {
330 return nil, err
331 }
332 if int64(len(body)) > limit {
333 return nil, errPrefixTooLarge
334 }
335 return body, nil
336}