From 5729f6513879c2734d3da8b5105f94a4d6e8ff64 Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Sat, 30 Jun 2018 23:07:31 -0700
Subject: [PATCH] Implement mDNS PC discovery

---
 app/backend/computermanager.cpp | 46 ++++++++++++++++++++++++++++++-
 app/backend/computermanager.h   | 49 +++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/app/backend/computermanager.cpp b/app/backend/computermanager.cpp
index 4e0a2c3a..b8549b9f 100644
--- a/app/backend/computermanager.cpp
+++ b/app/backend/computermanager.cpp
@@ -154,7 +154,8 @@ bool NvComputer::update(NvComputer& that)
 
 ComputerManager::ComputerManager(QObject *parent)
     : QObject(parent),
-      m_Polling(false)
+      m_Polling(false),
+      m_MdnsBrowser(nullptr)
 {
     QSettings settings;
 
@@ -186,8 +187,25 @@ void ComputerManager::startPolling()
 {
     QWriteLocker lock(&m_Lock);
 
+    if (m_Polling) {
+        return;
+    }
+
     m_Polling = true;
 
+    // Start an MDNS query for GameStream hosts
+    m_MdnsBrowser = new QMdnsEngine::Browser(&m_MdnsServer, "_nvstream._tcp.local.", &m_MdnsCache);
+    connect(m_MdnsBrowser, &QMdnsEngine::Browser::serviceAdded,
+            this, [this](const QMdnsEngine::Service& service) {
+        qDebug() << "Discovered mDNS host: " << service.hostname();
+
+        MdnsPendingComputer* pendingComputer = new MdnsPendingComputer(&m_MdnsServer, &m_MdnsCache, service);
+        connect(pendingComputer, SIGNAL(resolvedv4(MdnsPendingComputer*,QHostAddress)),
+                this, SLOT(handleMdnsServiceResolved(MdnsPendingComputer*,QHostAddress)));
+        m_PendingResolution.append(pendingComputer);
+    });
+
+    // Start polling threads for each known host
     QMapIterator<QString, NvComputer*> i(m_KnownHosts);
     while (i.hasNext()) {
         i.next();
@@ -195,6 +213,17 @@ void ComputerManager::startPolling()
     }
 }
 
+void ComputerManager::handleMdnsServiceResolved(MdnsPendingComputer* computer,
+                                                const QHostAddress& address)
+{
+    qDebug() << "Resolved " << computer->hostname() << " to " << address.toString();
+
+    addNewHost(address.toString(), true);
+
+    m_PendingResolution.removeOne(computer);
+    computer->deleteLater();
+}
+
 QVector<NvComputer*> ComputerManager::getComputers()
 {
     QReadLocker lock(&m_Lock);
@@ -229,8 +258,23 @@ void ComputerManager::stopPollingAsync()
 {
     QWriteLocker lock(&m_Lock);
 
+    if (!m_Polling) {
+        return;
+    }
+
     m_Polling = false;
 
+    // Delete machines that haven't been resolved yet
+    while (!m_PendingResolution.isEmpty()) {
+        MdnsPendingComputer* computer = m_PendingResolution.first();
+        computer->deleteLater();
+        m_PendingResolution.removeFirst();
+    }
+
+    // Delete the browser to stop discovery
+    delete m_MdnsBrowser;
+    m_MdnsBrowser = nullptr;
+
     // Interrupt all threads, but don't wait for them to terminate
     QMutableMapIterator<QString, QThread*> i(m_PollThreads);
     while (i.hasNext()) {
diff --git a/app/backend/computermanager.h b/app/backend/computermanager.h
index e7c0253d..e03700b3 100644
--- a/app/backend/computermanager.h
+++ b/app/backend/computermanager.h
@@ -1,6 +1,12 @@
 #pragma once
 #include "nvhttp.h"
 
+#include <qmdnsengine/server.h>
+#include <qmdnsengine/cache.h>
+#include <qmdnsengine/browser.h>
+#include <qmdnsengine/service.h>
+#include <qmdnsengine/resolver.h>
+
 #include <QThread>
 #include <QReadWriteLock>
 #include <QSettings>
@@ -207,6 +213,43 @@ private:
     NvComputer* m_Computer;
 };
 
+class MdnsPendingComputer : public QObject
+{
+    Q_OBJECT
+
+public:
+    explicit MdnsPendingComputer(QMdnsEngine::Server* server,
+                                 QMdnsEngine::Cache* cache,
+                                 const QMdnsEngine::Service& service)
+        : m_Hostname(service.hostname()),
+          m_Resolver(server, m_Hostname, cache)
+    {
+        connect(&m_Resolver, SIGNAL(resolved(QHostAddress)),
+                this, SLOT(handleResolved(QHostAddress)));
+    }
+
+    QString hostname()
+    {
+        return m_Hostname;
+    }
+
+private slots:
+    void handleResolved(const QHostAddress& address)
+    {
+        if (address.protocol() == QAbstractSocket::IPv4Protocol) {
+            m_Resolver.disconnect();
+            emit resolvedv4(this, address);
+        }
+    }
+
+signals:
+    void resolvedv4(MdnsPendingComputer*, const QHostAddress&);
+
+private:
+    QByteArray m_Hostname;
+    QMdnsEngine::Resolver m_Resolver;
+};
+
 class ComputerManager : public QObject
 {
     Q_OBJECT
@@ -231,6 +274,8 @@ signals:
 private slots:
     void handleComputerStateChanged(NvComputer* computer);
 
+    void handleMdnsServiceResolved(MdnsPendingComputer* computer, const QHostAddress& address);
+
 private:
     void saveHosts();
 
@@ -240,4 +285,8 @@ private:
     QReadWriteLock m_Lock;
     QMap<QString, NvComputer*> m_KnownHosts;
     QMap<QString, QThread*> m_PollThreads;
+    QMdnsEngine::Server m_MdnsServer;
+    QMdnsEngine::Browser* m_MdnsBrowser;
+    QMdnsEngine::Cache m_MdnsCache;
+    QVector<MdnsPendingComputer*> m_PendingResolution;
 };