diff --git a/api.go b/api.go index 3add686..03f63cb 100644 --- a/api.go +++ b/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 } diff --git a/api_key.go b/api_key.go index a968b26..41c8385 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/db.go b/db.go index a24e46b..17d8323 100644 --- a/db.go +++ b/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 } diff --git a/integration_cli_test.go b/integration_cli_test.go index 62e8c63..075c38a 100644 --- a/integration_cli_test.go +++ b/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), }, diff --git a/machine.go b/machine.go index db2c63f..c227645 100644 --- a/machine.go +++ b/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 } diff --git a/poll.go b/poll.go index a74d857..239f260 100644 --- a/poll.go +++ b/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 { diff --git a/preauth_keys.go b/preauth_keys.go index 55f6222..b32ff63 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 e8de299..d062f82 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 }