From 53e5c05b0a60a9cca9e5805e4c16325e45a6ee8d Mon Sep 17 00:00:00 2001 From: Juan Font Alonso Date: Mon, 20 Jun 2022 12:30:51 +0200 Subject: [PATCH] Remove gin from the poll handlers --- poll.go | 145 +++++++++++++++++++++++++++++--------------------------- 1 file changed, 74 insertions(+), 71 deletions(-) diff --git a/poll.go b/poll.go index 239f260b..1d215089 100644 --- a/poll.go +++ b/poll.go @@ -2,13 +2,14 @@ package headscale import ( "context" + "encoding/json" "errors" "fmt" "io" "net/http" "time" - "github.com/gin-gonic/gin" + "github.com/gorilla/mux" "github.com/rs/zerolog/log" "gorm.io/gorm" "tailscale.com/tailcfg" @@ -33,13 +34,25 @@ const machineNameContextKey = contextKey("machineName") // only after their first request (marked with the ReadOnly field). // // At this moment the updates are sent in a quite horrendous way, but they kinda work. -func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { +func (h *Headscale) PollNetMapHandler( + w http.ResponseWriter, + r *http.Request, +) { + vars := mux.Vars(r) + machineKeyStr, ok := vars["mkey"] + if !ok || machineKeyStr == "" { + log.Error(). + Str("handler", "PollNetMap"). + Msg("No machine key in request") + http.Error(w, "No machine key in request", http.StatusBadRequest) + + return + } log.Trace(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Msg("PollNetMapHandler called") - body, _ := io.ReadAll(ctx.Request.Body) - machineKeyStr := ctx.Param("id") + body, _ := io.ReadAll(r.Body) var machineKey key.MachinePublic err := machineKey.UnmarshalText([]byte(MachinePublicKeyEnsurePrefix(machineKeyStr))) @@ -48,7 +61,8 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { Str("handler", "PollNetMap"). Err(err). Msg("Cannot parse client key") - ctx.String(http.StatusBadRequest, "") + + http.Error(w, "Cannot parse client key", http.StatusBadRequest) return } @@ -59,7 +73,7 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { Str("handler", "PollNetMap"). Err(err). Msg("Cannot decode message") - ctx.String(http.StatusBadRequest, "") + http.Error(w, "Cannot decode message", http.StatusBadRequest) return } @@ -70,20 +84,21 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { log.Warn(). Str("handler", "PollNetMap"). Msgf("Ignoring request, cannot find machine with key %s", machineKey.String()) - ctx.String(http.StatusUnauthorized, "") + + http.Error(w, "", http.StatusUnauthorized) return } log.Error(). Str("handler", "PollNetMap"). Msgf("Failed to fetch machine from the database with Machine key: %s", machineKey.String()) - ctx.String(http.StatusInternalServerError, "") + http.Error(w, "", http.StatusInternalServerError) return } log.Trace(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Str("machine", machine.Hostname). Msg("Found machine in database") @@ -120,11 +135,11 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { if err != nil { log.Error(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Str("machine", machine.Hostname). Err(err). Msg("Failed to persist/update machine in the database") - ctx.String(http.StatusInternalServerError, ":(") + http.Error(w, "", http.StatusInternalServerError) return } @@ -134,11 +149,11 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { if err != nil { log.Error(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Str("machine", machine.Hostname). Err(err). Msg("Failed to get Map response") - ctx.String(http.StatusInternalServerError, ":(") + http.Error(w, "", http.StatusInternalServerError) return } @@ -150,7 +165,7 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { // Details on the protocol can be found in https://github.com/tailscale/tailscale/blob/main/tailcfg/tailcfg.go#L696 log.Debug(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Str("machine", machine.Hostname). Bool("readOnly", req.ReadOnly). Bool("omitPeers", req.OmitPeers). @@ -162,7 +177,10 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { Str("handler", "PollNetMap"). Str("machine", machine.Hostname). Msg("Client is starting up. Probably interested in a DERP map") - ctx.Data(http.StatusOK, "application/json; charset=utf-8", data) + + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(data) return } @@ -177,7 +195,7 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { // Only create update channel if it has not been created log.Trace(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Str("machine", machine.Hostname). Msg("Loading or creating update channel") @@ -194,8 +212,9 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { Str("handler", "PollNetMap"). Str("machine", machine.Hostname). Msg("Client sent endpoint update and is ok with a response without peer list") - ctx.Data(http.StatusOK, "application/json; charset=utf-8", data) - + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(data) // It sounds like we should update the nodes when we have received a endpoint update // even tho the comments in the tailscale code dont explicitly say so. updateRequestsFromNode.WithLabelValues(machine.Namespace.Name, machine.Hostname, "endpoint-update"). @@ -208,7 +227,7 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { Str("handler", "PollNetMap"). Str("machine", machine.Hostname). Msg("Ignoring request, don't know how to handle it") - ctx.String(http.StatusBadRequest, "") + http.Error(w, "", http.StatusBadRequest) return } @@ -232,7 +251,8 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { updateChan <- struct{}{} h.PollNetMapStream( - ctx, + w, + r, machine, req, machineKey, @@ -242,7 +262,7 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { ) log.Trace(). Str("handler", "PollNetMap"). - Str("id", ctx.Param("id")). + Str("id", machineKeyStr). Str("machine", machine.Hostname). Msg("Finished stream, closing PollNetMap session") } @@ -251,7 +271,8 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) { // stream logic, ensuring we communicate updates and data // to the connected clients. func (h *Headscale) PollNetMapStream( - ctx *gin.Context, + w http.ResponseWriter, + r *http.Request, machine *Machine, mapRequest tailcfg.MapRequest, machineKey key.MachinePublic, @@ -259,41 +280,21 @@ func (h *Headscale) PollNetMapStream( keepAliveChan chan []byte, updateChan chan struct{}, ) { - { - machine, err := h.GetMachineByMachineKey(machineKey) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - log.Warn(). - Str("handler", "PollNetMap"). - Msgf("Ignoring request, cannot find machine with key %s", machineKey.String()) - ctx.String(http.StatusUnauthorized, "") + ctx := context.WithValue(context.Background(), machineNameContextKey, machine.Hostname) - return - } - log.Error(). - Str("handler", "PollNetMap"). - Msgf("Failed to fetch machine from the database with Machine key: %s", machineKey.String()) - ctx.String(http.StatusInternalServerError, "") + ctx, cancel := context.WithCancel(ctx) + defer cancel() - return - } + go h.scheduledPollWorker( + ctx, + updateChan, + keepAliveChan, + machineKey, + mapRequest, + machine, + ) - ctx := context.WithValue(ctx.Request.Context(), machineNameContextKey, machine.Hostname) - - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - go h.scheduledPollWorker( - ctx, - updateChan, - keepAliveChan, - machineKey, - mapRequest, - machine, - ) - } - - ctx.Stream(func(writer io.Writer) bool { + for { log.Trace(). Str("handler", "PollNetMapStream"). Str("machine", machine.Hostname). @@ -312,7 +313,7 @@ func (h *Headscale) PollNetMapStream( Str("channel", "pollData"). Int("bytes", len(data)). Msg("Sending data received via pollData channel") - _, err := writer.Write(data) + _, err := w.Write(data) if err != nil { log.Error(). Str("handler", "PollNetMapStream"). @@ -321,7 +322,7 @@ func (h *Headscale) PollNetMapStream( Err(err). Msg("Cannot write data") - return false + break } log.Trace(). Str("handler", "PollNetMapStream"). @@ -343,7 +344,7 @@ func (h *Headscale) PollNetMapStream( // client has been removed from database // since the stream opened, terminate connection. - return false + break } now := time.Now().UTC() machine.LastSeen = &now @@ -369,7 +370,7 @@ func (h *Headscale) PollNetMapStream( Msg("Machine entry in database updated successfully after sending pollData") } - return true + break case data := <-keepAliveChan: log.Trace(). @@ -378,7 +379,7 @@ func (h *Headscale) PollNetMapStream( Str("channel", "keepAlive"). Int("bytes", len(data)). Msg("Sending keep alive message") - _, err := writer.Write(data) + _, err := w.Write(data) if err != nil { log.Error(). Str("handler", "PollNetMapStream"). @@ -387,7 +388,7 @@ func (h *Headscale) PollNetMapStream( Err(err). Msg("Cannot write keep alive message") - return false + break } log.Trace(). Str("handler", "PollNetMapStream"). @@ -409,7 +410,7 @@ func (h *Headscale) PollNetMapStream( // client has been removed from database // since the stream opened, terminate connection. - return false + break } now := time.Now().UTC() machine.LastSeen = &now @@ -430,7 +431,7 @@ func (h *Headscale) PollNetMapStream( Msg("Machine updated successfully after sending keep alive") } - return true + break case <-updateChan: log.Trace(). @@ -460,7 +461,7 @@ func (h *Headscale) PollNetMapStream( Err(err). Msg("Could not get the map update") } - _, err = writer.Write(data) + _, err = w.Write(data) if err != nil { log.Error(). Str("handler", "PollNetMapStream"). @@ -471,7 +472,7 @@ func (h *Headscale) PollNetMapStream( updateRequestsSentToNode.WithLabelValues(machine.Namespace.Name, machine.Hostname, "failed"). Inc() - return false + return } log.Trace(). Str("handler", "PollNetMapStream"). @@ -499,7 +500,7 @@ func (h *Headscale) PollNetMapStream( // client has been removed from database // since the stream opened, terminate connection. - return false + return } now := time.Now().UTC() @@ -529,9 +530,9 @@ func (h *Headscale) PollNetMapStream( Msgf("%s is up to date", machine.Hostname) } - return true + return - case <-ctx.Request.Context().Done(): + case <-ctx.Done(): log.Info(). Str("handler", "PollNetMapStream"). Str("machine", machine.Hostname). @@ -550,7 +551,7 @@ func (h *Headscale) PollNetMapStream( // client has been removed from database // since the stream opened, terminate connection. - return false + break } now := time.Now().UTC() machine.LastSeen = &now @@ -564,9 +565,11 @@ func (h *Headscale) PollNetMapStream( Msg("Cannot update machine LastSeen") } - return false + break } - }) + } + + log.Info().Msgf("Closing poll loop to %s", machine.Hostname) } func (h *Headscale) scheduledPollWorker(