Bound discovery probe response bodies

Change-source: pulse-maintainer

pkg/discovery/discovery.go

@@ -20,10 +20,16 @@ import (
 
 	"github.com/rs/zerolog/log"
 
+	"github.com/rcourtman/pulse-go-rewrite/internal/securityutil"
 	"github.com/rcourtman/pulse-go-rewrite/pkg/discovery/envdetect"
 	"github.com/rcourtman/pulse-go-rewrite/pkg/tlsutil"
 )
 
+const (
+	maxVersionProbeResponseBodyBytes int64 = 64 << 10 // 64 KiB
+	maxNodesProbeResponseBodyBytes   int64 = 1 << 20  // 1 MiB
+)
+
 // DiscoveredServer represents a discovered Proxmox/PBS/PMG server
 type DiscoveredServer struct {
 	IP       string `json:"ip"`
@@ -1240,6 +1246,9 @@ func (s *Scanner) fetchNodeHostname(ctx context.Context, ip string, port int) st
 		return ""
 	}
 	defer resp.Body.Close()
+	if err := securityutil.LimitResponseBody(resp, maxNodesProbeResponseBodyBytes); err != nil {
+		return ""
+	}
 
 	var nodesResp struct {
 		Data []struct {
@@ -1313,6 +1322,10 @@ func (s *Scanner) probeVersionEndpoint(ctx context.Context, httpClient *http.Cli
 	if resp.StatusCode != http.StatusOK {
 		return finding, "", ""
 	}
+	if err := securityutil.LimitResponseBody(resp, maxVersionProbeResponseBodyBytes); err != nil {
+		finding.Error = err
+		return finding, "", ""
+	}
 
 	var payload struct {
 		Data struct {

pkg/discovery/discovery_test.go

@@ -249,6 +249,61 @@ func TestCheckServerRetrievesVersion(t *testing.T) {
 	}
 }
 
+func TestProbeVersionEndpointRejectsOversizedChunkedResponse(t *testing.T) {
+	t.Parallel()
+
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		if flusher, ok := w.(http.Flusher); ok {
+			flusher.Flush() // Keep Content-Length unknown so the streaming limit is exercised.
+		}
+		_, _ = w.Write([]byte(`{"data":{"version":"`))
+		_, _ = w.Write([]byte(strings.Repeat("x", int(maxVersionProbeResponseBodyBytes))))
+		_, _ = w.Write([]byte(`"}}`))
+	}))
+	defer ts.Close()
+
+	scanner := newTestScanner(ts.Client())
+	finding, version, release := scanner.probeVersionEndpoint(
+		context.Background(),
+		ts.Client(),
+		strings.TrimPrefix(ts.URL, "https://"),
+	)
+
+	if finding.Error == nil || !strings.Contains(finding.Error.Error(), "response body exceeds 65536 bytes") {
+		t.Fatalf("expected bounded response error, got %+v", finding.Error)
+	}
+	if version != "" || release != "" {
+		t.Fatalf("oversized response returned version=%q release=%q", version, release)
+	}
+}
+
+func TestFetchNodeHostnameRejectsOversizedDeclaredResponse(t *testing.T) {
+	t.Parallel()
+
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.Header().Set("Content-Length", strconv.FormatInt(maxNodesProbeResponseBodyBytes+1, 10))
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer ts.Close()
+
+	host, portText, err := net.SplitHostPort(ts.Listener.Addr().String())
+	if err != nil {
+		t.Fatalf("SplitHostPort: %v", err)
+	}
+	port, err := strconv.Atoi(portText)
+	if err != nil {
+		t.Fatalf("Atoi: %v", err)
+	}
+
+	scanner := newTestScanner(ts.Client())
+	if got := scanner.fetchNodeHostname(context.Background(), host, port); got != "" {
+		t.Fatalf("fetchNodeHostname returned %q for oversized response", got)
+	}
+}
+
 func TestNewScannerWithProfileAcceptsSelfSignedProxmoxProbe(t *testing.T) {
 	t.Parallel()