lywsvip/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-10-18 21:45:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(discovery): bound browse resource usage

Browse Source
This commit is contained in:
sususu
2026-09-14 13:08:17 +08:00
parent c1b7c91f2f
commit f5c19d25bb
2 changed files with 167 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
internal/discovery/discovery_test.go
View File

@@ -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)

131
internal/discovery/zeroconf.go
View File

@@ -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
}