DataProxy: Restore Set-Cookie header after proxy request (#16838)

If Grafana rotates the user's auth token during a request to the data 
source proxy it will set the Set-Cookie header with new auth token in 
response before proxying the request to the datasource.
Before this fix the Set-Cookie response header was cleared after the 
proxied request was finished to make sure that proxied datasources 
cannot affect cookies in users browsers. This had the consequence 
of accidentally also clearing the new auth token set in Set-Cookie 
header.
With this fix the original Set-Cookie value in response header is now 
restored after the proxied datasource request is finished. The existing
logic of clearing Set-Cookie response header from proxied request 
have been left intact.

Fixes #16757

pkg/api/pluginproxy/ds_proxy.go

@@ -101,8 +101,14 @@ func (proxy *DataSourceProxy) HandleRequest() {
 		opentracing.HTTPHeaders,
 		opentracing.HTTPHeadersCarrier(proxy.ctx.Req.Request.Header))
 
+	originalSetCookie := proxy.ctx.Resp.Header().Get("Set-Cookie")
+
 	reverseProxy.ServeHTTP(proxy.ctx.Resp, proxy.ctx.Req.Request)
 	proxy.ctx.Resp.Header().Del("Set-Cookie")
+
+	if originalSetCookie != "" {
+		proxy.ctx.Resp.Header().Set("Set-Cookie", originalSetCookie)
+	}
 }
 
 func (proxy *DataSourceProxy) addTraceFromHeaderValue(span opentracing.Span, headerName string, tagName string) {

pkg/api/pluginproxy/ds_proxy_test.go

@@ -3,13 +3,15 @@ package pluginproxy
 import (
 	"bytes"
 	"fmt"
-	"github.com/grafana/grafana/pkg/components/securejsondata"
 	"io/ioutil"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"testing"
 	"time"
 
+	"github.com/grafana/grafana/pkg/components/securejsondata"
+
 	"golang.org/x/oauth2"
 	macaron "gopkg.in/macaron.v1"
 
@@ -496,9 +498,73 @@ func TestDSRouteRule(t *testing.T) {
 				runDatasourceAuthTest(test)
 			}
 		})
+
+		Convey("HandleRequest()", func() {
+			backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				http.SetCookie(w, &http.Cookie{Name: "flavor", Value: "chocolateChip"})
+				w.WriteHeader(200)
+				w.Write([]byte("I am the backend"))
+			}))
+			defer backend.Close()
+
+			plugin := &plugins.DataSourcePlugin{}
+			ds := &m.DataSource{Url: backend.URL, Type: m.DS_GRAPHITE}
+
+			responseRecorder := &CloseNotifierResponseRecorder{
+				ResponseRecorder: httptest.NewRecorder(),
+			}
+			defer responseRecorder.Close()
+
+			setupCtx := func(fn func(http.ResponseWriter)) *m.ReqContext {
+				responseWriter := macaron.NewResponseWriter("GET", responseRecorder)
+				if fn != nil {
+					fn(responseWriter)
+				}
+
+				return &m.ReqContext{
+					SignedInUser: &m.SignedInUser{},
+					Context: &macaron.Context{
+						Req: macaron.Request{
+							Request: httptest.NewRequest("GET", "/render", nil),
+						},
+						Resp: responseWriter,
+					},
+				}
+			}
+
+			Convey("When response header Set-Cookie is not set should remove proxied Set-Cookie header", func() {
+				ctx := setupCtx(nil)
+				proxy := NewDataSourceProxy(ds, plugin, ctx, "/render", &setting.Cfg{})
+				proxy.HandleRequest()
+				So(proxy.ctx.Resp.Header().Get("Set-Cookie"), ShouldBeEmpty)
+			})
+
+			Convey("When response header Set-Cookie is set should remove proxied Set-Cookie header and restore the original Set-Cookie header", func() {
+				ctx := setupCtx(func(w http.ResponseWriter) {
+					w.Header().Set("Set-Cookie", "important_cookie=important_value")
+				})
+				proxy := NewDataSourceProxy(ds, plugin, ctx, "/render", &setting.Cfg{})
+				proxy.HandleRequest()
+				So(proxy.ctx.Resp.Header().Get("Set-Cookie"), ShouldEqual, "important_cookie=important_value")
+			})
+		})
 	})
 }
 
+type CloseNotifierResponseRecorder struct {
+	*httptest.ResponseRecorder
+	closeChan chan bool
+}
+
+func (r *CloseNotifierResponseRecorder) CloseNotify() <-chan bool {
+	r.closeChan = make(chan bool)
+	return r.closeChan
+}
+
+func (r *CloseNotifierResponseRecorder) Close() {
+	close(r.closeChan)
+}
+
 // getDatasourceProxiedRequest is a helper for easier setup of tests based on global config and ReqContext.
 func getDatasourceProxiedRequest(ctx *m.ReqContext, cfg *setting.Cfg) *http.Request {
 	plugin := &plugins.DataSourcePlugin{}