| @@ -216,9 +216,13 @@ func NewZeroconfBrowser(ifaces ...net.Interface) *ZeroconfBrowser { · 215 unmodified lines |
| 216 | 216 | } |
| 217 | 217 | |
| 218 | 218 | const ( |
| 219 | | − maxDiscoveredServices = 256 |
| 220 | | − maxBrowseTXTRecords = 64 |
| 221 | | − maxBrowseTXTBytes = 16 * 1024 |
| 219 | + maxDiscoveredServices = 256 |
| 220 | + maxBrowseEntries = 256 |
| 221 | + maxBrowseTXTRecords = 64 |
| 222 | + maxBrowseTXTBytes = 16 * 1024 |
| 223 | + maxDiscoveredAddresses = 32 |
| 224 | + maxDiscoveredMetadataItems = 32 |
| 225 | + maxMetadataItemBytes = 64 |
| 222 | 226 | ) |
| 223 | 227 | |
| 224 | 228 | func browseEntryWithinLimits(entry *zeroconf.ServiceEntry) bool { |
| @@ -251,40 +255,48 @@ func (b *ZeroconfBrowser) Browse(ctx context.Context, serviceType, domain string · 26 unmodified lines |
| 251 | 255 | } |
| 252 | 256 | |
| 253 | 257 | entries := make(chan *zeroconf.ServiceEntry, 32) |
| 258 | + browseCtx, cancelBrowse := context.WithCancel(ctx) |
| 259 | + defer cancelBrowse() |
| 260 | + |
| 254 | 261 | var discovered []DiscoveredService |
| 255 | 262 | var mu sync.Mutex |
| 256 | 263 | seen := make(map[string]int) |
| 257 | 264 | |
| 258 | | − // Collect entries in background until channel is closed by zeroconf's params.done() |
| 265 | + // Stop the underlying zeroconf client after a bounded number of entries. This |
| 266 | + // is separate from the caller's context so a LAN flood cannot grow the |
| 267 | + // dependency's internal sentEntries cache for the full browse lifetime. |
| 259 | 268 | doneCh := make(chan struct{}) |
| 260 | 269 | go func() { |
| 261 | 270 | defer close(doneCh) |
| 271 | + entriesSeen := 0 |
| 262 | 272 | for entry := range entries { |
| 263 | | − if !browseEntryWithinLimits(entry) { |
| 264 | | − continue |
| 265 | | − } |
| 266 | | − svc := entryToDiscovered(entry) |
| 267 | | − if svc.Port == 0 || (len(svc.IPv4) == 0 && len(svc.IPv6) == 0) { |
| 273 | + if entriesSeen >= maxBrowseEntries { |
| 268 | 274 | continue |
| 269 | 275 | } |
| 270 | | − key := discoveredServiceKey(svc) |
| 271 | | − mu.Lock() |
| 272 | | − if index, ok := seen[key]; ok { |
| 273 | | − mergeDiscoveredService(&discovered[index], svc) |
| 274 | | − mu.Unlock() |
| 275 | | − continue |
| 276 | + entriesSeen++ |
| 277 | + |
| 278 | + if browseEntryWithinLimits(entry) { |
| 279 | + svc := entryToDiscovered(entry) |
| 280 | + if svc.Port != 0 && (len(svc.IPv4) != 0 || len(svc.IPv6) != 0) { |
| 281 | + key := discoveredServiceKey(svc) |
| 282 | + mu.Lock() |
| 283 | + if index, ok := seen[key]; ok { |
| 284 | + mergeDiscoveredService(&discovered[index], svc) |
| 285 | + } else if len(discovered) < maxDiscoveredServices { |
| 286 | + seen[key] = len(discovered) |
| 287 | + discovered = append(discovered, svc) |
| 288 | + } |
| 289 | + mu.Unlock() |
| 290 | + } |
| 276 | 291 | } |
| 277 | | − if len(discovered) >= maxDiscoveredServices { |
| 278 | | − mu.Unlock() |
| 279 | | − continue |
| 292 | + |
| 293 | + if entriesSeen == maxBrowseEntries { |
| 294 | + cancelBrowse() |
| 280 | 295 | } |
| 281 | | − seen[key] = len(discovered) |
| 282 | | − discovered = append(discovered, svc) |
| 283 | | − mu.Unlock() |
| 284 | 296 | } |
| 285 | 297 | }() |
| 286 | 298 | |
| 287 | | − errBrowse := zeroconf.Browse(ctx, serviceType, domain, entries, b.options...) |
| 299 | + errBrowse := zeroconf.Browse(browseCtx, serviceType, domain, entries, b.options...) |
| 288 | 300 | if errBrowse != nil { |
| 289 | 301 | // zeroconf may already have closed entries after a runtime error. |
| 290 | 302 | closeBrowseEntries(entries) |
| @@ -332,9 +344,7 @@ func mergeDiscoveredService(dst *DiscoveredService, src DiscoveredService) { · 41 unmodified lines |
| 332 | 344 | if dst.RawTXT == nil { |
| 333 | 345 | dst.RawTXT = make(map[string]string) |
| 334 | 346 | } |
| 335 | | − for key, value := range src.RawTXT { |
| 336 | | − dst.RawTXT[key] = value |
| 337 | | − } |
| 347 | + mergeRawTXTRecords(dst.RawTXT, src.RawTXT) |
| 338 | 348 | if value, ok := dst.RawTXT["auth_required"]; ok { |
| 339 | 349 | dst.AuthRequired = strings.EqualFold(strings.TrimSpace(value), "true") |
| 340 | 350 | } |
| @@ -349,9 +359,43 @@ func mergeDiscoveredService(dst *DiscoveredService, src DiscoveredService) { · 8 unmodified lines |
| 349 | 359 | } |
| 350 | 360 | } |
| 351 | 361 | |
| 362 | +func mergeRawTXTRecords(dst, src map[string]string) { |
| 363 | + totalBytes := rawTXTMapBytes(dst) |
| 364 | + for key, value := range src { |
| 365 | + if key == "" { |
| 366 | + continue |
| 367 | + } |
| 368 | + replacedBytes := 0 |
| 369 | + _, exists := dst[key] |
| 370 | + if exists { |
| 371 | + replacedBytes = rawTXTRecordBytes(key, dst[key]) |
| 372 | + } else if len(dst) >= maxBrowseTXTRecords { |
| 373 | + continue |
| 374 | + } |
| 375 | + nextBytes := totalBytes - replacedBytes + rawTXTRecordBytes(key, value) |
| 376 | + if nextBytes > maxBrowseTXTBytes { |
| 377 | + continue |
| 378 | + } |
| 379 | + dst[key] = value |
| 380 | + totalBytes = nextBytes |
| 381 | + } |
| 382 | +} |
| 383 | + |
| 384 | +func rawTXTMapBytes(records map[string]string) int { |
| 385 | + totalBytes := 0 |
| 386 | + for key, value := range records { |
| 387 | + totalBytes += rawTXTRecordBytes(key, value) |
| 388 | + } |
| 389 | + return totalBytes |
| 390 | +} |
| 391 | + |
| 392 | +func rawTXTRecordBytes(key, value string) int { |
| 393 | + return len(key) + len(value) + 1 |
| 394 | +} |
| 395 | + |
| 352 | 396 | func appendUniqueIPs(dst, src []net.IP) []net.IP { |
| 353 | 397 | for _, ip := range src { |
| 354 | | − if ip == nil { |
| 398 | + if ip == nil || len(dst) >= maxDiscoveredAddresses { |
| 355 | 399 | continue |
| 356 | 400 | } |
| 357 | 401 | duplicate := false |
| @@ -370,7 +414,8 @@ func appendUniqueIPs(dst, src []net.IP) []net.IP { · 12 unmodified lines |
| 370 | 414 | |
| 371 | 415 | func appendUniqueStrings(dst, src []string) []string { |
| 372 | 416 | for _, value := range src { |
| 373 | | − if value == "" { |
| 417 | + value = strings.TrimSpace(value) |
| 418 | + if value == "" || len(value) > maxMetadataItemBytes || len(dst) >= maxDiscoveredMetadataItems { |
| 374 | 419 | continue |
| 375 | 420 | } |
| 376 | 421 | duplicate := false |
| @@ -387,6 +432,25 @@ func appendUniqueStrings(dst, src []string) []string { · 10 unmodified lines |
| 387 | 432 | return dst |
| 388 | 433 | } |
| 389 | 434 | |
| 435 | +func parseTXTList(value string) []string { |
| 436 | + var result []string |
| 437 | + for len(value) > 0 && len(result) < maxDiscoveredMetadataItems { |
| 438 | + item := value |
| 439 | + if before, after, ok := strings.Cut(value, ","); ok { |
| 440 | + item = before |
| 441 | + value = after |
| 442 | + } else { |
| 443 | + value = "" |
| 444 | + } |
| 445 | + item = strings.TrimSpace(item) |
| 446 | + if item == "" || len(item) > maxMetadataItemBytes { |
| 447 | + continue |
| 448 | + } |
| 449 | + result = appendUniqueStrings(result, []string{item}) |
| 450 | + } |
| 451 | + return result |
| 452 | +} |
| 453 | + |
| 390 | 454 | // BrowseWithFallback discovers all AI gateways on the LAN (_ai-gateway._tcp) and prioritizes CPA instances. |
| 391 | 455 | func (b *ZeroconfBrowser) BrowseWithFallback(ctx context.Context) ([]DiscoveredService, error) { |
| 392 | 456 | allGateways, errMain := b.Browse(ctx, DefaultServiceType, DefaultDomain) |
| @@ -450,13 +514,13 @@ func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { · 57 unmodified lines |
| 450 | 514 | svc.AuthRequired = true |
| 451 | 515 | } |
| 452 | 516 | if methods := parsed["auth_methods"]; methods != "" { |
| 453 | | − svc.AuthMethods = strings.Split(methods, ",") |
| 517 | + svc.AuthMethods = parseTXTList(methods) |
| 454 | 518 | } |
| 455 | 519 | if protos := parsed["protocols"]; protos != "" { |
| 456 | | − svc.Protocols = strings.Split(protos, ",") |
| 520 | + svc.Protocols = parseTXTList(protos) |
| 457 | 521 | } |
| 458 | 522 | if feats := parsed["features"]; feats != "" { |
| 459 | | − svc.Features = strings.Split(feats, ",") |
| 523 | + svc.Features = parseTXTList(feats) |
| 460 | 524 | } |
| 461 | 525 | |
| 462 | 526 | if v, ok := parsed["api_openai"]; ok { |
| @@ -479,12 +543,15 @@ func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { · 16 unmodified lines |
| 479 | 543 | } |
| 480 | 544 | |
| 481 | 545 | func filterUsableIPs(ips []net.IP) []net.IP { |
| 482 | | − usable := make([]net.IP, 0, len(ips)) |
| 546 | + usable := make([]net.IP, 0, min(len(ips), maxDiscoveredAddresses)) |
| 483 | 547 | for _, ip := range ips { |
| 484 | 548 | if ip == nil || ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() { |
| 485 | 549 | continue |
| 486 | 550 | } |
| 487 | | − usable = append(usable, ip) |
| 551 | + usable = append(usable, append(net.IP(nil), ip...)) |
| 552 | + if len(usable) >= maxDiscoveredAddresses { |
| 553 | + break |
| 554 | + } |
| 488 | 555 | } |
| 489 | 556 | return usable |
| 490 | 557 | } |