Move DB call of pollmap to Machine inside a function

This commit is contained in:
Kristoffer Dalby 2021-10-02 21:58:28 +01:00
parent 42913e2c37
commit 0475eb6ef7
No known key found for this signature in database
GPG Key ID: 09F62DC067465735
4 changed files with 24 additions and 8 deletions

1
api.go
View File

@ -226,6 +226,7 @@ func (h *Headscale) getMapResponse(mKey wgkey.Key, req tailcfg.MapRequest, m Mac
Msg("Cannot convert to node") Msg("Cannot convert to node")
return nil, err return nil, err
} }
peers, err := h.getPeers(m) peers, err := h.getPeers(m)
if err != nil { if err != nil {
log.Error(). log.Error().

4
go.mod
View File

@ -11,7 +11,7 @@ require (
github.com/docker/docker v20.10.8+incompatible // indirect github.com/docker/docker v20.10.8+incompatible // indirect
github.com/efekarakus/termcolor v1.0.1 github.com/efekarakus/termcolor v1.0.1
github.com/gin-gonic/gin v1.7.4 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-github v17.0.0+incompatible // indirect
github.com/google/go-querystring v1.1.0 // indirect github.com/google/go-querystring v1.1.0 // indirect
github.com/hako/durafmt v0.0.0-20210608085754-5c1018a4e16b github.com/hako/durafmt v0.0.0-20210608085754-5c1018a4e16b
@ -26,7 +26,7 @@ require (
github.com/spf13/viper v1.8.1 github.com/spf13/viper v1.8.1
github.com/stretchr/testify v1.7.0 github.com/stretchr/testify v1.7.0
github.com/tailscale/hujson v0.0.0-20210818175511-7360507a6e88 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 github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
golang.org/x/crypto v0.0.0-20210817164053-32db794688a5 golang.org/x/crypto v0.0.0-20210817164053-32db794688a5
golang.org/x/net v0.0.0-20210913180222-943fd674d43e // indirect golang.org/x/net v0.0.0-20210913180222-943fd674d43e // indirect

View File

@ -240,6 +240,15 @@ func (h *Headscale) GetMachineByID(id uint64) (*Machine, error) {
return &m, nil 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 // UpdateMachine takes a Machine struct pointer (typically already loaded from database
// and updates it with the latest data from the database. // and updates it with the latest data from the database.
func (h *Headscale) UpdateMachine(m *Machine) error { func (h *Headscale) UpdateMachine(m *Machine) error {

18
poll.go
View File

@ -51,13 +51,19 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
return return
} }
var m Machine m, err := h.GetMachineByMachineKey(mKey.HexString())
if result := h.db.Preload("Namespace").First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) { if err != nil {
log.Warn(). 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"). Str("handler", "PollNetMap").
Msgf("Ignoring request, cannot find machine with key %s", mKey.HexString()) Msgf("Failed to fetch machine from the database with Machine key: %s", mKey.HexString())
c.String(http.StatusUnauthorized, "") c.String(http.StatusInternalServerError, "")
return
} }
log.Trace(). log.Trace().
Str("handler", "PollNetMap"). Str("handler", "PollNetMap").