Network process should be able to resolve its own registered mDNS names

https://bugs.webkit.org/show_bug.cgi?id=271368
rdar://112834753

Reviewed by Chris Dumez.

Store the registered mdns names in a HashSet.
In case of a mDNS resolution, we check the hashset and if there is a match, we use the current local address directly.
This ensures that mDNS resolution with Safari works even if mDNS registration fails.
We add a mechanism to remove the mdns names from the HashSet whenever the document goes away, to keep the same mDNS lifetime as before.

* Source/WebKit/NetworkProcess/webrtc/NetworkMDNSRegister.cpp:
(WebKit::NetworkMDNSRegister::unregisterMDNSNames):
(WebKit::NetworkMDNSRegister::hasRegisteredName):
(WebKit::NetworkMDNSRegister::registerMDNSName):
* Source/WebKit/NetworkProcess/webrtc/NetworkMDNSRegister.h:
* Source/WebKit/NetworkProcess/webrtc/NetworkRTCProvider.cpp:
(WebKit::NetworkRTCProvider::createResolver):
* Source/WebKit/Shared/RTCNetwork.h:
(WebKit::RTC::Network::IPAddress::isUnspecified const):

Canonical link: https://commits.webkit.org/276556@main

Source/WebKit/NetworkProcess/webrtc/NetworkMDNSRegister.cpp

@@ -55,6 +55,13 @@ struct NetworkMDNSRegister::DNSServiceDeallocator {
 void NetworkMDNSRegister::unregisterMDNSNames(WebCore::ScriptExecutionContextIdentifier documentIdentifier)
 {
     m_services.remove(documentIdentifier);
+    for (auto& name : m_perDocumentRegisteredNames.take(documentIdentifier))
+        m_registeredNames.remove(name);
+}
+
+bool NetworkMDNSRegister::hasRegisteredName(const String& name) const
+{
+    return m_registeredNames.contains(name);
 }
 
 struct PendingRegistrationRequest {
@@ -108,6 +115,11 @@ void NetworkMDNSRegister::registerMDNSName(WebCore::ScriptExecutionContextIdenti
 {
     auto name = makeString(WTF::UUID::createVersion4(), ".local"_s);
 
+    m_registeredNames.add(name);
+    m_perDocumentRegisteredNames.ensure(documentIdentifier, [] {
+        return Vector<String>();
+    }).iterator->value.append(name);
+
     DNSServiceRef service;
     auto iterator = m_services.find(documentIdentifier);
     if (iterator == m_services.end()) {

Source/WebKit/NetworkProcess/webrtc/NetworkMDNSRegister.h

@@ -34,6 +34,7 @@
 #include <wtf/Expected.h>
 #include <wtf/Forward.h>
 #include <wtf/HashMap.h>
+#include <wtf/HashSet.h>
 
 #if PLATFORM(COCOA) && defined __has_include && __has_include(<dns_sd.h>)
 #define ENABLE_MDNS 1
@@ -73,13 +74,19 @@ class NetworkMDNSRegister {
     void closeAndForgetService(DNSServiceRef);
 #endif
 
+    bool hasRegisteredName(const String&) const;
+
 private:
     void unregisterMDNSNames(WebCore::ScriptExecutionContextIdentifier);
     void registerMDNSName(WebCore::ScriptExecutionContextIdentifier, const String& ipAddress, CompletionHandler<void(const String&, std::optional<WebCore::MDNSRegisterError>)>&&);
 
     PAL::SessionID sessionID() const;
 
     WeakRef<NetworkConnectionToWebProcess> m_connection;
+    HashSet<String> m_registeredNames;
+
+    HashMap<WebCore::ScriptExecutionContextIdentifier, Vector<String>> m_perDocumentRegisteredNames;
+
 #if ENABLE_MDNS
     struct DNSServiceDeallocator;
     HashMap<WebCore::ScriptExecutionContextIdentifier, std::unique_ptr<_DNSServiceRef_t, DNSServiceDeallocator>> m_services;

Source/WebKit/NetworkProcess/webrtc/NetworkRTCMonitor.cpp

@@ -51,6 +51,9 @@ class NetworkManagerWrapper final : public sigslot::has_slots<> {
     void addListener(NetworkRTCMonitor&);
     void removeListener(NetworkRTCMonitor&);
 
+    const RTCNetwork::IPAddress& ipv4() const { return m_ipv4; }
+    const RTCNetwork::IPAddress& ipv6()  const { return m_ipv6; }
+
 private:
     void onNetworksChanged();
 
@@ -138,6 +141,16 @@ NetworkRTCMonitor::~NetworkRTCMonitor()
     ASSERT(!m_manager);
 }
 
+const RTCNetwork::IPAddress& NetworkRTCMonitor::ipv4() const
+{
+    return networkManager().ipv4();
+}
+
+const RTCNetwork::IPAddress& NetworkRTCMonitor::ipv6()  const
+{
+    return networkManager().ipv6();
+}
+
 void NetworkRTCMonitor::startUpdatingIfNeeded()
 {
     RTC_RELEASE_LOG("startUpdatingIfNeeded m_isStarted=%d", m_isStarted);

Source/WebKit/NetworkProcess/webrtc/NetworkRTCMonitor.h

@@ -60,6 +60,9 @@ class NetworkRTCMonitor final : public CanMakeWeakPtr<NetworkRTCMonitor> {
 
     void onNetworksChanged(const Vector<RTCNetwork>&, const RTCNetwork::IPAddress&, const RTCNetwork::IPAddress&);
 
+    const RTCNetwork::IPAddress& ipv4() const;
+    const RTCNetwork::IPAddress& ipv6()  const;
+
 private:
     void startUpdatingIfNeeded();
 

Source/WebKit/NetworkProcess/webrtc/NetworkRTCProvider.cpp

@@ -305,6 +305,18 @@ void NetworkRTCProvider::createResolver(LibWebRTCResolverIdentifier identifier,
         });
         return;
     }
+
+    RefPtr connection = m_connection.get();
+    if (connection && connection->mdnsRegister().hasRegisteredName(address)) {
+        Vector<WebKit::RTC::Network::IPAddress> ipAddresses;
+        if (!m_rtcMonitor.ipv4().isUnspecified())
+            ipAddresses.append(m_rtcMonitor.ipv4());
+        if (!m_rtcMonitor.ipv6().isUnspecified())
+            ipAddresses.append(m_rtcMonitor.ipv6());
+        connection->connection().send(Messages::WebRTCResolver::SetResolvedAddress(ipAddresses), identifier);
+        return;
+    }
+
     WebCore::DNSCompletionHandler completionHandler = [connection = m_connection, identifier](auto&& result) {
         ASSERT(isMainRunLoop());
         if (!connection)

Source/WebKit/Shared/RTCNetwork.h

@@ -53,6 +53,8 @@ struct IPAddress {
 
     rtc::IPAddress rtcAddress() const;
 
+    bool isUnspecified() const { return std::holds_alternative<UnspecifiedFamily>(value); }
+
     std::variant<UnspecifiedFamily, uint32_t, std::array<uint32_t, 4>> value;
 };