From 0113286b4fc5832144df8a2532b032ee65f6702c Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Sat, 30 Mar 2019 16:27:56 +0100
Subject: [PATCH] fix: gracefully handling hid receiver disconnection

---
 modules/hid/hid.go       | 10 ++++++++++
 modules/hid/hid_recon.go | 14 +++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/modules/hid/hid.go b/modules/hid/hid.go
index 424267f2..49b239c5 100644
--- a/modules/hid/hid.go
+++ b/modules/hid/hid.go
@@ -204,11 +204,21 @@ func (mod *HIDRecon) Configure() error {
 	return nil
 }
 
+func (mod *HIDRecon) forceStop() error {
+	return mod.SetRunning(false, func() {
+		if mod.dongle != nil {
+			mod.dongle.Close()
+			mod.dongle = nil
+			mod.Debug("device closed")
+		}
+	})
+}
 func (mod *HIDRecon) Stop() error {
 	return mod.SetRunning(false, func() {
 		mod.waitGroup.Wait()
 		if mod.dongle != nil {
 			mod.dongle.Close()
+			mod.dongle = nil
 			mod.Debug("device closed")
 		}
 	})
diff --git a/modules/hid/hid_recon.go b/modules/hid/hid_recon.go
index b4771a03..fc5d29e1 100644
--- a/modules/hid/hid_recon.go
+++ b/modules/hid/hid_recon.go
@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/bettercap/nrf24"
+	"github.com/google/gousb"
 )
 
 func (mod *HIDRecon) doHopping() {
@@ -26,7 +27,13 @@ func (mod *HIDRecon) doHopping() {
 			mod.channel = 1
 		}
 		if err := mod.dongle.SetChannel(mod.channel); err != nil {
-			mod.Warning("error hopping on channel %d: %v", mod.channel, err)
+			if err == gousb.ErrorNoDevice {
+				mod.Error("device disconnected, stopping module")
+				mod.forceStop()
+				return
+			} else {
+				mod.Warning("error hopping on channel %d: %v", mod.channel, err)
+			}
 		} else {
 			mod.lastHop = time.Now()
 		}
@@ -107,6 +114,11 @@ func (mod *HIDRecon) Start() error {
 
 			buf, err := mod.dongle.ReceivePayload()
 			if err != nil {
+				if err == gousb.ErrorNoDevice {
+					mod.Error("device disconnected, stopping module")
+					mod.forceStop()
+					return
+				}
 				mod.Warning("error receiving payload from channel %d: %v", mod.channel, err)
 				continue
 			}