fix(net): cap response body reads via configurable size limit

lambda/lambda.go

@@ -110,6 +110,7 @@ func handleRequest(ctx context.Context, req CheckRequest) (CheckResponse, error)
 		Headers:         req.Headers,
 		Method:          req.Method,
 		Body:            req.Body,
+		BodySizeLimit:   net.DefaultBodySizeLimit,
 	}
 
 	result := net.CheckWebsite(req.URL, netConfig)

net/net.go

@@ -3,10 +3,12 @@ package net
 import (
 	"bytes"
 	"crypto/tls"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
 	"log"
+	"math"
 	"net"
 	"net/http"
 	"net/http/httptrace"
@@ -23,23 +25,29 @@ const (
 	_defaultTimeout = 5 * time.Second
 	_userAgent      = "updo/1.0"
 	_httpsPort      = ":443"
+
+	// DefaultBodySizeLimit is the recommended response body cap (1 MiB) for
+	// probe responses. Callers should set NetworkConfig.BodySizeLimit to this
+	// value unless they have a reason to cap smaller or disable capping.
+	DefaultBodySizeLimit int64 = 1 << 20
 )
 
 type WebsiteCheckResult struct {
-	URL             string
-	ResolvedIP      string
-	IsUp            bool
-	StatusCode      int
-	ResponseTime    time.Duration
-	TraceInfo       *HttpTraceInfo
-	AssertionPassed bool
-	LastCheckTime   time.Time
-	AssertText      string
-	Method          string
-	RequestHeaders  http.Header
-	ResponseHeaders http.Header
-	RequestBody     string
-	ResponseBody    string
+	URL               string
+	ResolvedIP        string
+	IsUp              bool
+	StatusCode        int
+	ResponseTime      time.Duration
+	TraceInfo         *HttpTraceInfo
+	AssertionPassed   bool
+	LastCheckTime     time.Time
+	AssertText        string
+	Method            string
+	RequestHeaders    http.Header
+	ResponseHeaders   http.Header
+	RequestBody       string
+	ResponseBody      string
+	ResponseTruncated bool
 }
 type HttpTraceInfo struct {
 	Wait             time.Duration
@@ -59,6 +67,8 @@ type NetworkConfig struct {
 	Headers         []string
 	Method          string
 	Body            string
+	// BodySizeLimit caps bytes read from the response body. 0 means no limit.
+	BodySizeLimit int64
 }
 
 type HTTPRequestOptions struct {
@@ -68,20 +78,21 @@ type HTTPRequestOptions struct {
 }
 
 type HTTPResponse struct {
-	URL             string
-	ResolvedIP      string
-	StatusCode      int
-	StatusText      string
-	HTTPVersion     string
-	ResponseHeaders http.Header
-	ResponseBody    string
-	RequestHeaders  http.Header
-	RequestBody     string
-	Method          string
-	ResponseTime    time.Duration
-	TraceInfo       *HttpTraceInfo
-	LastCheckTime   time.Time
-	Error           error
+	URL               string
+	ResolvedIP        string
+	StatusCode        int
+	StatusText        string
+	HTTPVersion       string
+	ResponseHeaders   http.Header
+	ResponseBody      string
+	ResponseTruncated bool
+	RequestHeaders    http.Header
+	RequestBody       string
+	Method            string
+	ResponseTime      time.Duration
+	TraceInfo         *HttpTraceInfo
+	LastCheckTime     time.Time
+	Error             error
 }
 
 func CheckWebsite(urlStr string, config NetworkConfig) WebsiteCheckResult {
@@ -105,18 +116,19 @@ func CheckWebsite(urlStr string, config NetworkConfig) WebsiteCheckResult {
 	httpResp := makeHTTPRequest(urlStr, options, config)
 
 	result := WebsiteCheckResult{
-		URL:             urlStr,
-		ResolvedIP:      httpResp.ResolvedIP,
-		StatusCode:      httpResp.StatusCode,
-		ResponseTime:    httpResp.ResponseTime,
-		TraceInfo:       httpResp.TraceInfo,
-		LastCheckTime:   httpResp.LastCheckTime,
-		AssertText:      config.AssertText,
-		Method:          options.Method,
-		RequestHeaders:  httpResp.RequestHeaders,
-		ResponseHeaders: httpResp.ResponseHeaders,
-		RequestBody:     httpResp.RequestBody,
-		ResponseBody:    httpResp.ResponseBody,
+		URL:               urlStr,
+		ResolvedIP:        httpResp.ResolvedIP,
+		StatusCode:        httpResp.StatusCode,
+		ResponseTime:      httpResp.ResponseTime,
+		TraceInfo:         httpResp.TraceInfo,
+		LastCheckTime:     httpResp.LastCheckTime,
+		AssertText:        config.AssertText,
+		Method:            options.Method,
+		RequestHeaders:    httpResp.RequestHeaders,
+		ResponseHeaders:   httpResp.ResponseHeaders,
+		RequestBody:       httpResp.RequestBody,
+		ResponseBody:      httpResp.ResponseBody,
+		ResponseTruncated: httpResp.ResponseTruncated,
 	}
 
 	if httpResp.Error != nil {
@@ -333,10 +345,24 @@ func makeHTTPRequest(urlStr string, options HTTPRequestOptions, config NetworkCo
 		}
 	}()
 
-	bodyBytes, err := io.ReadAll(resp.Body)
+	body := resp.Body
+	if config.BodySizeLimit > 0 {
+		limit := config.BodySizeLimit
+		// MaxBytesReader internally does limit+1, so clamp MaxInt64 to avoid overflow.
+		if limit == math.MaxInt64 {
+			limit = math.MaxInt64 - 1
+		}
+		body = http.MaxBytesReader(nil, resp.Body, limit)
+	}
+	bodyBytes, err := io.ReadAll(body)
 	if err != nil {
-		result.Error = err
-		return result
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			result.ResponseTruncated = true
+		} else {
+			result.Error = err
+			return result
+		}
 	}
 
 	result.StatusCode = resp.StatusCode

net/net_test.go

@@ -379,3 +379,73 @@ func TestCheckWebsiteWithHeaders(t *testing.T) {
 		t.Errorf("CheckWebsite() StatusCode = %d, want 200", result.StatusCode)
 	}
 }
+
+func TestCheckWebsiteBodyLimit(t *testing.T) {
+	tests := []struct {
+		name           string
+		bodySize       int
+		bodySizeLimit  int64
+		wantTruncated  bool
+		wantBodyLength int
+	}{
+		{
+			name:           "body under custom limit",
+			bodySize:       500,
+			bodySizeLimit:  1000,
+			wantTruncated:  false,
+			wantBodyLength: 500,
+		},
+		{
+			name:           "body equals custom limit",
+			bodySize:       1000,
+			bodySizeLimit:  1000,
+			wantTruncated:  false,
+			wantBodyLength: 1000,
+		},
+		{
+			name:           "body exceeds custom limit",
+			bodySize:       2000,
+			bodySizeLimit:  1000,
+			wantTruncated:  true,
+			wantBodyLength: 1000,
+		},
+		{
+			name:           "zero BodySizeLimit means unlimited",
+			bodySize:       2 * 1024 * 1024,
+			bodySizeLimit:  0,
+			wantTruncated:  false,
+			wantBodyLength: 2 * 1024 * 1024,
+		},
+		{
+			name:           "DefaultBodySizeLimit caps at 1 MiB",
+			bodySize:       2 * 1024 * 1024,
+			bodySizeLimit:  DefaultBodySizeLimit,
+			wantTruncated:  true,
+			wantBodyLength: 1 << 20,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			payload := make([]byte, tt.bodySize)
+			for i := range payload {
+				payload[i] = 'x'
+			}
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(200)
+				_, _ = w.Write(payload)
+			}))
+			defer server.Close()
+
+			config := NetworkConfig{Timeout: 5 * time.Second, BodySizeLimit: tt.bodySizeLimit}
+			result := CheckWebsite(server.URL, config)
+
+			if result.ResponseTruncated != tt.wantTruncated {
+				t.Errorf("ResponseTruncated = %v, want %v", result.ResponseTruncated, tt.wantTruncated)
+			}
+			if len(result.ResponseBody) != tt.wantBodyLength {
+				t.Errorf("ResponseBody length = %d, want %d", len(result.ResponseBody), tt.wantBodyLength)
+			}
+		})
+	}
+}

