From bfd5d3b16aec9612db1f7aef58fe2466cec02b82 Mon Sep 17 00:00:00 2001
From: Alexander Emelin <frvzmb@gmail.com>
Date: Thu, 6 May 2021 22:28:14 +0300
Subject: [PATCH] Live: streamId in path for ws push endpoint (#33786)

---
 pkg/services/live/live.go                |  3 ++-
 pkg/services/live/livecontext/context.go | 15 +++++++++++++++
 pkg/services/live/pushws/push.go         | 12 ++++--------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/pkg/services/live/live.go b/pkg/services/live/live.go
index e8bbfb27b1e..24010553e84 100644
--- a/pkg/services/live/live.go
+++ b/pkg/services/live/live.go
@@ -248,6 +248,7 @@ func (g *GrafanaLive) Init() error {
 		user := ctx.SignedInUser
 		newCtx := livecontext.SetContextSignedUser(ctx.Req.Context(), user)
 		newCtx = livecontext.SetContextValues(newCtx, ctx.Req.URL.Query())
+		newCtx = livecontext.SetContextStreamID(newCtx, ctx.Params(":streamId"))
 		r := ctx.Req.Request
 		r = r.WithContext(newCtx)
 		pushWSHandler.ServeHTTP(ctx.Resp, r)
@@ -258,7 +259,7 @@ func (g *GrafanaLive) Init() error {
 	}, middleware.ReqSignedIn)
 
 	g.RouteRegister.Group("/api/live", func(group routing.RouteRegister) {
-		group.Get("/push", g.pushWebsocketHandler)
+		group.Get("/push/:streamId", g.pushWebsocketHandler)
 	}, middleware.ReqOrgAdmin)
 
 	return nil
diff --git a/pkg/services/live/livecontext/context.go b/pkg/services/live/livecontext/context.go
index 2cb5bea71c7..e459294595c 100644
--- a/pkg/services/live/livecontext/context.go
+++ b/pkg/services/live/livecontext/context.go
@@ -38,3 +38,18 @@ func GetContextValues(ctx context.Context) (url.Values, bool) {
 	}
 	return nil, false
 }
+
+type streamIDContextKey struct{}
+
+func SetContextStreamID(ctx context.Context, streamID string) context.Context {
+	ctx = context.WithValue(ctx, streamIDContextKey{}, streamID)
+	return ctx
+}
+
+func GetContextStreamID(ctx context.Context) (string, bool) {
+	if val := ctx.Value(streamIDContextKey{}); val != nil {
+		values, ok := val.(string)
+		return values, ok
+	}
+	return "", false
+}
diff --git a/pkg/services/live/pushws/push.go b/pkg/services/live/pushws/push.go
index 31f9f2b80b6..8535e467cc2 100644
--- a/pkg/services/live/pushws/push.go
+++ b/pkg/services/live/pushws/push.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/live/convert"
+	"github.com/grafana/grafana/pkg/services/live/livecontext"
 	"github.com/grafana/grafana/pkg/services/live/managedstream"
 	"github.com/grafana/grafana/pkg/services/live/pushurl"
 
@@ -101,15 +102,10 @@ const (
 )
 
 func (s *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
-	var streamID string
-
-	streamID = r.Header.Get("X-Grafana-Live-Stream")
-	if streamID == "" {
-		streamID = r.URL.Query().Get("gf_live_stream")
-	}
-	if streamID == "" {
+	streamID, ok := livecontext.GetContextStreamID(r.Context())
+	if !ok || streamID == "" {
 		logger.Warn("Push request without stream ID")
-		rw.WriteHeader(http.StatusBadRequest)
+		rw.WriteHeader(http.StatusInternalServerError)
 		return
 	}