api: use alerting headers (#95118)

* api: use alerting headers

* improve code, add integration test

* better comment

* fixed test

* merged tests

pkg/services/pluginsintegration/clientmiddleware/usealertingheaders_middleware.go

@@ -0,0 +1,57 @@
+package clientmiddleware
+
+import (
+	"context"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana/pkg/services/contexthandler"
+	ngalertmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
+)
+
+func NewUseAlertHeadersMiddleware() backend.HandlerMiddleware {
+	return backend.HandlerMiddlewareFunc(func(next backend.Handler) backend.Handler {
+		return &UseAlertHeadersMiddleware{
+			BaseHandler: backend.NewBaseHandler(next),
+		}
+	})
+}
+
+type UseAlertHeadersMiddleware struct {
+	backend.BaseHandler
+}
+
+var alertHeaders = []string{
+	"X-Rule-Name",
+	"X-Rule-Folder",
+	"X-Rule-Source",
+	"X-Rule-Type",
+	"X-Rule-Version",
+	ngalertmodels.FromAlertHeaderName,
+}
+
+func applyAlertHeaders(ctx context.Context, req *backend.QueryDataRequest) {
+	reqCtx := contexthandler.FromContext(ctx)
+	if reqCtx == nil || reqCtx.Req == nil {
+		return
+	}
+	incomingHeaders := reqCtx.Req.Header
+
+	for _, key := range alertHeaders {
+		incomingValue := incomingHeaders.Get(key)
+		if incomingValue != "" {
+			// FromAlert must be set directly, because we need
+			// to keep the incorrect capitalization for backwards-compatibility
+			// reasons. otherwise Go would normalize it to "Fromalert"
+			if key == ngalertmodels.FromAlertHeaderName {
+				req.Headers[key] = incomingValue
+			} else {
+				req.SetHTTPHeader(key, incomingValue)
+			}
+		}
+	}
+}
+
+func (m *UseAlertHeadersMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	applyAlertHeaders(ctx, req)
+	return m.BaseHandler.QueryData(ctx, req)
+}

pkg/services/pluginsintegration/clientmiddleware/usealertingheaders_middleware_test.go

@@ -0,0 +1,70 @@
+package clientmiddleware
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/handlertest"
+	"github.com/grafana/grafana/pkg/services/user"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUserAlertingHeadersMiddleware(t *testing.T) {
+	testQueryDataReq := func(t *testing.T, req *http.Request) *backend.QueryDataRequest {
+		cdt := handlertest.NewHandlerMiddlewareTest(t,
+			WithReqContext(req, &user.SignedInUser{}),
+			handlertest.WithMiddlewares(NewUseAlertHeadersMiddleware()),
+		)
+
+		pluginCtx := backend.PluginContext{
+			DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
+		}
+
+		_, err := cdt.MiddlewareHandler.QueryData(req.Context(), &backend.QueryDataRequest{
+			PluginContext: pluginCtx,
+			Headers:       map[string]string{},
+		})
+		require.NoError(t, err)
+		return cdt.QueryDataReq
+	}
+
+	t.Run("Handle non-alerting case without problems", func(t *testing.T) {
+		req, err := http.NewRequest(http.MethodGet, "/some/thing", nil)
+		require.NoError(t, err)
+		outReq := testQueryDataReq(t, req)
+
+		// special marker
+		require.Equal(t, "", outReq.Headers["FromAlert"])
+
+		// the normal http headers
+		require.Equal(t, "", outReq.GetHTTPHeader("X-Rule-Name"))
+		require.Equal(t, "", outReq.GetHTTPHeader("X-Rule-Folder"))
+		require.Equal(t, "", outReq.GetHTTPHeader("X-Rule-Source"))
+		require.Equal(t, "", outReq.GetHTTPHeader("X-Rule-Type"))
+		require.Equal(t, "", outReq.GetHTTPHeader("X-Rule-Version"))
+	})
+
+	t.Run("Use Alerting headers when they exist", func(t *testing.T) {
+		req, err := http.NewRequest(http.MethodGet, "/some/thing", nil)
+		require.NoError(t, err)
+		req.Header.Set("Fromalert", "true")
+		req.Header.Set("X-Rule-Name", "n1")
+		req.Header.Set("X-Rule-Folder", "f1")
+		req.Header.Set("X-Rule-Source", "s1")
+		req.Header.Set("X-Rule-Type", "t1")
+		req.Header.Set("X-Rule-Version", "v1")
+
+		outReq := testQueryDataReq(t, req)
+
+		// special marker
+		require.Equal(t, "true", outReq.Headers["FromAlert"])
+
+		// normal http headers
+		require.Equal(t, "n1", outReq.GetHTTPHeader("X-Rule-Name"))
+		require.Equal(t, "f1", outReq.GetHTTPHeader("X-Rule-Folder"))
+		require.Equal(t, "s1", outReq.GetHTTPHeader("X-Rule-Source"))
+		require.Equal(t, "t1", outReq.GetHTTPHeader("X-Rule-Type"))
+		require.Equal(t, "v1", outReq.GetHTTPHeader("X-Rule-Version"))
+	})
+}

pkg/services/pluginsintegration/pluginsintegration.go

@@ -189,6 +189,7 @@ func CreateMiddlewares(cfg *setting.Cfg, oAuthTokenService oauthtoken.OAuthToken
 		clientmiddleware.NewCookiesMiddleware(skipCookiesNames),
 		clientmiddleware.NewCachingMiddlewareWithFeatureManager(cachingService, features),
 		clientmiddleware.NewForwardIDMiddleware(),
+		clientmiddleware.NewUseAlertHeadersMiddleware(),
 	)
 
 	if cfg.SendUserHeader {

pkg/tests/api/loki/loki_test.go

@@ -91,7 +91,17 @@ func TestIntegrationLoki(t *testing.T) {
 		require.NoError(t, err)
 		u := fmt.Sprintf("http://admin:admin@%s/api/ds/query", grafanaListeningAddr)
 		// nolint:gosec
-		resp, err := http.Post(u, "application/json", buf1)
+		req, err := http.NewRequest("POST", u, buf1)
+		require.NoError(t, err)
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("X-Unspported-Header", "uh")
+		req.Header.Set("Fromalert", "true")
+		req.Header.Set("X-Rule-Name", "n1")
+		req.Header.Set("X-Rule-Folder", "f1")
+		req.Header.Set("X-Rule-Source", "s1")
+		req.Header.Set("X-Rule-Type", "t1")
+		req.Header.Set("X-Rule-Version", "v1")
+		resp, err := http.DefaultClient.Do(req)
 		require.NoError(t, err)
 		require.Equal(t, http.StatusBadRequest, resp.StatusCode)
 		t.Cleanup(func() {
@@ -107,6 +117,13 @@ func TestIntegrationLoki(t *testing.T) {
 		require.True(t, ok)
 		require.Equal(t, "basicAuthUser", username)
 		require.Equal(t, "basicAuthPassword", pwd)
+		require.Equal(t, "", outgoingRequest.Header.Get("X-Unspported-Header"))
+		require.Equal(t, "true", outgoingRequest.Header.Get("Fromalert"))
+		require.Equal(t, "n1", outgoingRequest.Header.Get("X-Rule-Name"))
+		require.Equal(t, "f1", outgoingRequest.Header.Get("X-Rule-Folder"))
+		require.Equal(t, "s1", outgoingRequest.Header.Get("X-Rule-Source"))
+		require.Equal(t, "t1", outgoingRequest.Header.Get("X-Rule-Type"))
+		require.Equal(t, "v1", outgoingRequest.Header.Get("X-Rule-Version"))
 	})
 
 	t.Run("should forward `X-Dashboard-Title` header but no `X-Panel-Title`", func(t *testing.T) {