From 0475eb6ef7ffc70fa544c9f31f1e697d4235323e Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sat, 2 Oct 2021 21:58:28 +0100
Subject: [PATCH] Move DB call of pollmap to Machine inside a function

---
 api.go     |  1 +
 go.mod     |  4 ++--
 machine.go |  9 +++++++++
 poll.go    | 18 ++++++++++++------
 4 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/api.go b/api.go
index e2a5618..304f9a8 100644
--- a/api.go
+++ b/api.go
@@ -226,6 +226,7 @@ func (h *Headscale) getMapResponse(mKey wgkey.Key, req tailcfg.MapRequest, m Mac
 			Msg("Cannot convert to node")
 		return nil, err
 	}
+
 	peers, err := h.getPeers(m)
 	if err != nil {
 		log.Error().
diff --git a/go.mod b/go.mod
index 51acff8..3a49c0f 100644
--- a/go.mod
+++ b/go.mod
@@ -11,7 +11,7 @@ require (
 	github.com/docker/docker v20.10.8+incompatible // indirect
 	github.com/efekarakus/termcolor v1.0.1
 	github.com/gin-gonic/gin v1.7.4
-	github.com/gofrs/uuid v4.0.0+incompatible // indirect
+	github.com/gofrs/uuid v4.0.0+incompatible
 	github.com/google/go-github v17.0.0+incompatible // indirect
 	github.com/google/go-querystring v1.1.0 // indirect
 	github.com/hako/durafmt v0.0.0-20210608085754-5c1018a4e16b
@@ -26,7 +26,7 @@ require (
 	github.com/spf13/viper v1.8.1
 	github.com/stretchr/testify v1.7.0
 	github.com/tailscale/hujson v0.0.0-20210818175511-7360507a6e88
-	github.com/tcnksm/go-latest v0.0.0-20170313132115-e3007ae9052e // indirect
+	github.com/tcnksm/go-latest v0.0.0-20170313132115-e3007ae9052e
 	github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
 	golang.org/x/crypto v0.0.0-20210817164053-32db794688a5
 	golang.org/x/net v0.0.0-20210913180222-943fd674d43e // indirect
diff --git a/machine.go b/machine.go
index 1d4939c..7e266b6 100644
--- a/machine.go
+++ b/machine.go
@@ -240,6 +240,15 @@ func (h *Headscale) GetMachineByID(id uint64) (*Machine, error) {
 	return &m, nil
 }
 
+// GetMachineByMachineKey finds a Machine by ID and returns the Machine struct
+func (h *Headscale) GetMachineByMachineKey(mKey string) (*Machine, error) {
+	m := Machine{}
+	if result := h.db.Preload("Namespace").First(&m, "machine_key = ?", mKey); result.Error != nil {
+		return nil, result.Error
+	}
+	return &m, nil
+}
+
 // UpdateMachine takes a Machine struct pointer (typically already loaded from database
 // and updates it with the latest data from the database.
 func (h *Headscale) UpdateMachine(m *Machine) error {
diff --git a/poll.go b/poll.go
index 60bfa9e..ca4f676 100644
--- a/poll.go
+++ b/poll.go
@@ -51,13 +51,19 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
 		return
 	}
 
-	var m Machine
-	if result := h.db.Preload("Namespace").First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
-		log.Warn().
+	m, err := h.GetMachineByMachineKey(mKey.HexString())
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			log.Warn().
+				Str("handler", "PollNetMap").
+				Msgf("Ignoring request, cannot find machine with key %s", mKey.HexString())
+			c.String(http.StatusUnauthorized, "")
+			return
+		}
+		log.Error().
 			Str("handler", "PollNetMap").
-			Msgf("Ignoring request, cannot find machine with key %s", mKey.HexString())
-		c.String(http.StatusUnauthorized, "")
-		return
+			Msgf("Failed to fetch machine from the database with Machine key: %s", mKey.HexString())
+		c.String(http.StatusInternalServerError, "")
 	}
 	log.Trace().
 		Str("handler", "PollNetMap").