From a176bad16c1a0b84bd3e1f45b1fcfc816df75840 Mon Sep 17 00:00:00 2001
From: Ebihara Yusuke <yusuke@ebihara.me>
Date: Sat, 30 Sep 2023 15:55:05 +0900
Subject: [PATCH] Implement LookupStrategy on client.Browse function

In some case, only IPv4 or IPv6 is returned.
For the usecase which requre IPv4 address, made possible to specify
stratedy.
This may resolve #27
---
 client.go  | 33 +++++++++++++++++++++++++++++++++
 service.go |  1 +
 2 files changed, 34 insertions(+)

diff --git a/client.go b/client.go
index 1ceac4bd..746d8584 100644
--- a/client.go
+++ b/client.go
@@ -81,14 +81,39 @@ func NewResolver(options ...ClientOption) (*Resolver, error) {
 	}, nil
 }
 
+// Specify strategy to get A and AAAA records.
+// TODO: support stategy to Lookup* functions.
+type LookupStrategy struct {
+	// If true, wait until the ipv4 query is responded even if ipv6 query is responded.
+	// If false, query may be responded even if ipv4 query is not responded.
+	ForceIPv4 bool
+	// If true, wait until the ipv6 query is responded even if ipv4 query is responded.
+	// If false, query may be responded even if ipv6 query is not responded.
+	ForceIPv6 bool
+}
+
+var (
+	ReturnFirst = LookupStrategy{false, false}
+	ForceIPv4   = LookupStrategy{true, false}
+	ForceIPv6   = LookupStrategy{false, true}
+	ForceBoth   = LookupStrategy{true, true}
+)
+
 // Browse for all services of a given type in a given domain.
+// This method is left for backward compatibility.
 func (r *Resolver) Browse(ctx context.Context, service, domain string, entries chan<- *ServiceEntry) error {
+	return r.BrowseWithStrategy(ctx, service, domain, LookupStrategy{false, false}, entries)
+}
+
+// Browse for all services of a given type in a given domain.
+func (r *Resolver) BrowseWithStrategy(ctx context.Context, service, domain string, strategy LookupStrategy, entries chan<- *ServiceEntry) error {
 	params := defaultParams(service)
 	if domain != "" {
 		params.Domain = domain
 	}
 	params.Entries = entries
 	params.isBrowsing = true
+	params.Strategy = strategy
 	ctx, cancel := context.WithCancel(ctx)
 	go r.c.mainloop(ctx, params)
 
@@ -285,6 +310,14 @@ func (c *client) mainloop(ctx context.Context, params *lookupParams) {
 				// If this is an DNS-SD query do not throw PTR away.
 				// It is expected to have only PTR for enumeration
 				if params.ServiceRecord.ServiceTypeName() != params.ServiceRecord.ServiceName() {
+					// wait based on lookup strategy
+					if params.Strategy.ForceIPv4 && len(e.AddrIPv4) == 0 {
+						continue
+					}
+					if params.Strategy.ForceIPv6 && len(e.AddrIPv6) == 0 {
+						continue
+					}
+
 					// Require at least one resolved IP address for ServiceEntry
 					// TODO: wait some more time as chances are high both will arrive.
 					if len(e.AddrIPv4) == 0 && len(e.AddrIPv6) == 0 {
diff --git a/service.go b/service.go
index 6253c543..12faf11a 100644
--- a/service.go
+++ b/service.go
@@ -70,6 +70,7 @@ type lookupParams struct {
 	ServiceRecord
 	Entries chan<- *ServiceEntry // Entries Channel
 
+	Strategy    LookupStrategy
 	isBrowsing  bool
 	stopProbing chan struct{}
 	once        sync.Once