Enforce response body limits in kubelet-to-gcm

kubelet-to-gcm/monitor/config/initialize.go

@@ -18,7 +18,7 @@ package config
 
 import (
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net/http"
 	"strings"
 	"time"
@@ -29,6 +29,9 @@ import (
 const (
 	gceMetaDataEndpoint = "http://169.254.169.254"
 	gceMetaDataPrefix   = "/computeMetadata/v1"
+	// maxResponseBodySize is the maximum size of the response body we will read.
+	// 10MB is more than enough for GCE metadata responses.
+	maxResponseBodySize = 10 * 1024 * 1024
 )
 
 // NewConfigs returns the SourceConfigs for all monitored endpoints, and
@@ -117,7 +120,7 @@ func getGCEMetaData(uri string) ([]byte, error) {
 		return nil, fmt.Errorf("Failed request %q for GCE metadata: %v", uri, err)
 	}
 	defer resp.Body.Close()
-	body, err := ioutil.ReadAll(resp.Body)
+	body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBodySize))
 	if err != nil {
 		return nil, fmt.Errorf("Failed to read body for request %q for GCE metadata: %v", uri, err)
 	}

kubelet-to-gcm/monitor/controller/client.go

@@ -19,7 +19,6 @@ package controller
 import (
 	"fmt"
 	"io"
-	"io/ioutil"
 	"net/http"
 	"net/url"
 	"strings"
@@ -72,6 +71,12 @@ func NewMetrics(body []byte) (*Metrics, error) {
 	return metrics, nil
 }
 
+const (
+	// maxResponseBodySize is the maximum size of the response body we will read.
+	// 10MB is more than enough for most prometheus metrics responses.
+	maxResponseBodySize = 10 * 1024 * 1024
+)
+
 // Client queries metrics from the controller process.
 type Client struct {
 	client     *http.Client
@@ -99,7 +104,7 @@ func (c *Client) doRequestAndParse(req *http.Request) (*Metrics, error) {
 		return nil, err
 	}
 	defer response.Body.Close()
-	body, err := ioutil.ReadAll(response.Body)
+	body, err := io.ReadAll(io.LimitReader(response.Body, maxResponseBodySize))
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body - %v", err)
 	}

kubelet-to-gcm/monitor/controller/client_test.go

@@ -0,0 +1,91 @@
+/*
+Copyright 2017 Google Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package controller
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+)
+
+func TestDoRequestAndParse_SizeLimit(t *testing.T) {
+	// Create a mock server that returns a response larger than maxResponseBodySize.
+	largeDataSize := maxResponseBodySize + 1024
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Write more than maxResponseBodySize bytes.
+		data := make([]byte, largeDataSize)
+		w.Write(data)
+	}))
+	defer ts.Close()
+
+	u, _ := url.Parse(ts.URL)
+	client := &Client{
+		client:     ts.Client(),
+		metricsURL: u,
+	}
+
+	req, _ := http.NewRequest("GET", ts.URL, nil)
+	_, err := client.doRequestAndParse(req)
+
+	if err == nil {
+		t.Fatal("Expected error due to size limit, but got none")
+	}
+
+	// The error should indicate a problem reading the body or parsing it.
+	// Since we are writing random null bytes, parsing should fail if it gets that far.
+	// But io.LimitReader will just stop at maxResponseBodySize.
+	// io.ReadAll doesn't return an error if it reaches EOF.
+	// However, NewMetrics(body) will fail to parse the truncated/invalid data.
+	fmt.Printf("Got expected error: %v\n", err)
+}
+
+func TestDoRequestAndParse_Success(t *testing.T) {
+	metricsData := `
+# HELP node_collector_evictions_number Number of evictions
+# TYPE node_collector_evictions_number counter
+node_collector_evictions_number 10
+# HELP process_start_time_seconds Start time
+# TYPE process_start_time_seconds gauge
+process_start_time_seconds 1234567890
+`
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprint(w, metricsData)
+	}))
+	defer ts.Close()
+
+	u, _ := url.Parse(ts.URL)
+	client := &Client{
+		client:     ts.Client(),
+		metricsURL: u,
+	}
+
+	req, _ := http.NewRequest("GET", ts.URL, nil)
+	metrics, err := client.doRequestAndParse(req)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if metrics.NodeEvictions != 10 {
+		t.Errorf("Expected 10 evictions, got %d", metrics.NodeEvictions)
+	}
+	if metrics.CreateTime != 1234567890 {
+		t.Errorf("Expected 1234567890 create time, got %d", metrics.CreateTime)
+	}
+}

kubelet-to-gcm/monitor/kubelet/client.go

@@ -19,13 +19,19 @@ package kubelet
 import (
 	"encoding/json"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net/http"
 	"net/url"
 
 	stats "k8s.io/kubelet/pkg/apis/stats/v1alpha1"
 )
 
+const (
+	// maxResponseBodySize is the maximum size of the response body we will read.
+	// 50MB should be enough for any Kubelet stats summary response.
+	maxResponseBodySize = 50 * 1024 * 1024
+)
+
 // Client contains all the information and methods to encapsulate
 // communication with the Kubelet.
 type Client struct {
@@ -58,7 +64,7 @@ func (k *Client) doRequestAndUnmarshal(client *http.Client, req *http.Request, v
 		return err
 	}
 	defer response.Body.Close()
-	body, err := ioutil.ReadAll(response.Body)
+	body, err := io.ReadAll(io.LimitReader(response.Body, maxResponseBodySize))
 	if err != nil {
 		return fmt.Errorf("failed to read response body - %v", err)
 	}

kubelet-to-gcm/monitor/kubelet/client_test.go

@@ -0,0 +1,84 @@
+/*
+Copyright 2017 Google Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kubelet
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	stats "k8s.io/kubelet/pkg/apis/stats/v1alpha1"
+)
+
+func TestDoRequestAndUnmarshal_SizeLimit(t *testing.T) {
+	largeDataSize := maxResponseBodySize + 1024
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		data := make([]byte, largeDataSize)
+		w.Write(data)
+	}))
+	defer ts.Close()
+
+	u, _ := url.Parse(ts.URL)
+	k := &Client{
+		client:     ts.Client(),
+		summaryURL: u,
+	}
+
+	req, _ := http.NewRequest("GET", ts.URL, nil)
+	var value stats.Summary
+	err := k.doRequestAndUnmarshal(ts.Client(), req, &value)
+
+	if err == nil {
+		t.Fatal("Expected error due to size limit, but got none")
+	}
+	fmt.Printf("Got expected error: %v\n", err)
+}
+
+func TestDoRequestAndUnmarshal_Success(t *testing.T) {
+	summary := &stats.Summary{
+		Node: stats.NodeStats{
+			NodeName: "test-node",
+		},
+	}
+	data, _ := json.Marshal(summary)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write(data)
+	}))
+	defer ts.Close()
+
+	u, _ := url.Parse(ts.URL)
+	k := &Client{
+		client:     ts.Client(),
+		summaryURL: u,
+	}
+
+	req, _ := http.NewRequest("GET", ts.URL, nil)
+	var result stats.Summary
+	err := k.doRequestAndUnmarshal(ts.Client(), req, &result)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if result.Node.NodeName != "test-node" {
+		t.Errorf("Expected test-node, got %s", result.Node.NodeName)
+	}
+}