Move DB call of pollmap to Machine inside a function

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().

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

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 {

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").