Cap HTTP request and websocket message sizes

Unauthenticated clients could exhaust process memory (and, for /store,
disk) by posting arbitrarily large bodies, and could OOM or pin
goroutines via the /ws/echo websocket.

- Wrap request bodies in http.MaxBytesReader (10 MiB) on the echo,
  store, cache and token handlers via a shared readLimitedBody helper
- Bound /ws/echo: per-message read limit, idle read deadline with
  ping/pong keepalive, and write deadlines
- Add regression tests for both limits

Author: stefanprodan

pkg/api/http/body_limit_test.go

@@ -0,0 +1,42 @@
+package http
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+// TestRequestBodySizeLimit verifies body-reading handlers reject payloads larger
+// than maxRequestBodySize instead of buffering them into memory. The echo
+// handler is used because with no backends configured it simply reflects the
+// body and needs no external dependencies.
+func TestRequestBodySizeLimit(t *testing.T) {
+	srv := NewMockServer()
+	srv.router.HandleFunc("/echo", srv.echoHandler)
+
+	// A body within the limit is accepted (202).
+	within := httptest.NewRequest("POST", "/echo", bytes.NewReader(make([]byte, 1024)))
+	rr := httptest.NewRecorder()
+	srv.router.ServeHTTP(rr, within)
+	if rr.Code != http.StatusAccepted {
+		t.Errorf("within-limit body: got status %d want %d", rr.Code, http.StatusAccepted)
+	}
+
+	// A body over the limit is rejected with a 413 code in the response body.
+	over := httptest.NewRequest("POST", "/echo", bytes.NewReader(make([]byte, maxRequestBodySize+1)))
+	rr = httptest.NewRecorder()
+	srv.router.ServeHTTP(rr, over)
+
+	var resp struct {
+		Code    int    `json:"code"`
+		Message string `json:"message"`
+	}
+	if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
+		t.Fatalf("oversize body: response is not the expected error JSON: %v (body=%q)", err, rr.Body.String())
+	}
+	if resp.Code != http.StatusRequestEntityTooLarge {
+		t.Errorf("oversize body: got code %d want %d", resp.Code, http.StatusRequestEntityTooLarge)
+	}
+}

pkg/api/http/cache.go

