From a09633e859300fbf00c079f7ec7e284d37081b7e Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Mon, 30 May 2022 15:39:24 +0200
Subject: [PATCH] Check errors of more database calls

---
 api.go  | 13 +++++++++++--
 db.go   |  4 +++-
 poll.go | 15 ++++++++++++++-
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/api.go b/api.go
index 461492a7..f68b225c 100644
--- a/api.go
+++ b/api.go
@@ -475,7 +475,16 @@ func (h *Headscale) handleMachineRefreshKey(
 		Str("machine", machine.Name).
 		Msg("We have the OldNodeKey in the database. This is a key refresh")
 	machine.NodeKey = NodePublicKeyStripPrefix(registerRequest.NodeKey)
-	h.db.Save(&machine)
+
+	if err := h.db.Save(&machine).Error; err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Failed to update machine key in the database")
+		ctx.String(http.StatusInternalServerError, "Internal server error")
+
+		return
+	}
 
 	resp.AuthURL = ""
 	resp.User = *machine.Namespace.toUser()
@@ -485,7 +494,7 @@ func (h *Headscale) handleMachineRefreshKey(
 			Caller().
 			Err(err).
 			Msg("Cannot encode message")
-		ctx.String(http.StatusInternalServerError, "Extremely sad!")
+		ctx.String(http.StatusInternalServerError, "Internal server error")
 
 		return
 	}
diff --git a/db.go b/db.go
index 9130d90e..0c1bfec6 100644
--- a/db.go
+++ b/db.go
@@ -175,7 +175,9 @@ func (h *Headscale) setValue(key string, value string) error {
 		return nil
 	}
 
-	h.db.Create(keyValue)
+	if err := h.db.Create(keyValue).Error; err != nil {
+		return fmt.Errorf("failed to create key value pair in the database: %w", err)
+	}
 
 	return nil
 }
diff --git a/poll.go b/poll.go
index f6dc1387..4ea10a28 100644
--- a/poll.go
+++ b/poll.go
@@ -126,7 +126,20 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) {
 		machine.Endpoints = req.Endpoints
 		machine.LastSeen = &now
 	}
-	h.db.Updates(machine)
+
+	if err := h.db.Updates(machine).Error; err != nil {
+		if err != nil {
+			log.Error().
+				Str("handler", "PollNetMap").
+				Str("id", ctx.Param("id")).
+				Str("machine", machine.Name).
+				Err(err).
+				Msg("Failed to persist/update machine in the database")
+			ctx.String(http.StatusInternalServerError, ":(")
+
+			return
+		}
+	}
 
 	data, err := h.getMapResponse(machineKey, req, machine)
 	if err != nil {