Merge pull request #597 from kradalby/db-error-handling

api.go

@@ -474,7 +474,16 @@ func (h *Headscale) handleMachineRefreshKey(
 		Str("machine", machine.Hostname).
 		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()
@@ -484,7 +493,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
 	}

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
 }

db.go

@@ -214,7 +214,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
 }

integration_cli_test.go

@@ -60,7 +60,7 @@ func (s *IntegrationCLITestSuite) SetupTest() {
 	}
 
 	headscaleOptions := &dockertest.RunOptions{
-		Name: "headscale",
+		Name: "headscale-cli",
 		Mounts: []string{
 			fmt.Sprintf("%s/integration_test/etc:/etc/headscale", currentPath),
 		},

machine.go

@@ -378,19 +378,26 @@ 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
 }
 
 // RenameMachine takes a Machine struct and a new GivenName for the machines
@@ -413,13 +420,15 @@ func (h *Headscale) RenameMachine(machine *Machine, newName string) error {
 
 	h.setLastStateChangeToNow(machine.Namespace.Name)
 
-	h.db.Save(machine)
+	if err := h.db.Save(machine).Error; err != nil {
+		return fmt.Errorf("failed to rename 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) {
+func (h *Headscale) RefreshMachine(machine *Machine, expiry time.Time) error {
 	now := time.Now()
 
 	machine.LastSuccessfulUpdate = &now
@@ -427,7 +436,14 @@ 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.
@@ -793,7 +809,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().
@@ -853,7 +871,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
 }

poll.go

@@ -115,7 +115,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.Hostname).
+				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 {

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

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
 }