From bd9bcef3a932714b89024e4dae89d5dda48a2332 Mon Sep 17 00:00:00 2001
From: Sergey Krashevich <svk@svk.su>
Date: Sat, 29 Apr 2023 19:00:36 +0300
Subject: [PATCH 1/4] Add WebSocket functionality to HomeKit API for device
 discovery

---
 cmd/homekit/api.go | 94 +++++++++++++++++++++++++++-------------------
 www/add.html       | 56 +++++++++++++++++----------
 2 files changed, 91 insertions(+), 59 deletions(-)

diff --git a/cmd/homekit/api.go b/cmd/homekit/api.go
index a055871a..cde3aeb3 100644
--- a/cmd/homekit/api.go
+++ b/cmd/homekit/api.go
@@ -1,60 +1,39 @@
 package homekit
 
 import (
-	"encoding/json"
 	"fmt"
 	"github.com/AlexxIT/go2rtc/cmd/app/store"
 	"github.com/AlexxIT/go2rtc/cmd/streams"
 	"github.com/AlexxIT/go2rtc/pkg/hap"
 	"github.com/AlexxIT/go2rtc/pkg/hap/mdns"
+	"github.com/gorilla/websocket"
 	"net/http"
-	"net/url"
+
 	"strings"
+	"time"
 )
 
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		return true
+	},
+}
+
 func apiHandler(w http.ResponseWriter, r *http.Request) {
 	switch r.Method {
 	case "GET":
-		items := make([]any, 0)
 
-		for name, src := range store.GetDict("streams") {
-			if src := src.(string); strings.HasPrefix(src, "homekit") {
-				u, err := url.Parse(src)
-				if err != nil {
-					continue
-				}
-				device := Device{
-					Name:   name,
-					Addr:   u.Host,
-					Paired: true,
-				}
-				items = append(items, device)
-			}
-		}
-
-		for info := range mdns.GetAll() {
-			if !strings.HasSuffix(info.Name, mdns.Suffix) {
-				continue
-			}
-			name := info.Name[:len(info.Name)-len(mdns.Suffix)]
-			device := Device{
-				Name: strings.ReplaceAll(name, "\\", ""),
-				Addr: fmt.Sprintf("%s:%d", info.AddrV4, info.Port),
-			}
-			for _, field := range info.InfoFields {
-				switch field[:2] {
-				case "id":
-					device.ID = field[3:]
-				case "md":
-					device.Model = field[3:]
-				case "sf":
-					device.Paired = field[3] == '0'
-				}
-			}
-			items = append(items, device)
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			log.Error().Err(err).Caller().Send()
+			_, err = w.Write([]byte(err.Error()))
+			return
 		}
+		defer conn.Close()
 
-		_ = json.NewEncoder(w).Encode(items)
+		hkDiscoverDevices(conn)
 
 	case "POST":
 		// TODO: post params...
@@ -76,6 +55,43 @@ func apiHandler(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func hkDiscoverDevices(conn *websocket.Conn) {
+	for {
+		entries := mdns.GetAll()
+
+		for entry := range entries {
+			if !strings.HasSuffix(entry.Name, mdns.Suffix) {
+				continue
+			}
+
+			name := entry.Name[:len(entry.Name)-len(mdns.Suffix)]
+			device := Device{
+				Name: strings.ReplaceAll(name, "\\", ""),
+				Addr: fmt.Sprintf("%s:%d", entry.AddrV4, entry.Port),
+			}
+			for _, field := range entry.InfoFields {
+				switch field[:2] {
+				case "id":
+					device.ID = field[3:]
+				case "md":
+					device.Model = field[3:]
+				case "sf":
+					device.Paired = field[3] == '0'
+				}
+			}
+
+			err := conn.WriteJSON(device)
+			if err != nil {
+				log.Error().Err(err).Caller().Send()
+
+				return
+			}
+		}
+
+		time.Sleep(1 * time.Second)
+	}
+}
+
 func hkPair(deviceID, pin, name string) (err error) {
 	var conn *hap.Conn
 
diff --git a/www/add.html b/www/add.html
index 476c65c5..ed1226d9 100644
--- a/www/add.html
+++ b/www/add.html
@@ -124,29 +124,45 @@
 <script>
     document.getElementById('homekit').addEventListener('click', async ev => {
         ev.target.nextElementSibling.style.display = 'block'
-        const r = await fetch('api/homekit', {cache: 'no-cache'})
-
-        /** @type {Array<{id:string,name:string,addr:string,model:string,paired:boolean}>} */
-        const data = await r.json()
+        const homekitTable = document.getElementById("homekit-body");
+
+        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+        const socket = new WebSocket(`${protocol}//${window.location.host}/api/homekit`);
+
+        socket.onmessage = function (event) {
+            const device = JSON.parse(event.data);
+            console.log("Discovered device:", device);
+            addOrUpdateRow(device);
+        };
+
+        function addOrUpdateRow(device) {
+            const existingRow = document.getElementById(`device-${device.id}`);
+            if (existingRow) {
+                existingRow.innerHTML = createRow(device);
+            } else {
+                const newRow = document.createElement("tr");
+                newRow.id = `device-${device.id}`;
+                newRow.innerHTML = createRow(device);
+                homekitTable.appendChild(newRow);
+            }
+        }
 
-        const tbody = document.getElementById('homekit-body')
-        tbody.innerHTML =
-            data.reduce((res, item) => {
-                let commands = ''
-                if (item.id === "") {
-                    commands = `<a href="#" onclick="unpair('${item.name}')">unpair</a>`
-                } else if (item.paired === false) {
-                    commands = `<a href="#" onclick="pair('${item.id}','${item.name}')">pair</a>`
-                }
-                return res + `<tr>
-                        <td>${item.name}</td>
-                        <td>${item.addr}</td>
-                        <td>${item.model}</td>
-                        <td>${commands}</td>
-                    </tr>`
-            }, '')
+        function createRow(device) {
+            let commands = "";
+            if (device.id === "") {
+                commands = `<a href="#" onclick="unpair('${device.name}')">unpair</a>`;
+            } else if (device.paired === false) {
+                commands = `<a href="#" onclick="pair('${device.id}','${device.name}')">pair</a>`;
+            }
+            return `<td>${device.name}</td>
+                    <td>${device.addr}</td>
+                    <td>${device.model}</td>
+                    <td>${commands}</td>`;
+        }
     })
 
+    
+
     function pair(id, name) {
         const pin = document.querySelector('#pin').value
         fetch(`api/homekit?id=${id}&name=${name}&pin=${pin}`, {method: 'POST'})

From 7b522500798f0581d83b418eb645a982eb8288a3 Mon Sep 17 00:00:00 2001
From: Sergey Krashevich <svk@svk.su>
Date: Sat, 29 Apr 2023 19:12:09 +0300
Subject: [PATCH 2/4] =?UTF-8?q?the=20small=20pro=D1=91b=20fixed?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/homekit/api.go | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/cmd/homekit/api.go b/cmd/homekit/api.go
index cde3aeb3..d2352042 100644
--- a/cmd/homekit/api.go
+++ b/cmd/homekit/api.go
@@ -2,12 +2,14 @@ package homekit
 
 import (
 	"fmt"
+	"net/http"
+	"net/url"
+
 	"github.com/AlexxIT/go2rtc/cmd/app/store"
 	"github.com/AlexxIT/go2rtc/cmd/streams"
 	"github.com/AlexxIT/go2rtc/pkg/hap"
 	"github.com/AlexxIT/go2rtc/pkg/hap/mdns"
 	"github.com/gorilla/websocket"
-	"net/http"
 
 	"strings"
 	"time"
@@ -59,6 +61,26 @@ func hkDiscoverDevices(conn *websocket.Conn) {
 	for {
 		entries := mdns.GetAll()
 
+		for name, src := range store.GetDict("streams") {
+			if src := src.(string); strings.HasPrefix(src, "homekit") {
+				u, err := url.Parse(src)
+				if err != nil {
+					continue
+				}
+				device := Device{
+					Name:   name,
+					Addr:   u.Host,
+					Paired: true,
+				}
+				err = conn.WriteJSON(device)
+				if err != nil {
+					log.Error().Err(err).Caller().Send()
+
+					return
+				}
+			}
+		}
+
 		for entry := range entries {
 			if !strings.HasSuffix(entry.Name, mdns.Suffix) {
 				continue

From 6abbf723ca41c2d07d0c9eaa0c1c11d9451d037d Mon Sep 17 00:00:00 2001
From: Sergey Krashevich <svk@svk.su>
Date: Sat, 29 Apr 2023 19:33:31 +0300
Subject: [PATCH 3/4] Add activeConnections and activeConnectionsMutex to track
 the number of active connections and use them in hkDiscoverDevices to prevent
 unnecessary discovery when there are no active connections

---
 cmd/homekit/api.go | 121 +++++++++++++++++++++++++++++----------------
 1 file changed, 78 insertions(+), 43 deletions(-)

diff --git a/cmd/homekit/api.go b/cmd/homekit/api.go
index d2352042..01b27f9d 100644
--- a/cmd/homekit/api.go
+++ b/cmd/homekit/api.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"sync"
 
 	"github.com/AlexxIT/go2rtc/cmd/app/store"
 	"github.com/AlexxIT/go2rtc/cmd/streams"
@@ -23,6 +24,9 @@ var upgrader = websocket.Upgrader{
 	},
 }
 
+var activeConnections int
+var activeConnectionsMutex sync.Mutex
+
 func apiHandler(w http.ResponseWriter, r *http.Request) {
 	switch r.Method {
 	case "GET":
@@ -33,9 +37,27 @@ func apiHandler(w http.ResponseWriter, r *http.Request) {
 			_, err = w.Write([]byte(err.Error()))
 			return
 		}
-		defer conn.Close()
+		activeConnectionsMutex.Lock()
+		activeConnections++
+		activeConnectionsMutex.Unlock()
 
-		hkDiscoverDevices(conn)
+		done := make(chan struct{})
+		go hkDiscoverDevices(conn, done)
+
+		for {
+			_, _, err := conn.ReadMessage()
+			if err != nil {
+				log.Debug().Err(err).Caller().Send()
+				_, err = w.Write([]byte(err.Error()))
+				break
+			}
+		}
+
+		close(done)
+		activeConnectionsMutex.Lock()
+		activeConnections--
+		activeConnectionsMutex.Unlock()
+		conn.Close()
 
 	case "POST":
 		// TODO: post params...
@@ -57,60 +79,73 @@ func apiHandler(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func hkDiscoverDevices(conn *websocket.Conn) {
+func hkDiscoverDevices(conn *websocket.Conn, done chan struct{}) {
 	for {
-		entries := mdns.GetAll()
+		select {
+		case <-done:
+			return
+		default:
+			log.Trace().Int("active connections: ", activeConnections).Msg("[homekit] ")
+			activeConnectionsMutex.Lock()
+			if activeConnections <= 0 {
+				activeConnectionsMutex.Unlock()
+				return
+			}
+			activeConnectionsMutex.Unlock()
+
+			entries := mdns.GetAll()
+
+			for name, src := range store.GetDict("streams") {
+				if src := src.(string); strings.HasPrefix(src, "homekit") {
+					u, err := url.Parse(src)
+					if err != nil {
+						continue
+					}
+					device := Device{
+						Name:   name,
+						Addr:   u.Host,
+						Paired: true,
+					}
+					err = conn.WriteJSON(device)
+					if err != nil {
+						log.Error().Err(err).Caller().Send()
+
+						return
+					}
+				}
+			}
 
-		for name, src := range store.GetDict("streams") {
-			if src := src.(string); strings.HasPrefix(src, "homekit") {
-				u, err := url.Parse(src)
-				if err != nil {
+			for entry := range entries {
+				if !strings.HasSuffix(entry.Name, mdns.Suffix) {
 					continue
 				}
+
+				name := entry.Name[:len(entry.Name)-len(mdns.Suffix)]
 				device := Device{
-					Name:   name,
-					Addr:   u.Host,
-					Paired: true,
+					Name: strings.ReplaceAll(name, "\\", ""),
+					Addr: fmt.Sprintf("%s:%d", entry.AddrV4, entry.Port),
 				}
-				err = conn.WriteJSON(device)
-				if err != nil {
-					log.Error().Err(err).Caller().Send()
-
-					return
+				for _, field := range entry.InfoFields {
+					switch field[:2] {
+					case "id":
+						device.ID = field[3:]
+					case "md":
+						device.Model = field[3:]
+					case "sf":
+						device.Paired = field[3] == '0'
+					}
 				}
-			}
-		}
 
-		for entry := range entries {
-			if !strings.HasSuffix(entry.Name, mdns.Suffix) {
-				continue
-			}
+				err := conn.WriteJSON(device)
+				if err != nil {
+					log.Debug().Err(err).Caller().Send()
 
-			name := entry.Name[:len(entry.Name)-len(mdns.Suffix)]
-			device := Device{
-				Name: strings.ReplaceAll(name, "\\", ""),
-				Addr: fmt.Sprintf("%s:%d", entry.AddrV4, entry.Port),
-			}
-			for _, field := range entry.InfoFields {
-				switch field[:2] {
-				case "id":
-					device.ID = field[3:]
-				case "md":
-					device.Model = field[3:]
-				case "sf":
-					device.Paired = field[3] == '0'
+					return
 				}
 			}
 
-			err := conn.WriteJSON(device)
-			if err != nil {
-				log.Error().Err(err).Caller().Send()
-
-				return
-			}
+			time.Sleep(1 * time.Second)
 		}
-
-		time.Sleep(1 * time.Second)
 	}
 }
 

From c653306943897cc739c7355099024bd7f8fb20e0 Mon Sep 17 00:00:00 2001
From: Sergey Krashevich <svk@svk.su>
Date: Sat, 29 Apr 2023 19:47:52 +0300
Subject: [PATCH 4/4] Set mDNS query timeout to 5 seconds for every tenth query

...for handle really slowpoke devices
---
 cmd/homekit/api.go     | 11 +++++++++--
 pkg/hap/mdns/client.go | 12 +++++++++---
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/cmd/homekit/api.go b/cmd/homekit/api.go
index 01b27f9d..8a3503c5 100644
--- a/cmd/homekit/api.go
+++ b/cmd/homekit/api.go
@@ -80,6 +80,7 @@ func apiHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 func hkDiscoverDevices(conn *websocket.Conn, done chan struct{}) {
+	queryCounter := 0
 	for {
 		select {
 		case <-done:
@@ -93,7 +94,13 @@ func hkDiscoverDevices(conn *websocket.Conn, done chan struct{}) {
 			}
 			activeConnectionsMutex.Unlock()
 
-			entries := mdns.GetAll()
+			queryCounter++
+			timeout := time.Second
+			if queryCounter%10 == 0 {
+				timeout = 5 * time.Second
+			}
+
+			entries := mdns.GetAll(timeout)
 
 			for name, src := range store.GetDict("streams") {
 				if src := src.(string); strings.HasPrefix(src, "homekit") {
@@ -144,7 +151,7 @@ func hkDiscoverDevices(conn *websocket.Conn, done chan struct{}) {
 				}
 			}
 
-			time.Sleep(1 * time.Second)
+			time.Sleep(timeout)
 		}
 	}
 }
diff --git a/pkg/hap/mdns/client.go b/pkg/hap/mdns/client.go
index c5befa45..f0ae78a6 100644
--- a/pkg/hap/mdns/client.go
+++ b/pkg/hap/mdns/client.go
@@ -2,16 +2,22 @@ package mdns
 
 import (
 	"fmt"
-	"github.com/hashicorp/mdns"
 	"strings"
+	"time"
+
+	"github.com/hashicorp/mdns"
 )
 
 const Suffix = "._hap._tcp.local."
 
-func GetAll() chan *mdns.ServiceEntry {
+func GetAll(timeout_opt ...time.Duration) chan *mdns.ServiceEntry {
+	timeout := time.Second
+	if len(timeout_opt) > 0 {
+		timeout = timeout_opt[0]
+	}
 	entries := make(chan *mdns.ServiceEntry)
 	params := &mdns.QueryParam{
-		Service: "_hap._tcp", Entries: entries, DisableIPv6: true,
+		Service: "_hap._tcp", Entries: entries, DisableIPv6: true, Timeout: timeout,
 	}
 
 	go func() {