fix(rw2.0): reject remote write 2.0 based on content type

The current solution returns 2xx , but doesn't actually ingest the samples.

Prometheus does detect this
prometheus-1  | time=2025-01-13T13:01:35.028Z level=ERROR
source=queue_manager.go:1670 msg="non-recoverable error"
component=remote remote_name=150c10
 url=http://mimir-1:8001/api/v1/push failedSampleCount=2000
failedHistogramCount=0 failedExemplarCount=0
err="sent v2 request with 2000 samples, 0 histograms and 0 exemplars;
 got 2xx, but PRW 2.0 response header statistics indicate 0 samples,
 0 histograms and 0 exemplars were accepted; assumining failure e.g.
 the target only supports PRW 1.0 prometheus.WriteRequest, but does not
 check the Content-Type header correctly"

But we can do better and also start working towards RW2.0 support.

Signed-off-by: György Krajcsovits <[email protected]>

Author: krajorama

pkg/distributor/push.go

@@ -13,6 +13,7 @@ import (
 	"math/rand"
 	"net/http"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/go-kit/log"
@@ -149,49 +150,63 @@ func handler(
 				logger = utillog.WithSourceIPs(source, logger)
 			}
 		}
-		supplier := func() (*mimirpb.WriteRequest, func(), error) {
-			rb := util.NewRequestBuffers(requestBufferPool)
-			var req mimirpb.PreallocWriteRequest
 
-			userID, err := tenant.TenantID(ctx)
-			if err != nil && !errors.Is(err, user.ErrNoOrgID) { // ignore user.ErrNoOrgID
-				return nil, nil, errors.Wrap(err, "failed to get tenant ID")
-			}
-
-			// userID might be empty if none was in the ctx, in this case just use the default setting.
-			if limits.MaxGlobalExemplarsPerUser(userID) == 0 {
-				// The user is not allowed to send exemplars, so there is no need to unmarshal them.
-				// Optimization to avoid the allocations required for unmarshaling exemplars.
-				req.SkipUnmarshalingExemplars = true
+		var supplier supplierFunc
+		isRW2, err := isRemoteWrite2(r)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+		}
+		if isRW2 {
+			supplier = func() (*mimirpb.WriteRequest, func(), error) {
+				// Return 415 Unsupported Media Type for remote-write v2 requests for now. This is not retryable
+				// unless the client switches to remote-write v1.
+				return nil, nil, httpgrpc.Error(http.StatusUnsupportedMediaType, "remote-write v2 is not supported")
 			}
+		} else {
+			supplier = func() (*mimirpb.WriteRequest, func(), error) {
+				rb := util.NewRequestBuffers(requestBufferPool)
+				var req mimirpb.PreallocWriteRequest
+
+				userID, err := tenant.TenantID(ctx)
+				if err != nil && !errors.Is(err, user.ErrNoOrgID) { // ignore user.ErrNoOrgID
+					return nil, nil, errors.Wrap(err, "failed to get tenant ID")
+				}
 
-			if err := parser(ctx, r, maxRecvMsgSize, rb, &req, logger); err != nil {
-				// Check for httpgrpc error, default to client error if parsing failed
-				if _, ok := httpgrpc.HTTPResponseFromError(err); !ok {
-					err = httpgrpc.Error(http.StatusBadRequest, err.Error())
+				// userID might be empty if none was in the ctx, in this case just use the default setting.
+				if limits.MaxGlobalExemplarsPerUser(userID) == 0 {
+					// The user is not allowed to send exemplars, so there is no need to unmarshal them.
+					// Optimization to avoid the allocations required for unmarshaling exemplars.
+					req.SkipUnmarshalingExemplars = true
 				}
 
-				rb.CleanUp()
-				return nil, nil, err
-			}
+				if err := parser(ctx, r, maxRecvMsgSize, rb, &req, logger); err != nil {
+					// Check for httpgrpc error, default to client error if parsing failed
+					if _, ok := httpgrpc.HTTPResponseFromError(err); !ok {
+						err = httpgrpc.Error(http.StatusBadRequest, err.Error())
+					}
 
-			if allowSkipLabelNameValidation {
-				req.SkipLabelValidation = req.SkipLabelValidation && r.Header.Get(SkipLabelNameValidationHeader) == "true"
-			} else {
-				req.SkipLabelValidation = false
-			}
+					rb.CleanUp()
+					return nil, nil, err
+				}
 
-			if allowSkipLabelCountValidation {
-				req.SkipLabelCountValidation = req.SkipLabelCountValidation && r.Header.Get(SkipLabelCountValidationHeader) == "true"
-			} else {
-				req.SkipLabelCountValidation = false
-			}
+				if allowSkipLabelNameValidation {
+					req.SkipLabelValidation = req.SkipLabelValidation && r.Header.Get(SkipLabelNameValidationHeader) == "true"
+				} else {
+					req.SkipLabelValidation = false
+				}
 
-			cleanup := func() {
-				mimirpb.ReuseSlice(req.Timeseries)
-				rb.CleanUp()
+				if allowSkipLabelCountValidation {
+					req.SkipLabelCountValidation = req.SkipLabelCountValidation && r.Header.Get(SkipLabelCountValidationHeader) == "true"
+				} else {
+					req.SkipLabelCountValidation = false
+				}
+
+				cleanup := func() {
+					mimirpb.ReuseSlice(req.Timeseries)
+					rb.CleanUp()
+				}
+				return &req.WriteRequest, cleanup, nil
 			}
-			return &req.WriteRequest, cleanup, nil
 		}
 		req := newRequest(supplier)
 		if err := push(ctx, req); err != nil {
@@ -226,6 +241,40 @@ func handler(
 	})
 }
 
+func isRemoteWrite2(r *http.Request) (bool, error) {
+	const appProtoContentType = "application/x-protobuf"
+
+	contentType := r.Header.Get("Content-Type")
+	if contentType == "" {
+		// If the content type is not set, we assume it is remote write v1.
+		return false, nil
+	}
+	parts := strings.Split(contentType, ";")
+	if parts[0] != appProtoContentType {
+		return false, fmt.Errorf("expected %v as the first (media) part, got %v content-type", appProtoContentType, contentType)
+	}
+
+	// Parse potential https://www.rfc-editor.org/rfc/rfc9110#parameter
+	for _, p := range parts[1:] {
+		pair := strings.Split(p, "=")
+		if len(pair) != 2 {
+			return false, fmt.Errorf("as per https://www.rfc-editor.org/rfc/rfc9110#parameter expected parameters to be key-values, got %v in %v content-type", p, contentType)
+		}
+		if pair[0] == "proto" {
+			switch pair[1] {
+			case "prometheus.WriteRequest":
+				return false, nil
+			case "io.prometheus.write.v2.Request":
+				return true, nil
+			default:
+				return false, fmt.Errorf("got %v content type; expected prometheus.WriteRequest or io.prometheus.write.v2.Request", contentType)
+			}
+		}
+	}
+	// No "proto=" parameter, assuming v1.
+	return false, nil
+}
+
 func calculateRetryAfter(retryAttemptHeader string, minBackoff, maxBackoff time.Duration) string {
 	const jitterFactor = 0.5