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.go
| 1 | package main |
| 2 | |
| 3 | import ( |
| 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. |
| 29 | const ( |
| 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 | |
| 39 | var ( |
| 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. |
| 48 | type prefixSet struct { |
| 49 | set map[netip.Prefix]struct{} |
| 50 | lens []int |
| 51 | } |
| 52 | |
| 53 | func 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 | |
| 64 | func (p *prefixSet) len() int { return len(p.set) } |
| 65 | |
| 66 | func (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). |
| 79 | type 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. |
| 86 | func 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). |
| 142 | func 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. |
| 161 | type 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 | |
| 173 | func 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). |
| 191 | func (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). |
| 199 | func (o *openaiPrefixes) size() int { |
| 200 | if set := o.cur.Load(); set != nil { |
| 201 | return set.len() |
| 202 | } |
| 203 | return 0 |
| 204 | } |
| 205 | |
| 206 | func (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. |
| 216 | func (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. |
| 233 | func 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. |
| 256 | func (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. |
| 285 | func (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. |
| 312 | func 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 | } |