| @@ -216,9 +216,13 @@ func NewZeroconfBrowser(ifaces ...net.Interface) *ZeroconfBrowser { · 215 unmodified lines |
| 216 | } | 216 | } |
| 217 | | 217 | |
| 218 | const ( | 218 | const ( |
| 219 | maxDiscoveredServices = 256 | 219 | maxDiscoveredServices = 256 |
| 220 | maxBrowseTXTRecords = 64 | 220 | maxBrowseEntries = 256 |
| 221 | maxBrowseTXTBytes = 16 * 1024 | 221 | maxBrowseTXTRecords = 64 |
| | 222 | maxBrowseTXTBytes = 16 * 1024 |
| | 223 | maxDiscoveredAddresses = 32 |
| | 224 | maxDiscoveredMetadataItems = 32 |
| | 225 | maxMetadataItemBytes = 64 |
| 222 | ) | 226 | ) |
| 223 | | 227 | |
| 224 | func browseEntryWithinLimits(entry *zeroconf.ServiceEntry) bool { | 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 | entries := make(chan *zeroconf.ServiceEntry, 32) | 257 | entries := make(chan *zeroconf.ServiceEntry, 32) |
| | 258 | browseCtx, cancelBrowse := context.WithCancel(ctx) |
| | 259 | defer cancelBrowse() |
| | 260 | |
| 254 | var discovered []DiscoveredService | 261 | var discovered []DiscoveredService |
| 255 | var mu sync.Mutex | 262 | var mu sync.Mutex |
| 256 | seen := make(map[string]int) | 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 | doneCh := make(chan struct{}) | 268 | doneCh := make(chan struct{}) |
| 260 | go func() { | 269 | go func() { |
| 261 | defer close(doneCh) | 270 | defer close(doneCh) |
| | 271 | entriesSeen := 0 |
| 262 | for entry := range entries { | 272 | for entry := range entries { |
| 263 | if !browseEntryWithinLimits(entry) { | 273 | if entriesSeen >= maxBrowseEntries { |
| 264 | continue | | |
| 265 | } | | |
| 266 | svc := entryToDiscovered(entry) | | |
| 267 | if svc.Port == 0 || (len(svc.IPv4) == 0 && len(svc.IPv6) == 0) { | | |
| 268 | continue | 274 | continue |
| 269 | } | 275 | } |
| 270 | key := discoveredServiceKey(svc) | 276 | entriesSeen++ |
| 271 | mu.Lock() | 277 | |
| 272 | if index, ok := seen[key]; ok { | 278 | if browseEntryWithinLimits(entry) { |
| 273 | mergeDiscoveredService(&discovered[index], svc) | 279 | svc := entryToDiscovered(entry) |
| 274 | mu.Unlock() | 280 | if svc.Port != 0 && (len(svc.IPv4) != 0 || len(svc.IPv6) != 0) { |
| 275 | continue | 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 { | 292 | |
| 278 | mu.Unlock() | 293 | if entriesSeen == maxBrowseEntries { |
| 279 | continue | 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 | if errBrowse != nil { | 300 | if errBrowse != nil { |
| 289 | // zeroconf may already have closed entries after a runtime error. | 301 | // zeroconf may already have closed entries after a runtime error. |
| 290 | closeBrowseEntries(entries) | 302 | closeBrowseEntries(entries) |
| @@ -332,9 +344,7 @@ func mergeDiscoveredService(dst *DiscoveredService, src DiscoveredService) { · 41 unmodified lines |
| 332 | if dst.RawTXT == nil { | 344 | if dst.RawTXT == nil { |
| 333 | dst.RawTXT = make(map[string]string) | 345 | dst.RawTXT = make(map[string]string) |
| 334 | } | 346 | } |
| 335 | for key, value := range src.RawTXT { | 347 | mergeRawTXTRecords(dst.RawTXT, src.RawTXT) |
| 336 | dst.RawTXT[key] = value | | |
| 337 | } | | |
| 338 | if value, ok := dst.RawTXT["auth_required"]; ok { | 348 | if value, ok := dst.RawTXT["auth_required"]; ok { |
| 339 | dst.AuthRequired = strings.EqualFold(strings.TrimSpace(value), "true") | 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 | func appendUniqueIPs(dst, src []net.IP) []net.IP { | 396 | func appendUniqueIPs(dst, src []net.IP) []net.IP { |
| 353 | for _, ip := range src { | 397 | for _, ip := range src { |
| 354 | if ip == nil { | 398 | if ip == nil || len(dst) >= maxDiscoveredAddresses { |
| 355 | continue | 399 | continue |
| 356 | } | 400 | } |
| 357 | duplicate := false | 401 | duplicate := false |
| @@ -370,7 +414,8 @@ func appendUniqueIPs(dst, src []net.IP) []net.IP { · 12 unmodified lines |
| 370 | | 414 | |
| 371 | func appendUniqueStrings(dst, src []string) []string { | 415 | func appendUniqueStrings(dst, src []string) []string { |
| 372 | for _, value := range src { | 416 | for _, value := range src { |
| 373 | if value == "" { | 417 | value = strings.TrimSpace(value) |
| | 418 | if value == "" || len(value) > maxMetadataItemBytes || len(dst) >= maxDiscoveredMetadataItems { |
| 374 | continue | 419 | continue |
| 375 | } | 420 | } |
| 376 | duplicate := false | 421 | duplicate := false |
| @@ -387,6 +432,25 @@ func appendUniqueStrings(dst, src []string) []string { · 10 unmodified lines |
| 387 | return dst | 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 | // BrowseWithFallback discovers all AI gateways on the LAN (_ai-gateway._tcp) and prioritizes CPA instances. | 454 | // BrowseWithFallback discovers all AI gateways on the LAN (_ai-gateway._tcp) and prioritizes CPA instances. |
| 391 | func (b *ZeroconfBrowser) BrowseWithFallback(ctx context.Context) ([]DiscoveredService, error) { | 455 | func (b *ZeroconfBrowser) BrowseWithFallback(ctx context.Context) ([]DiscoveredService, error) { |
| 392 | allGateways, errMain := b.Browse(ctx, DefaultServiceType, DefaultDomain) | 456 | allGateways, errMain := b.Browse(ctx, DefaultServiceType, DefaultDomain) |
| @@ -450,13 +514,13 @@ func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { · 57 unmodified lines |
| 450 | svc.AuthRequired = true | 514 | svc.AuthRequired = true |
| 451 | } | 515 | } |
| 452 | if methods := parsed["auth_methods"]; methods != "" { | 516 | if methods := parsed["auth_methods"]; methods != "" { |
| 453 | svc.AuthMethods = strings.Split(methods, ",") | 517 | svc.AuthMethods = parseTXTList(methods) |
| 454 | } | 518 | } |
| 455 | if protos := parsed["protocols"]; protos != "" { | 519 | if protos := parsed["protocols"]; protos != "" { |
| 456 | svc.Protocols = strings.Split(protos, ",") | 520 | svc.Protocols = parseTXTList(protos) |
| 457 | } | 521 | } |
| 458 | if feats := parsed["features"]; feats != "" { | 522 | if feats := parsed["features"]; feats != "" { |
| 459 | svc.Features = strings.Split(feats, ",") | 523 | svc.Features = parseTXTList(feats) |
| 460 | } | 524 | } |
| 461 | | 525 | |
| 462 | if v, ok := parsed["api_openai"]; ok { | 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 | func filterUsableIPs(ips []net.IP) []net.IP { | 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 | for _, ip := range ips { | 547 | for _, ip := range ips { |
| 484 | if ip == nil || ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() { | 548 | if ip == nil || ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() { |
| 485 | continue | 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 | return usable | 556 | return usable |
| 490 | } | 557 | } |