diff --git a/internal/discovery/discovery_test.go b/internal/discovery/discovery_test.go index 565de5ef6..2b5ea7147 100644 --- a/internal/discovery/discovery_test.go +++ b/internal/discovery/discovery_test.go @@ -2,6 +2,7 @@ package discovery import ( "context" + "fmt" "net" "os" "strings" @@ -383,6 +384,73 @@ func TestMergeDiscoveredService(t *testing.T) { } } +func TestDiscoveryMergeLimits(t *testing.T) { + dst := DiscoveredService{ + IPv4: []net.IP{net.ParseIP("192.0.2.1")}, + RawTXT: map[string]string{"auth_required": "true"}, + } + var srcIPs []net.IP + for i := 0; i < maxDiscoveredAddresses*2; i++ { + srcIPs = append(srcIPs, net.IPv4(198, 18, byte(i/256), byte(i%256))) + } + metadata := make([]string, 0, maxDiscoveredMetadataItems*2) + for i := 0; i < maxDiscoveredMetadataItems*2; i++ { + metadata = append(metadata, fmt.Sprintf("feature-%02d", i)) + } + rawTXT := make(map[string]string, maxBrowseTXTRecords*2) + for i := 0; i < maxBrowseTXTRecords*2; i++ { + rawTXT[fmt.Sprintf("key-%03d", i)] = strings.Repeat("v", maxTXTRecordBytes) + } + + mergeDiscoveredService(&dst, DiscoveredService{ + IPv4: srcIPs, + RawTXT: rawTXT, + AuthMethods: metadata, + Protocols: metadata, + Features: metadata, + }) + + if len(dst.IPv4) != maxDiscoveredAddresses { + t.Fatalf("merged IPv4 count = %d, want %d", len(dst.IPv4), maxDiscoveredAddresses) + } + if len(dst.AuthMethods) != maxDiscoveredMetadataItems || len(dst.Protocols) != maxDiscoveredMetadataItems || len(dst.Features) != maxDiscoveredMetadataItems { + t.Fatalf("merged metadata counts = auth %d, protocols %d, features %d; want %d each", len(dst.AuthMethods), len(dst.Protocols), len(dst.Features), maxDiscoveredMetadataItems) + } + if len(dst.RawTXT) > maxBrowseTXTRecords || rawTXTMapBytes(dst.RawTXT) > maxBrowseTXTBytes { + t.Fatalf("merged TXT size = %d records/%d bytes, limits are %d records/%d bytes", len(dst.RawTXT), rawTXTMapBytes(dst.RawTXT), maxBrowseTXTRecords, maxBrowseTXTBytes) + } +} + +func TestEntryToDiscoveredLimitsTXTLists(t *testing.T) { + var methods strings.Builder + for i := 0; i < maxDiscoveredMetadataItems*2; i++ { + if i > 0 { + methods.WriteByte(',') + } + methods.WriteString(fmt.Sprintf("method-%02d", i)) + } + addresses := make([]net.IP, 0, maxDiscoveredAddresses*2) + for i := 0; i < maxDiscoveredAddresses*2; i++ { + addresses = append(addresses, net.IPv4(198, 18, byte(i/256), byte(i%256))) + } + svc := entryToDiscovered(&zeroconf.ServiceEntry{ + ServiceRecord: zeroconf.ServiceRecord{ + Instance: "node", + Service: DefaultServiceType, + Domain: DefaultDomain, + }, + Port: 8317, + AddrIPv4: addresses, + Text: []string{"auth_methods=" + methods.String()}, + }) + if len(svc.AuthMethods) != maxDiscoveredMetadataItems { + t.Fatalf("auth method count = %d, want %d", len(svc.AuthMethods), maxDiscoveredMetadataItems) + } + if len(svc.IPv4) != maxDiscoveredAddresses { + t.Fatalf("IPv4 count = %d, want %d", len(svc.IPv4), maxDiscoveredAddresses) + } +} + func TestFilterInterfaces_RealMachine(t *testing.T) { // Testing real interface filter on current test environment ifaces, err := FilterInterfaces(nil, nil) diff --git a/internal/discovery/zeroconf.go b/internal/discovery/zeroconf.go index 70b710ba2..2a39ae495 100644 --- a/internal/discovery/zeroconf.go +++ b/internal/discovery/zeroconf.go @@ -216,9 +216,13 @@ func NewZeroconfBrowser(ifaces ...net.Interface) *ZeroconfBrowser { } const ( - maxDiscoveredServices = 256 - maxBrowseTXTRecords = 64 - maxBrowseTXTBytes = 16 * 1024 + maxDiscoveredServices = 256 + maxBrowseEntries = 256 + maxBrowseTXTRecords = 64 + maxBrowseTXTBytes = 16 * 1024 + maxDiscoveredAddresses = 32 + maxDiscoveredMetadataItems = 32 + maxMetadataItemBytes = 64 ) func browseEntryWithinLimits(entry *zeroconf.ServiceEntry) bool { @@ -251,40 +255,48 @@ func (b *ZeroconfBrowser) Browse(ctx context.Context, serviceType, domain string } entries := make(chan *zeroconf.ServiceEntry, 32) + browseCtx, cancelBrowse := context.WithCancel(ctx) + defer cancelBrowse() + var discovered []DiscoveredService var mu sync.Mutex seen := make(map[string]int) - // Collect entries in background until channel is closed by zeroconf's params.done() + // Stop the underlying zeroconf client after a bounded number of entries. This + // is separate from the caller's context so a LAN flood cannot grow the + // dependency's internal sentEntries cache for the full browse lifetime. doneCh := make(chan struct{}) go func() { defer close(doneCh) + entriesSeen := 0 for entry := range entries { - if !browseEntryWithinLimits(entry) { + if entriesSeen >= maxBrowseEntries { continue } - svc := entryToDiscovered(entry) - if svc.Port == 0 || (len(svc.IPv4) == 0 && len(svc.IPv6) == 0) { - continue + entriesSeen++ + + if browseEntryWithinLimits(entry) { + svc := entryToDiscovered(entry) + if svc.Port != 0 && (len(svc.IPv4) != 0 || len(svc.IPv6) != 0) { + key := discoveredServiceKey(svc) + mu.Lock() + if index, ok := seen[key]; ok { + mergeDiscoveredService(&discovered[index], svc) + } else if len(discovered) < maxDiscoveredServices { + seen[key] = len(discovered) + discovered = append(discovered, svc) + } + mu.Unlock() + } } - key := discoveredServiceKey(svc) - mu.Lock() - if index, ok := seen[key]; ok { - mergeDiscoveredService(&discovered[index], svc) - mu.Unlock() - continue + + if entriesSeen == maxBrowseEntries { + cancelBrowse() } - if len(discovered) >= maxDiscoveredServices { - mu.Unlock() - continue - } - seen[key] = len(discovered) - discovered = append(discovered, svc) - mu.Unlock() } }() - errBrowse := zeroconf.Browse(ctx, serviceType, domain, entries, b.options...) + errBrowse := zeroconf.Browse(browseCtx, serviceType, domain, entries, b.options...) if errBrowse != nil { // zeroconf may already have closed entries after a runtime error. closeBrowseEntries(entries) @@ -332,9 +344,7 @@ func mergeDiscoveredService(dst *DiscoveredService, src DiscoveredService) { if dst.RawTXT == nil { dst.RawTXT = make(map[string]string) } - for key, value := range src.RawTXT { - dst.RawTXT[key] = value - } + mergeRawTXTRecords(dst.RawTXT, src.RawTXT) if value, ok := dst.RawTXT["auth_required"]; ok { dst.AuthRequired = strings.EqualFold(strings.TrimSpace(value), "true") } @@ -349,9 +359,43 @@ func mergeDiscoveredService(dst *DiscoveredService, src DiscoveredService) { } } +func mergeRawTXTRecords(dst, src map[string]string) { + totalBytes := rawTXTMapBytes(dst) + for key, value := range src { + if key == "" { + continue + } + replacedBytes := 0 + _, exists := dst[key] + if exists { + replacedBytes = rawTXTRecordBytes(key, dst[key]) + } else if len(dst) >= maxBrowseTXTRecords { + continue + } + nextBytes := totalBytes - replacedBytes + rawTXTRecordBytes(key, value) + if nextBytes > maxBrowseTXTBytes { + continue + } + dst[key] = value + totalBytes = nextBytes + } +} + +func rawTXTMapBytes(records map[string]string) int { + totalBytes := 0 + for key, value := range records { + totalBytes += rawTXTRecordBytes(key, value) + } + return totalBytes +} + +func rawTXTRecordBytes(key, value string) int { + return len(key) + len(value) + 1 +} + func appendUniqueIPs(dst, src []net.IP) []net.IP { for _, ip := range src { - if ip == nil { + if ip == nil || len(dst) >= maxDiscoveredAddresses { continue } duplicate := false @@ -370,7 +414,8 @@ func appendUniqueIPs(dst, src []net.IP) []net.IP { func appendUniqueStrings(dst, src []string) []string { for _, value := range src { - if value == "" { + value = strings.TrimSpace(value) + if value == "" || len(value) > maxMetadataItemBytes || len(dst) >= maxDiscoveredMetadataItems { continue } duplicate := false @@ -387,6 +432,25 @@ func appendUniqueStrings(dst, src []string) []string { return dst } +func parseTXTList(value string) []string { + var result []string + for len(value) > 0 && len(result) < maxDiscoveredMetadataItems { + item := value + if before, after, ok := strings.Cut(value, ","); ok { + item = before + value = after + } else { + value = "" + } + item = strings.TrimSpace(item) + if item == "" || len(item) > maxMetadataItemBytes { + continue + } + result = appendUniqueStrings(result, []string{item}) + } + return result +} + // BrowseWithFallback discovers all AI gateways on the LAN (_ai-gateway._tcp) and prioritizes CPA instances. func (b *ZeroconfBrowser) BrowseWithFallback(ctx context.Context) ([]DiscoveredService, error) { allGateways, errMain := b.Browse(ctx, DefaultServiceType, DefaultDomain) @@ -450,13 +514,13 @@ func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { svc.AuthRequired = true } if methods := parsed["auth_methods"]; methods != "" { - svc.AuthMethods = strings.Split(methods, ",") + svc.AuthMethods = parseTXTList(methods) } if protos := parsed["protocols"]; protos != "" { - svc.Protocols = strings.Split(protos, ",") + svc.Protocols = parseTXTList(protos) } if feats := parsed["features"]; feats != "" { - svc.Features = strings.Split(feats, ",") + svc.Features = parseTXTList(feats) } if v, ok := parsed["api_openai"]; ok { @@ -479,12 +543,15 @@ func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { } func filterUsableIPs(ips []net.IP) []net.IP { - usable := make([]net.IP, 0, len(ips)) + usable := make([]net.IP, 0, min(len(ips), maxDiscoveredAddresses)) for _, ip := range ips { if ip == nil || ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() { continue } - usable = append(usable, ip) + usable = append(usable, append(net.IP(nil), ip...)) + if len(usable) >= maxDiscoveredAddresses { + break + } } return usable }