fix: relay middleware no longer eats local writes

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
 }

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 {