From bfdfe989e05c88b81569fad2f23aa9d0031c007a Mon Sep 17 00:00:00 2001 From: ZverGuy Date: Mon, 13 Apr 2026 18:20:12 +0300 Subject: [PATCH] fix(controller): add HTTP status check and limit response body size - Check resp.StatusCode before parsing Prometheus response - Limit response body read to 1 MB via io.LimitReader - Use strings.HasPrefix in test instead of fragile slice indexing - Add TestQueryPrometheusMetricServerError for 500 responses Co-Authored-By: Claude Signed-off-by: ZverGuy --- .../controller/workloadmonitor_controller.go | 7 ++++++- .../workloadmonitor_controller_test.go | 16 +++++++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/internal/controller/workloadmonitor_controller.go b/internal/controller/workloadmonitor_controller.go index d7514552..900610e4 100644 --- a/internal/controller/workloadmonitor_controller.go +++ b/internal/controller/workloadmonitor_controller.go @@ -146,7 +146,12 @@ func (r *WorkloadMonitorReconciler) queryPrometheusMetric(ctx context.Context, p } defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + logger.V(1).Info("Prometheus returned non-OK status", "query", promQL, "status", resp.StatusCode) + return 0 + } + + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) if err != nil { logger.Error(err, "Failed to read Prometheus response") return 0 diff --git a/internal/controller/workloadmonitor_controller_test.go b/internal/controller/workloadmonitor_controller_test.go index de9b78e5..e8feb0b0 100644 --- a/internal/controller/workloadmonitor_controller_test.go +++ b/internal/controller/workloadmonitor_controller_test.go @@ -5,6 +5,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "strings" "testing" cozyv1alpha1 "github.com/cozystack/cozystack/api/v1alpha1" @@ -412,6 +413,19 @@ func TestQueryPrometheusMetricEmpty(t *testing.T) { } } +func TestQueryPrometheusMetricServerError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + reconciler := &WorkloadMonitorReconciler{PrometheusURL: srv.URL} + size := reconciler.queryPrometheusMetric(context.TODO(), `SeaweedFS_s3_bucket_size_bytes{bucket="test"}`) + if size != 0 { + t.Errorf("expected 0 for server error, got %d", size) + } +} + func TestQueryPrometheusMetricNoURL(t *testing.T) { reconciler := &WorkloadMonitorReconciler{PrometheusURL: ""} size := reconciler.queryPrometheusMetric(context.TODO(), `anything`) @@ -424,7 +438,7 @@ func TestReconcileBucketClaimWithPrometheus(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { query := r.URL.Query().Get("query") switch { - case len(query) > 0 && query[0:len("SeaweedFS_s3_bucket_physical")] == "SeaweedFS_s3_bucket_physical": + case strings.HasPrefix(query, "SeaweedFS_s3_bucket_physical"): fmt.Fprint(w, `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"bucket":"cosi-abc123"},"value":[1713000000,"2147483648"]}]}}`) default: fmt.Fprint(w, `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"bucket":"cosi-abc123"},"value":[1713000000,"1073741824"]}]}}`)