simple/monitoring.go

@@ -214,6 +214,7 @@ func monitorTargetSimple(ctx context.Context, target config.Target, targetIndex
 			Headers:         target.Headers,
 			Method:          target.Method,
 			Body:            target.Body,
+			BodySizeLimit:   net.DefaultBodySizeLimit,
 		}
 
 		regions := target.Regions
@@ -277,6 +278,9 @@ func monitorTargetSimple(ctx context.Context, target config.Target, targetIndex
 
 			if monitor, exists := monitors[keyStr]; exists {
 				result := net.CheckWebsite(target.URL, netConfig)
+				if result.ResponseTruncated {
+					log.Printf("Warning: response body from %s truncated at BodySizeLimit of %d bytes", target.URL, netConfig.BodySizeLimit)
+				}
 				monitor.AddResult(result)
 				if sequence, exists := sequences[keyStr]; exists {
 					*sequence++

tui/manager.go

@@ -454,6 +454,11 @@ func (m *Manager) UpdateTarget(data TargetData) {
 		logAdded = true
 	}
 
+	if data.Result.ResponseTruncated {
+		m.logBuffer.AddLogEntry(LogLevelWarning, "Response body truncated", "Exceeded BodySizeLimit; assertion checks may be unreliable", data.TargetKey)
+		logAdded = true
+	}
+
 	if !data.Result.IsUp && data.LambdaError == nil {
 		level := LogLevelError
 		message := "Request failed"

tui/monitoring.go

@@ -279,6 +279,7 @@ func monitorTargetTUI(ctx context.Context, target config.Target, targetIndex int
 			Headers:         target.Headers,
 			Method:          target.Method,
 			Body:            target.Body,
+			BodySizeLimit:   net.DefaultBodySizeLimit,
 		}
 
 		regions := target.Regions