fix: trace query RBAC: do not try to unmarshal unsupported content types

api/traces/v1/trace_rbac.go

@@ -85,25 +85,31 @@ func WithTraceQLNamespaceSelectAndForbidOtherAPIs(enabled bool) func(http.Handle
 }
 
 func unmarshal(response *http.Response, body []byte, pb proto.Message) error {
-	switch response.Header.Get(HeaderContentType) {
+	contentTypeHeader := response.Header.Get(HeaderContentType)
+	switch contentTypeHeader {
 	case ContentTypeProtobuf:
 		return proto.Unmarshal(body, pb)
-	default:
+	case ContentTypeJSON:
 		return (&jsonpb.Unmarshaler{}).Unmarshal(bytes.NewReader(body), pb)
+	default:
+		return fmt.Errorf("unsupported content type: %q", contentTypeHeader)
 	}
 }
 
 func marshal(response *http.Response, buf *bytes.Buffer, pb proto.Message) error {
-	switch response.Header.Get(HeaderContentType) {
+	contentTypeHeader := response.Header.Get(HeaderContentType)
+	switch contentTypeHeader {
 	case ContentTypeProtobuf:
 		b, err := proto.Marshal(pb)
 		if err != nil {
 			return err
 		}
 		buf.Write(b)
 		return nil
-	default:
+	case ContentTypeJSON:
 		return (&jsonpb.Marshaler{}).Marshal(buf, pb)
+	default:
+		return fmt.Errorf("unsupported content type: %q", contentTypeHeader)
 	}
 }
 
@@ -128,14 +134,18 @@ func responseRBACModifier(log log.Logger) func(response *http.Response) error {
 				responseBuffer := &bytes.Buffer{}
 				switch {
 				case routeQueryV1.MatchString(request.URL.Path):
+					contentTypeHeader := response.Header.Get(HeaderContentType)
+
 					// do not use unmarshal here, because we must use
 					// UnmarshalFromJSONV1 instead of (&jsonpb.Unmarshaler{}).Unmarshal
 					trace := &tempopb.Trace{}
-					switch response.Header.Get(HeaderContentType) {
+					switch contentTypeHeader {
 					case ContentTypeProtobuf:
 						err = proto.Unmarshal(b, trace)
-					default:
+					case ContentTypeJSON:
 						err = tempopb.UnmarshalFromJSONV1(b, trace)
+					default:
+						return fmt.Errorf("unsupported content type: %q", contentTypeHeader)
 					}
 					if err != nil {
 						return err
@@ -145,21 +155,23 @@ func responseRBACModifier(log log.Logger) func(response *http.Response) error {
 
 					// do not use marshal here, because we must use
 					// MarshalToJSONV1 instead of (&jsonpb.Marshaler{}).Marshal
-					switch response.Header.Get(HeaderContentType) {
+					switch contentTypeHeader {
 					case ContentTypeProtobuf:
 						var out []byte
 						out, err = proto.Marshal(trace)
 						if err != nil {
 							return err
 						}
 						responseBuffer = bytes.NewBuffer(out)
-					default:
+					case ContentTypeJSON:
 						var traceResponseBody []byte
 						traceResponseBody, err = tempopb.MarshalToJSONV1(trace)
 						if err != nil {
 							return err
 						}
 						responseBuffer = bytes.NewBuffer(traceResponseBody)
+					default:
+						return fmt.Errorf("unsupported content type: %q", contentTypeHeader)
 					}
 
 				case routeQueryV2.MatchString(request.URL.Path):
@@ -209,6 +221,10 @@ func responseRBACModifier(log log.Logger) func(response *http.Response) error {
 }
 
 func traceRBAC(allowedNamespaces map[string]bool, trace *tempopb.Trace) *tempopb.Trace {
+	if trace == nil {
+		return nil
+	}
+
 	for _, rs := range trace.ResourceSpans {
 		notAllowedNamespace := ""
 		missingNamespaceAttribute := true

api/traces/v1/trace_rbac_test.go

@@ -564,7 +564,7 @@ func contextWithAllowedNamespaces(t *testing.T, namespaces []string) context.Con
 
 func makeResponse(ctx context.Context, statusCode int, path string, body string, header http.Header) *http.Response {
 	if header == nil {
-		header = http.Header{}
+		header = http.Header{HeaderContentType: []string{ContentTypeJSON}}
 	}
 	return &http.Response{
 		StatusCode: statusCode,
@@ -823,7 +823,7 @@ func TestResponseRBACModifier(t *testing.T) {
       ]
     }
   ]
-}`, http.Header{"Content-Encoding": []string{"gzip"}})
+}`, http.Header{"Content-Encoding": []string{"gzip"}, HeaderContentType: []string{ContentTypeJSON}})
 
 		require.NoError(t, modifier(resp))
 
@@ -833,6 +833,64 @@ func TestResponseRBACModifier(t *testing.T) {
 	})
 }
 
+func TestMarshalUnmarshalContentType(t *testing.T) {
+	tests := []struct {
+		name        string
+		contentType string
+		wantErr     bool
+	}{
+		{"protobuf", ContentTypeProtobuf, false},
+		{"json", ContentTypeJSON, false},
+		{"unsupported grafana llm json", "application/vnd.grafana.llm+json", true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			response := &http.Response{
+				Header: http.Header{HeaderContentType: []string{tt.contentType}},
+			}
+
+			original := &tempopb.TraceByIDResponse{
+				Trace: &tempopb.Trace{
+					ResourceSpans: []*tracev1.ResourceSpans{
+						{
+							Resource: &resourcev1.Resource{
+								Attributes: []*commonv1.KeyValue{
+									createStringAttribute("k8s.namespace.name", "example"),
+								},
+							},
+						},
+					},
+				},
+			}
+
+			buf := &bytes.Buffer{}
+			err := marshal(response, buf, original)
+			if tt.wantErr {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), "unsupported content type")
+				assert.Contains(t, err.Error(), tt.contentType)
+			} else {
+				require.NoError(t, err)
+			}
+
+			unmarshalled := &tempopb.TraceByIDResponse{}
+			err = unmarshal(response, buf.Bytes(), unmarshalled)
+			if tt.wantErr {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), "unsupported content type")
+				assert.Contains(t, err.Error(), tt.contentType)
+			} else {
+				require.NoError(t, err)
+			}
+
+			if !tt.wantErr {
+				assert.Equal(t, original, unmarshalled)
+			}
+		})
+	}
+}
+
 func makeProtobufResponse(ctx context.Context, statusCode int, path string, pb proto.Message) *http.Response {
 	b, err := proto.Marshal(pb)
 	if err != nil {