From f411a0795af67630a0a70a88ce64fa071de50a56 Mon Sep 17 00:00:00 2001
From: Callum Styan <callumstyan@gmail.com>
Date: Mon, 28 Oct 2024 10:45:28 -0700
Subject: [PATCH] fix: fix bug in query result marshaling for invalid utf8
 characters (#14585)

Signed-off-by: Callum Styan <callumstyan@gmail.com>
---
 pkg/util/marshal/query.go      | 16 ++++++++++++++++
 pkg/util/marshal/query_test.go | 19 +++++++++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 pkg/util/marshal/query_test.go

diff --git a/pkg/util/marshal/query.go b/pkg/util/marshal/query.go
index 7f06cd6ebe..cbf3b40d94 100644
--- a/pkg/util/marshal/query.go
+++ b/pkg/util/marshal/query.go
@@ -1,8 +1,11 @@
 package marshal
 
 import (
+	"bytes"
 	"fmt"
 	"strconv"
+	"strings"
+	"unicode/utf8"
 	"unsafe"
 
 	jsoniter "github.com/json-iterator/go"
@@ -20,6 +23,16 @@ import (
 	"github.com/grafana/loki/v3/pkg/util/httpreq"
 )
 
+var (
+	// The rune error replacement is rejected by Prometheus hence replacing them with space.
+	removeInvalidUtf = func(r rune) rune {
+		if r == utf8.RuneError {
+			return 32 // rune value for space
+		}
+		return r
+	}
+)
+
 // NewResultValue constructs a ResultValue from a promql.Value
 func NewResultValue(v parser.Value) (loghttp.ResultValue, error) {
 	var err error
@@ -77,6 +90,9 @@ func NewStreams(s logqlmodel.Streams) (loghttp.Streams, error) {
 	ret := make([]loghttp.Stream, len(s))
 
 	for i, stream := range s {
+		if strings.ContainsRune(stream.Labels, utf8.RuneError) {
+			stream.Labels = string(bytes.Map(removeInvalidUtf, []byte(stream.Labels)))
+		}
 		ret[i], err = NewStream(stream)
 
 		if err != nil {
diff --git a/pkg/util/marshal/query_test.go b/pkg/util/marshal/query_test.go
new file mode 100644
index 0000000000..15764760c2
--- /dev/null
+++ b/pkg/util/marshal/query_test.go
@@ -0,0 +1,19 @@
+package marshal
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/v3/pkg/logqlmodel"
+)
+
+func TestNewStreams(t *testing.T) {
+	s, err := NewStreams(logqlmodel.Streams{
+		{
+			Labels: "{asdf=\"�\"}",
+		},
+	})
+	require.NoError(t, err)
+	require.Equal(t, " ", s[0].Labels["asdf"], "expected only a space for label who only contained invalid UTF8 rune")
+}