From 711d868b39e799913829a01ee38456b4212ca2a0 Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Wed, 16 Feb 2022 02:44:14 -0500
Subject: [PATCH] On Darwin, wait until GetAddrInfo returns all the results
 before calling the resolve callback.

---
 src/platform/Darwin/DnssdImpl.cpp | 28 ++++++++++++++++++++++++----
 src/platform/Darwin/DnssdImpl.h   |  1 +
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/platform/Darwin/DnssdImpl.cpp b/src/platform/Darwin/DnssdImpl.cpp
index f892d03a6f8411..413de03a7dbcfb 100644
--- a/src/platform/Darwin/DnssdImpl.cpp
+++ b/src/platform/Darwin/DnssdImpl.cpp
@@ -403,18 +403,35 @@ static void OnGetAddrInfo(DNSServiceRef sdRef, DNSServiceFlags flags, uint32_t i
 
     ChipLogDetail(DeviceLayer, "Mdns: %s hostname:%s", __func__, hostname);
 
+    chip::Inet::IPAddress ip;
+    CHIP_ERROR status = chip::Inet::IPAddress::GetIPAddressFromSockAddr(*address, ip);
+    if (status == CHIP_NO_ERROR)
+    {
+        sdCtx->addresses.push_back(ip);
+    }
+
+    if (flags & kDNSServiceFlagsMoreComing)
+    {
+        // Wait for that.
+        return;
+    }
+
     DnssdService service   = {};
     service.mPort          = sdCtx->port;
     service.mTextEntries   = sdCtx->textEntries.empty() ? nullptr : sdCtx->textEntries.data();
     service.mTextEntrySize = sdCtx->textEntries.empty() ? 0 : sdCtx->textEntries.size();
-    chip::Inet::IPAddress ip;
-    CHIP_ERROR status = chip::Inet::IPAddress::GetIPAddressFromSockAddr(*address, ip);
-    service.mAddress.SetValue(ip);
+    // Use the first IP we got for the DnssdService.
+    if (sdCtx->addresses.size() != 0)
+    {
+        service.mAddress.SetValue(sdCtx->addresses.front());
+        sdCtx->addresses.erase(sdCtx->addresses.begin());
+    }
     Platform::CopyString(service.mName, sdCtx->name);
     Platform::CopyString(service.mHostName, hostname);
     service.mInterface = Inet::InterfaceId(sdCtx->interfaceId);
 
-    sdCtx->callback(sdCtx->context, &service, Span<Inet::IPAddress>(), status);
+    // TODO: Does it really make sense to pass in the status from our last "get the IP address" operation?
+    sdCtx->callback(sdCtx->context, &service, Span<Inet::IPAddress>(sdCtx->addresses.data(), sdCtx->addresses.size()), status);
     MdnsContexts::GetInstance().Remove(sdCtx);
 }
 
@@ -512,6 +529,9 @@ static void OnResolve(DNSServiceRef sdRef, DNSServiceFlags flags, uint32_t inter
 
     GetAddrInfo(sdCtx->context, sdCtx->callback, interfaceId, sdCtx->addressType, sdCtx->name, hostname, ntohs(port), txtLen,
                 txtRecord);
+
+    // TODO: If flags & kDNSServiceFlagsMoreComing should we keep waiting to see
+    // what else we resolve instead of calling Remove() here?
     MdnsContexts::GetInstance().Remove(sdCtx);
 }
 
diff --git a/src/platform/Darwin/DnssdImpl.h b/src/platform/Darwin/DnssdImpl.h
index a0f1bdc4be78cb..0c002ba9de9dd6 100644
--- a/src/platform/Darwin/DnssdImpl.h
+++ b/src/platform/Darwin/DnssdImpl.h
@@ -91,6 +91,7 @@ struct ResolveContext : public GenericContext
 struct GetAddrInfoContext : public GenericContext
 {
     DnssdResolveCallback callback;
+    std::vector<Inet::IPAddress> addresses;
     std::vector<TextEntry> textEntries;
     char name[Common::kInstanceNameMaxLength + 1];
     uint32_t interfaceId;