From fd619b693196c3614310b25a39efa956d0b31b71 Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Fri, 22 Jan 2021 14:12:40 +0100
Subject: [PATCH] Fixes tail api marshalling for v1. (#3211)

Introduce a bug by removing the default marshalling (#3163) but the tail api was using the default json. This fixes it by forcing the usage of jsoniter package.
Added a missing test so that it doesn't happen again.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
---
 pkg/logql/marshal/marshal.go      | 15 ++++++++++++---
 pkg/logql/marshal/marshal_test.go | 22 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/pkg/logql/marshal/marshal.go b/pkg/logql/marshal/marshal.go
index 1964e6bc8755..f7973b3079bc 100644
--- a/pkg/logql/marshal/marshal.go
+++ b/pkg/logql/marshal/marshal.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/gorilla/websocket"
 	json "github.com/json-iterator/go"
+	jsoniter "github.com/json-iterator/go"
 
 	"github.com/grafana/loki/pkg/loghttp"
 	legacy "github.com/grafana/loki/pkg/loghttp/legacy"
@@ -46,15 +47,23 @@ func WriteLabelResponseJSON(l logproto.LabelResponse, w io.Writer) error {
 	return json.NewEncoder(w).Encode(v1Response)
 }
 
+// WebsocketWriter knows how to write message to a websocket connection.
+type WebsocketWriter interface {
+	WriteMessage(int, []byte) error
+}
+
 // WriteTailResponseJSON marshals the legacy.TailResponse to v1 loghttp JSON and
 // then writes it to the provided connection.
-func WriteTailResponseJSON(r legacy.TailResponse, c *websocket.Conn) error {
+func WriteTailResponseJSON(r legacy.TailResponse, c WebsocketWriter) error {
 	v1Response, err := NewTailResponse(r)
 	if err != nil {
 		return err
 	}
-
-	return c.WriteJSON(v1Response)
+	data, err := jsoniter.Marshal(v1Response)
+	if err != nil {
+		return err
+	}
+	return c.WriteMessage(websocket.TextMessage, data)
 }
 
 // WriteSeriesResponseJSON marshals a logproto.SeriesResponse to v1 loghttp JSON and then
diff --git a/pkg/logql/marshal/marshal_test.go b/pkg/logql/marshal/marshal_test.go
index 6564ffe25e86..ee9c4d0e69af 100644
--- a/pkg/logql/marshal/marshal_test.go
+++ b/pkg/logql/marshal/marshal_test.go
@@ -521,3 +521,25 @@ func Benchmark_Encode(b *testing.B) {
 		}
 	}
 }
+
+type WebsocketWriterFunc func(int, []byte) error
+
+func (w WebsocketWriterFunc) WriteMessage(t int, d []byte) error { return w(t, d) }
+
+func Test_WriteTailResponseJSON(t *testing.T) {
+	require.NoError(t,
+		WriteTailResponseJSON(legacy.TailResponse{
+			Streams: []logproto.Stream{
+				{Labels: `{app="foo"}`, Entries: []logproto.Entry{{Timestamp: time.Unix(0, 1), Line: `foobar`}}},
+			},
+			DroppedEntries: []legacy.DroppedEntry{
+				{Timestamp: time.Unix(0, 2), Labels: `{app="dropped"}`},
+			},
+		},
+			WebsocketWriterFunc(func(i int, b []byte) error {
+				require.Equal(t, `{"streams":[{"stream":{"app":"foo"},"values":[["1","foobar"]]}],"dropped_entries":[{"timestamp":"2","labels":{"app":"dropped"}}]}`, string(b))
+				return nil
+			}),
+		),
+	)
+}