From c82a42ce196156c887d6b1fc5a0ee5d7c8aa4e91 Mon Sep 17 00:00:00 2001 From: Alexander Adam Date: Thu, 18 Jun 2026 16:56:14 +0200 Subject: [PATCH] fix: relay middleware no longer eats local writes filterByCache was replacing r.Body with the filtered slice. then the local store handler ran right after, read r.Body, and dutifully saved only the heartbeats that hadn't been relayed yet. everything else got quietly tossed. surprise. now we return the filtered bytes for the upstream send and leave r.Body where we found it. --- middlewares/custom/wakatime.go | 52 +++++++++++++++++------------ middlewares/custom/wakatime_test.go | 12 +++++-- 2 files changed, 40 insertions(+), 24 deletions(-) diff --git a/middlewares/custom/wakatime.go b/middlewares/custom/wakatime.go index 3e610797b..28f36d7c6 100644 --- a/middlewares/custom/wakatime.go +++ b/middlewares/custom/wakatime.go @@ -21,6 +21,9 @@ import ( const maxFailuresPerDay = 100 +// not really an error. it's the "seen all of these already, go back to sleep" signal. +var errNoNewHeartbeats = errors.New("no new heartbeats to relay") + // WakatimeRelayMiddleware is a middleware to conditionally relay heartbeats to Wakatime (and other compatible services) type WakatimeRelayMiddleware struct { httpClient *http.Client @@ -69,16 +72,17 @@ func (m *WakatimeRelayMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reque return } - err := m.filterByCache(r) + // relayBody goes upstream, r.Body still has the full thing for the + // local store handler. mix them up and your own heartbeats vanish into the void. + relayBody, err := m.filterByCache(r) + if errors.Is(err, errNoNewHeartbeats) { + return + } if err != nil { slog.Warn("filter cache error", "error", err) return } - body, _ := io.ReadAll(r.Body) - r.Body.Close() - r.Body = io.NopCloser(bytes.NewBuffer(body)) - // prevent cycles downstreamInstanceId := ownInstanceId if originInstanceId != "" { @@ -104,7 +108,7 @@ func (m *WakatimeRelayMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reque go m.send( http.MethodPost, url, - bytes.NewReader(body), + bytes.NewReader(relayBody), headers, user, ) @@ -147,23 +151,30 @@ func (m *WakatimeRelayMiddleware) send(method, url string, body io.Reader, heade } } -// filterByCache takes an HTTP request, tries to parse the body contents as heartbeats, checks against a local cache for whether a heartbeat has already been relayed before according to its hash and in-place filters these from the request's raw json body. -// This method operates on the raw body data (interface{}), because serialization of models.Heartbeat is not necessarily identical to what the CLI has actually sent. -// Purpose of this mechanism is mainly to prevent cyclic relays / loops. -// Caution: this method does in-place changes to the request. -func (m *WakatimeRelayMiddleware) filterByCache(r *http.Request) error { +// filterByCache returns the json body for the relay request, minus any +// heartbeats we've already forwarded. works on the raw decoded form +// (interface{}) since models.Heartbeat doesn't round-trip 1:1 with what +// the cli ships. point of all this: stop two linked instances from playing +// heartbeat ping-pong forever. +// r.Body is left alone so whoever runs after us still sees the full list. +func (m *WakatimeRelayMiddleware) filterByCache(r *http.Request) ([]byte, error) { heartbeats, err := routeutils.ParseHeartbeats(r) if err != nil { - return err + return nil, err } - body, _ := io.ReadAll(r.Body) - r.Body.Close() - r.Body = io.NopCloser(bytes.NewBuffer(body)) + // ParseHeartbeats already drained r.Body and put it back. read again here, + // we need the raw form because models.Heartbeat throws away fields the + // cli sends that we'd rather forward as-is. + body, err := io.ReadAll(r.Body) + if err != nil { + return nil, err + } + r.Body = io.NopCloser(bytes.NewReader(body)) var rawData interface{} - if err := json.NewDecoder(io.NopCloser(bytes.NewBuffer(body))).Decode(&rawData); err != nil { - return err + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&rawData); err != nil { + return nil, err } newData := make([]interface{}, 0, len(heartbeats)) @@ -186,7 +197,7 @@ func (m *WakatimeRelayMiddleware) filterByCache(r *http.Request) error { } if len(newData) == 0 { - return errors.New("no new heartbeats to relay") + return nil, errNoNewHeartbeats } if len(newData) != len(heartbeats) { @@ -196,9 +207,8 @@ func (m *WakatimeRelayMiddleware) filterByCache(r *http.Request) error { buf := bytes.Buffer{} if err := json.NewEncoder(&buf).Encode(newData); err != nil { - return err + return nil, err } - r.Body = io.NopCloser(&buf) - return nil + return buf.Bytes(), nil } diff --git a/middlewares/custom/wakatime_test.go b/middlewares/custom/wakatime_test.go index 32f7b6336..fe980b0c0 100644 --- a/middlewares/custom/wakatime_test.go +++ b/middlewares/custom/wakatime_test.go @@ -207,7 +207,7 @@ func (suite *WakatimeRelayMiddlewareTestSuite) TestFilterByCache_Bulk() { req, _ := http.NewRequest(http.MethodPost, "/api/heartbeats", bytes.NewBuffer(body)) req = suite.withUser(req, user) - err := suite.sut.filterByCache(req) + _, err := suite.sut.filterByCache(req) suite.NoError(err) // Now hb1 and hb2 are in cache. Try again with hb2 and hb3 @@ -220,13 +220,19 @@ func (suite *WakatimeRelayMiddlewareTestSuite) TestFilterByCache_Bulk() { req2, _ := http.NewRequest(http.MethodPost, "/api/heartbeats", bytes.NewBuffer(body2)) req2 = suite.withUser(req2, user) - err = suite.sut.filterByCache(req2) + relayBody, err := suite.sut.filterByCache(req2) suite.NoError(err) + // upstream gets only hb3, hb2 is old news var filtered []interface{} - json.NewDecoder(req2.Body).Decode(&filtered) + json.NewDecoder(bytes.NewReader(relayBody)).Decode(&filtered) suite.Len(filtered, 1) suite.Equal("/tmp/3.go", filtered[0].(map[string]interface{})["entity"]) + + // but r.Body still carries both, otherwise the local store eats one + var passedThrough []interface{} + json.NewDecoder(req2.Body).Decode(&passedThrough) + suite.Len(passedThrough, 2) } func (suite *WakatimeRelayMiddlewareTestSuite) withUser(r *http.Request, user *models.User) *http.Request {