From 52cc3bc8eb8e50c2d4bd4e420a54d845e8872dcc Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Mon, 30 May 2022 15:31:06 +0200
Subject: [PATCH] Check all errors for db.Save

---
 api_key.go      |  5 ++++-
 machine.go      | 32 ++++++++++++++++++++++++--------
 preauth_keys.go | 14 +++++++++++---
 routes.go       |  7 ++++++-
 4 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/api_key.go b/api_key.go
index a968b260..41c83857 100644
--- a/api_key.go
+++ b/api_key.go
@@ -57,7 +57,10 @@ func (h *Headscale) CreateAPIKey(
 		Hash:       hash,
 		Expiration: expiration,
 	}
-	h.db.Save(&key)
+
+	if err := h.db.Save(&key).Error; err != nil {
+		return "", nil, fmt.Errorf("failed to save API key to database: %w", err)
+	}
 
 	return keyStr, &key, nil
 }
diff --git a/machine.go b/machine.go
index 845b6491..ef4b86a3 100644
--- a/machine.go
+++ b/machine.go
@@ -367,23 +367,30 @@ func (h *Headscale) SetTags(machine *Machine, tags []string) error {
 		return err
 	}
 	h.setLastStateChangeToNow(machine.Namespace.Name)
-	h.db.Save(machine)
+
+	if err := h.db.Save(machine).Error; err != nil {
+		return fmt.Errorf("failed to update tags for machine in the database: %w", err)
+	}
 
 	return nil
 }
 
 // ExpireMachine takes a Machine struct and sets the expire field to now.
-func (h *Headscale) ExpireMachine(machine *Machine) {
+func (h *Headscale) ExpireMachine(machine *Machine) error {
 	now := time.Now()
 	machine.Expiry = &now
 
 	h.setLastStateChangeToNow(machine.Namespace.Name)
 
-	h.db.Save(machine)
+	if err := h.db.Save(machine).Error; err != nil {
+		return fmt.Errorf("failed to expire machine in the database: %w", err)
+	}
+
+	return nil
 }
 
-// RefreshMachine takes a Machine struct and sets the expire field to now.
-func (h *Headscale) RefreshMachine(machine *Machine, expiry time.Time) {
+// RefreshMachine takes a Machine struct and sets the expire field.
+func (h *Headscale) RefreshMachine(machine *Machine, expiry time.Time) error {
 	now := time.Now()
 
 	machine.LastSuccessfulUpdate = &now
@@ -391,7 +398,11 @@ func (h *Headscale) RefreshMachine(machine *Machine, expiry time.Time) {
 
 	h.setLastStateChangeToNow(machine.Namespace.Name)
 
-	h.db.Save(machine)
+	if err := h.db.Save(machine).Error; err != nil {
+		return fmt.Errorf("failed to refresh machine (update expiration) in the database: %w", err)
+	}
+
+	return nil
 }
 
 // DeleteMachine softs deletes a Machine from the database.
@@ -761,7 +772,9 @@ func (h *Headscale) RegisterMachine(machine Machine,
 
 	machine.IPAddresses = ips
 
-	h.db.Save(&machine)
+	if err := h.db.Save(machine).Error; err != nil {
+		return nil, fmt.Errorf("failed register(save) machine in the database: %w", err)
+	}
 
 	log.Trace().
 		Caller().
@@ -821,7 +834,10 @@ func (h *Headscale) EnableRoutes(machine *Machine, routeStrs ...string) error {
 	}
 
 	machine.EnabledRoutes = newRoutes
-	h.db.Save(&machine)
+
+	if err := h.db.Save(machine).Error; err != nil {
+		return fmt.Errorf("failed enable routes for machine in the database: %w", err)
+	}
 
 	return nil
 }
diff --git a/preauth_keys.go b/preauth_keys.go
index 55f62226..b32ff636 100644
--- a/preauth_keys.go
+++ b/preauth_keys.go
@@ -4,6 +4,7 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"errors"
+	"fmt"
 	"strconv"
 	"time"
 
@@ -60,7 +61,10 @@ func (h *Headscale) CreatePreAuthKey(
 		CreatedAt:   &now,
 		Expiration:  expiration,
 	}
-	h.db.Save(&key)
+
+	if err := h.db.Save(&key).Error; err != nil {
+		return nil, fmt.Errorf("failed to create key in the database: %w", err)
+	}
 
 	return &key, nil
 }
@@ -114,9 +118,13 @@ func (h *Headscale) ExpirePreAuthKey(k *PreAuthKey) error {
 }
 
 // UsePreAuthKey marks a PreAuthKey as used.
-func (h *Headscale) UsePreAuthKey(k *PreAuthKey) {
+func (h *Headscale) UsePreAuthKey(k *PreAuthKey) error {
 	k.Used = true
-	h.db.Save(k)
+	if err := h.db.Save(k).Error; err != nil {
+		return fmt.Errorf("failed to update key used status in the database: %w", err)
+	}
+
+	return nil
 }
 
 // checkKeyValidity does the heavy lifting for validation of the PreAuthKey coming from a node
diff --git a/routes.go b/routes.go
index e8de299b..d062f827 100644
--- a/routes.go
+++ b/routes.go
@@ -1,6 +1,8 @@
 package headscale
 
 import (
+	"fmt"
+
 	"inet.af/netaddr"
 )
 
@@ -108,7 +110,10 @@ func (h *Headscale) EnableNodeRoute(
 	}
 
 	machine.EnabledRoutes = enabledRoutes
-	h.db.Save(&machine)
+
+	if err := h.db.Save(&machine).Error; err != nil {
+		return fmt.Errorf("failed to update node routes in the database: %w", err)
+	}
 
 	return nil
 }