fix(observability-proxy): cut the connection when a log answer breaks off

Once the status line has gone out, a read error from the store left
the handler returning normally, so the chunked body ended cleanly and
the client took a partial answer for a complete one.

The handler now aborts the response in that case. The request is
counted before the abort, because the abort skips the rest of the
handler.

Assisted-by: LLM
Signed-off-by: ZverGuy <maximbel2003@gmail.com>

internal/observabilityproxy/server.go

@@ -41,6 +41,11 @@ const tailChunkBytes = 64 << 10
 // it wraps DeadlineExceeded so a store that never answered reads as 504.
 var errTailIdle = fmt.Errorf("tail idle: %w", context.DeadlineExceeded)
 
+// errAborted ends a response whose status line already went out by
+// cutting the connection, so a client sees a broken body instead of a
+// complete one with lines missing.
+var errAborted = &HTTPError{Code: http.StatusBadGateway, Message: "response aborted"}
+
 // maxFormBytes bounds a POST body; LogsQL queries are a few KiB at most.
 const maxFormBytes = 1 << 20
 
@@ -129,10 +134,13 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if herr == nil {
 		herr = s.serve(rec, r, route)
 	}
-	if herr != nil {
+	if herr != nil && herr != errAborted {
 		writeError(rec, herr)
 	}
 	s.requests.WithLabelValues(route.Kind.String(), strconv.Itoa(rec.code)).Inc()
+	if herr == errAborted {
+		panic(http.ErrAbortHandler)
+	}
 }
 
 func writeError(w http.ResponseWriter, e *HTTPError) {
@@ -251,8 +259,9 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, route Route) *HTT
 	}
 	if tail {
 		s.streamLines(w, rc, resp.Body, idle, arm)
-	} else if _, err := io.Copy(w, resp.Body); err != nil && !errors.Is(err, context.Canceled) {
+	} else if _, err := io.Copy(w, resp.Body); err != nil {
 		slog.Debug("copying upstream response", "error", err)
+		return errAborted
 	}
 	return nil
 }

internal/observabilityproxy/server_deadline_test.go

@@ -18,6 +18,7 @@ package observabilityproxy
 
 import (
 	"bufio"
+	"io"
 	"net"
 	"net/http"
 	"sync"
@@ -228,3 +229,18 @@ func TestServerTailUpstreamWithoutHeaders(t *testing.T) {
 		t.Fatalf("open tails = %d after the tail ended", n)
 	}
 }
+
+// TestServerAbortsTruncatedResponse cuts the client connection when the
+// store's body breaks off after the headers went out, so a partial answer
+// does not end as a clean 200.
+func TestServerAbortsTruncatedResponse(t *testing.T) {
+	h := newHarness(t, Config{}, func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Length", "100")
+		_, _ = io.WriteString(w, `{"_msg":"a"}`+"\n")
+	})
+	resp := h.do(http.MethodGet, logsRoot+"query?query=*", "good", nil, nil)
+	defer func() { _ = resp.Body.Close() }()
+	if _, err := io.ReadAll(resp.Body); err == nil {
+		t.Fatal("a truncated store answer reached the client as a complete response")
+	}
+}