@@ -2,7 +2,6 @@ package http
 
 import (
 	"fmt"
-	"io"
 	"net/http"
 	"net/url"
 	"time"
@@ -33,15 +32,14 @@ func (s *Server) cacheWriteHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	key := mux.Vars(r)["key"]
-	body, err := io.ReadAll(r.Body)
-	if err != nil {
-		s.ErrorResponse(w, r, span, "reading the request body failed", http.StatusBadRequest)
+	body, ok := s.readLimitedBody(w, r, span)
+	if !ok {
 		return
 	}
 
 	conn := s.pool.Get()
 	defer conn.Close()
-	_, err = conn.Do("SET", key, string(body))
+	_, err := conn.Do("SET", key, string(body))
 	if err != nil {
 		s.logger.Warn("cache set failed", zap.Error(err))
 		s.ErrorResponse(w, r, span, "cache set failed", http.StatusInternalServerError)

pkg/api/http/echo.go

@@ -27,10 +27,8 @@ func (s *Server) echoHandler(w http.ResponseWriter, r *http.Request) {
 	ctx, span := s.tracer.Start(r.Context(), "echoHandler")
 	defer span.End()
 
-	body, err := io.ReadAll(r.Body)
-	if err != nil {
-		s.logger.Error("reading the request body failed", zap.Error(err))
-		s.ErrorResponse(w, r, span, "invalid request body", http.StatusBadRequest)
+	body, ok := s.readLimitedBody(w, r, span)
+	if !ok {
 		return
 	}
 	defer r.Body.Close()

pkg/api/http/echows.go

@@ -12,6 +12,20 @@ import (
 
 var wsCon = websocket.Upgrader{}
 
+const (
+	// wsMaxMessageSize caps a single inbound websocket message so one large
+	// frame cannot exhaust process memory.
+	wsMaxMessageSize = 1 << 20 // 1 MiB
+	// wsReadTimeout is how long the server waits for the next client message or
+	// pong before closing an idle connection.
+	wsReadTimeout = 60 * time.Second
+	// wsWriteTimeout bounds a single write so a slow reader cannot block forever.
+	wsWriteTimeout = 10 * time.Second
+	// wsPingInterval is how often the server pings the client to keep the
+	// connection alive and detect dead peers; it must be shorter than wsReadTimeout.
+	wsPingInterval = 30 * time.Second
+)
+
 // EchoWS godoc
 // @Summary Echo over websockets
 // @Description echos content via websockets
@@ -24,11 +38,18 @@ var wsCon = websocket.Upgrader{}
 func (s *Server) echoWsHandler(w http.ResponseWriter, r *http.Request) {
 	c, err := wsCon.Upgrade(w, r, nil)
 	if err != nil {
-		if err != nil {
-			s.logger.Warn("websocket upgrade error", zap.Error(err))
-			return
-		}
+		s.logger.Warn("websocket upgrade error", zap.Error(err))
+		return
 	}
+
+	// Bound per-message size and idle time; refresh the read deadline whenever
+	// the client responds to a ping so live connections stay open.
+	c.SetReadLimit(wsMaxMessageSize)
+	_ = c.SetReadDeadline(time.Now().Add(wsReadTimeout))
+	c.SetPongHandler(func(string) error {
+		return c.SetReadDeadline(time.Now().Add(wsReadTimeout))
+	})
+
 	var wg sync.WaitGroup
 	wg.Add(1)
 
@@ -84,15 +105,23 @@ func (s *Server) sendHostWs(ws *websocket.Conn, in chan interface{}, done chan s
 }
 
 func (s *Server) writeWs(ws *websocket.Conn, in chan interface{}) {
+	ping := time.NewTicker(wsPingInterval)
+	defer ping.Stop()
 	for {
 		select {
 		case msg := <-in:
+			_ = ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
 			if err := ws.WriteJSON(msg); err != nil {
 				if !strings.Contains(err.Error(), "close") {
 					s.logger.Warn("websocket write error", zap.Error(err))
 				}
 				return
 			}
+		case <-ping.C:
+			_ = ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
+			if err := ws.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
 		}
 	}
 }

pkg/api/http/echows_test.go

@@ -4,6 +4,7 @@ import (
 	"net/http/httptest"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/gorilla/websocket"
 )
@@ -34,3 +35,40 @@ func TestEchoWsHandler(t *testing.T) {
 		t.Error("received empty message")
 	}
 }
+
+// TestEchoWsReadLimit verifies the server caps inbound message size: a message
+// larger than wsMaxMessageSize must cause the server to close the connection
+// instead of buffering it into memory.
+func TestEchoWsReadLimit(t *testing.T) {
+	srv := NewMockServer()
+	srv.router.HandleFunc("/ws/echo", srv.echoWsHandler)
+	server := httptest.NewServer(srv.router)
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws/echo"
+	ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("websocket dial failed: %v", err)
+	}
+	defer ws.Close()
+
+	// A message just over the limit must not be echoed; the server closes it.
+	oversize := make([]byte, wsMaxMessageSize+1)
+	if err := ws.WriteMessage(websocket.TextMessage, oversize); err != nil {
+		t.Fatalf("write failed: %v", err)
+	}
+
+	// Read until the connection errors. With the read limit in place the server
+	// sends a 1009 (message too big) close frame; without it the server would
+	// instead echo the oversize payload back. The 5s deadline bounds the loop.
+	// Status frames from the periodic ticker (err == nil) are skipped.
+	ws.SetReadDeadline(time.Now().Add(5 * time.Second))
+	for {
+		if _, _, err = ws.ReadMessage(); err != nil {
+			break
+		}
+	}
+	if !websocket.IsCloseError(err, websocket.CloseMessageTooBig) {
+		t.Errorf("expected close code %d (message too big), got: %v", websocket.CloseMessageTooBig, err)
+	}
+}

pkg/api/http/http.go

@@ -3,6 +3,8 @@ package http
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
+	"io"
 	"math/rand"
 	"net/http"
 	"time"
@@ -13,6 +15,11 @@ import (
 	"go.uber.org/zap"
 )
 
+// maxRequestBodySize caps how much of a request body the body-reading handlers
+// buffer into memory. Without this bound an unauthenticated client can POST an
+// arbitrarily large body and exhaust process memory (and, for /store, disk).
+const maxRequestBodySize = 10 << 20 // 10 MiB
+
 func randomErrorMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		rand.Seed(time.Now().Unix())
@@ -87,6 +94,25 @@ func (s *Server) ErrorResponse(w http.ResponseWriter, r *http.Request, span trac
 	w.Write(prettyJSON(body))
 }
 
+// readLimitedBody reads the request body up to maxRequestBodySize. It returns
+// the body and true on success. On an oversized body it writes a 413 response,
+// on any other read error a 400, and returns false so the caller returns early.
+func (s *Server) readLimitedBody(w http.ResponseWriter, r *http.Request, span trace.Span) ([]byte, bool) {
+	r.Body = http.MaxBytesReader(w, r.Body, maxRequestBodySize)
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		var maxErr *http.MaxBytesError
+		if errors.As(err, &maxErr) {
+			s.ErrorResponse(w, r, span, "request body too large", http.StatusRequestEntityTooLarge)
+			return nil, false
+		}
+		s.logger.Error("reading the request body failed", zap.Error(err))
+		s.ErrorResponse(w, r, span, "invalid request body", http.StatusBadRequest)
+		return nil, false
+	}
+	return body, true
+}
+
 // setRawResponseHeaders prevents XSS by ensuring browsers never interpret raw responses as HTML.
 func setRawResponseHeaders(w http.ResponseWriter) {
 	w.Header().Set("Content-Type", "application/octet-stream")

pkg/api/http/store.go

@@ -3,7 +3,6 @@ package http
 import (
 	"crypto/sha1"
 	"encoding/hex"
-	"io"
 	"net/http"
 	"os"
 	"path"
@@ -27,14 +26,13 @@ func (s *Server) storeWriteHandler(w http.ResponseWriter, r *http.Request) {
 	_, span := s.tracer.Start(r.Context(), "storeWriteHandler")
 	defer span.End()
 
-	body, err := io.ReadAll(r.Body)
-	if err != nil {
-		s.ErrorResponse(w, r, span, "reading the request body failed", http.StatusBadRequest)
+	body, ok := s.readLimitedBody(w, r, span)
+	if !ok {
 		return
 	}
 
 	hash := hash(string(body))
-	err = os.WriteFile(path.Join(s.config.DataPath, hash), body, 0644)
+	err := os.WriteFile(path.Join(s.config.DataPath, hash), body, 0644)
 	if err != nil {
 		s.logger.Warn("writing file failed", zap.Error(err), zap.String("file", path.Join(s.config.DataPath, hash)))
 		s.ErrorResponse(w, r, span, "writing file failed", http.StatusInternalServerError)

pkg/api/http/token.go

@@ -2,13 +2,11 @@ package http
 
 import (
 	"fmt"
-	"io"
 	"net/http"
 	"strings"
 	"time"
 
 	"github.com/golang-jwt/jwt/v4"
-	"go.uber.org/zap"
 )
 
 type jwtCustomClaims struct {
@@ -28,10 +26,8 @@ func (s *Server) tokenGenerateHandler(w http.ResponseWriter, r *http.Request) {
 	_, span := s.tracer.Start(r.Context(), "tokenGenerateHandler")
 	defer span.End()
 
-	body, err := io.ReadAll(r.Body)
-	if err != nil {
-		s.logger.Error("reading the request body failed", zap.Error(err))
-		s.ErrorResponse(w, r, span, "invalid request body", http.StatusBadRequest)
+	body, ok := s.readLimitedBody(w, r, span)
+	if !ok {
 		return
 	}
 	defer r.Body.Close()