Data Sources: Add QueryData OAuth & cookie forwarding middleware (#50466)

pkg/infra/httpclient/httpclientprovider/forwarded_cookie_middleware_test.go

@@ -0,0 +1,65 @@
+package httpclientprovider_test
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/infra/httpclient/httpclientprovider"
+	"github.com/stretchr/testify/require"
+)
+
+func TestForwardedCookiesMiddleware(t *testing.T) {
+	tcs := []struct {
+		desc                 string
+		allowedCookies       []string
+		expectedCookieHeader string
+	}{
+		{
+			desc:                 "With nil allowedCookies should not populate Cookie header",
+			allowedCookies:       nil,
+			expectedCookieHeader: "",
+		},
+		{
+			desc:                 "With empty allowed cookies should not populate Cookie header",
+			allowedCookies:       []string{},
+			expectedCookieHeader: "",
+		},
+		{
+			desc:                 "When provided with allowed cookies should populate Cookie header",
+			allowedCookies:       []string{"c1", "c3"},
+			expectedCookieHeader: "c1=1; c3=3",
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			ctx := &testContext{}
+			finalRoundTripper := ctx.createRoundTripper()
+			forwarded := []*http.Cookie{
+				{Name: "c1", Value: "1"},
+				{Name: "c2", Value: "2"},
+				{Name: "c3", Value: "3"},
+			}
+			mw := httpclientprovider.ForwardedCookiesMiddleware(forwarded, tc.allowedCookies)
+			opts := httpclient.Options{}
+			rt := mw.CreateMiddleware(opts, finalRoundTripper)
+			require.NotNil(t, rt)
+			middlewareName, ok := mw.(httpclient.MiddlewareName)
+			require.True(t, ok)
+			require.Equal(t, "forwarded-cookies", middlewareName.MiddlewareName())
+
+			req, err := http.NewRequest(http.MethodGet, "http://", nil)
+			require.NoError(t, err)
+			res, err := rt.RoundTrip(req)
+			require.NoError(t, err)
+			require.NotNil(t, res)
+			if res.Body != nil {
+				require.NoError(t, res.Body.Close())
+			}
+			require.Len(t, ctx.callChain, 1)
+			require.ElementsMatch(t, []string{"final"}, ctx.callChain)
+			require.Equal(t, tc.expectedCookieHeader, ctx.req.Header.Get("Cookie"))
+		})
+	}
+}

pkg/infra/httpclient/httpclientprovider/forwarded_cookies_middleware.go

@@ -0,0 +1,24 @@
+package httpclientprovider
+
+import (
+	"net/http"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/util/proxyutil"
+)
+
+const ForwardedCookiesMiddlewareName = "forwarded-cookies"
+
+// ForwardedCookiesMiddleware middleware that sets Cookie header on the
+// outgoing request, if forwarded cookies configured/provided.
+func ForwardedCookiesMiddleware(forwardedCookies []*http.Cookie, allowedCookies []string) httpclient.Middleware {
+	return httpclient.NamedMiddlewareFunc(ForwardedCookiesMiddlewareName, func(opts httpclient.Options, next http.RoundTripper) http.RoundTripper {
+		return httpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+			for _, cookie := range forwardedCookies {
+				req.AddCookie(cookie)
+			}
+			proxyutil.ClearCookieHeader(req, allowedCookies)
+			return next.RoundTrip(req)
+		})
+	})
+}

pkg/infra/httpclient/httpclientprovider/forwarded_oauth_identity_middleware.go

@@ -0,0 +1,31 @@
+package httpclientprovider
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"golang.org/x/oauth2"
+)
+
+const ForwardedOAuthIdentityMiddlewareName = "forwarded-oauth-identity"
+
+// ForwardedOAuthIdentityMiddleware middleware that sets Authorization/X-ID-Token
+// headers on the outgoing request if an OAuth Token is provided
+func ForwardedOAuthIdentityMiddleware(token *oauth2.Token) httpclient.Middleware {
+	return httpclient.NamedMiddlewareFunc(ForwardedOAuthIdentityMiddlewareName, func(opts httpclient.Options, next http.RoundTripper) http.RoundTripper {
+		if token == nil {
+			return next
+		}
+		return httpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+			req.Header.Set("Authorization", fmt.Sprintf("%s %s", token.Type(), token.AccessToken))
+
+			idToken, ok := token.Extra("id_token").(string)
+			if ok && idToken != "" {
+				req.Header.Set("X-ID-Token", idToken)
+			}
+
+			return next.RoundTrip(req)
+		})
+	})
+}

pkg/infra/httpclient/httpclientprovider/forwarded_oauth_identity_middleware_test.go

@@ -0,0 +1,82 @@
+package httpclientprovider_test
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/infra/httpclient/httpclientprovider"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/oauth2"
+)
+
+func TestForwardedOAuthIdentityMiddleware(t *testing.T) {
+	at := &oauth2.Token{
+		AccessToken: "access-token",
+	}
+	tcs := []struct {
+		desc                        string
+		token                       *oauth2.Token
+		expectedAuthorizationHeader string
+		expectedIDTokenHeader       string
+	}{
+		{
+			desc:                        "With nil token should not populate Cookie headers",
+			token:                       nil,
+			expectedAuthorizationHeader: "",
+			expectedIDTokenHeader:       "",
+		},
+		{
+			desc:                        "With access token set should populate Authorization header",
+			token:                       at,
+			expectedAuthorizationHeader: "Bearer access-token",
+			expectedIDTokenHeader:       "",
+		},
+		{
+			desc:                        "With Authorization and X-ID-Token header set should populate Authorization and X-Id-Token header",
+			token:                       at.WithExtra(map[string]interface{}{"id_token": "id-token"}),
+			expectedAuthorizationHeader: "Bearer access-token",
+			expectedIDTokenHeader:       "id-token",
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			ctx := &testContext{}
+			finalRoundTripper := ctx.createRoundTripper()
+			mw := httpclientprovider.ForwardedOAuthIdentityMiddleware(tc.token)
+			opts := httpclient.Options{}
+			rt := mw.CreateMiddleware(opts, finalRoundTripper)
+			require.NotNil(t, rt)
+			middlewareName, ok := mw.(httpclient.MiddlewareName)
+			require.True(t, ok)
+			require.Equal(t, "forwarded-oauth-identity", middlewareName.MiddlewareName())
+
+			req, err := http.NewRequest(http.MethodGet, "http://", nil)
+			require.NoError(t, err)
+			res, err := rt.RoundTrip(req)
+			require.NoError(t, err)
+			require.NotNil(t, res)
+			if res.Body != nil {
+				require.NoError(t, res.Body.Close())
+			}
+			require.Len(t, ctx.callChain, 1)
+			require.ElementsMatch(t, []string{"final"}, ctx.callChain)
+			require.Equal(t, tc.expectedAuthorizationHeader, ctx.req.Header.Get("Authorization"))
+			require.Equal(t, tc.expectedIDTokenHeader, ctx.req.Header.Get("X-ID-Token"))
+		})
+	}
+}
+
+type testContext struct {
+	callChain []string
+	req       *http.Request
+}
+
+func (c *testContext) createRoundTripper() http.RoundTripper {
+	return httpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		c.callChain = append(c.callChain, "final")
+		c.req = req
+		return &http.Response{StatusCode: http.StatusOK}, nil
+	})
+}

pkg/infra/httpclient/httpclientprovider/http_client_provider.go

@@ -5,12 +5,11 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/grafana/grafana/pkg/models"
-
 	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/metrics/metricutil"
 	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/mwitkow/go-conntrack"
 )

pkg/services/query/query.go

@@ -10,6 +10,7 @@ import (
 	"github.com/grafana/grafana/pkg/api/dtos"
 	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/expr"
+	"github.com/grafana/grafana/pkg/infra/httpclient/httpclientprovider"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
@@ -22,6 +23,7 @@ import (
 	"github.com/grafana/grafana/pkg/util/proxyutil"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
 )
 
 const (
@@ -136,6 +138,13 @@ func (s *Service) handleQueryData(ctx context.Context, user *models.SignedInUser
 		Queries: []backend.DataQuery{},
 	}
 
+	middlewares := []httpclient.Middleware{}
+	if parsedReq.httpRequest != nil {
+		middlewares = append(middlewares,
+			httpclientprovider.ForwardedCookiesMiddleware(parsedReq.httpRequest.Cookies(), ds.AllowedCookies()),
+		)
+	}
+
 	if s.oAuthTokenService.IsOAuthPassThruEnabled(ds) {
 		if token := s.oAuthTokenService.GetCurrentOAuthToken(ctx, user); token != nil {
 			req.Headers["Authorization"] = fmt.Sprintf("%s %s", token.Type(), token.AccessToken)
@@ -144,6 +153,7 @@ func (s *Service) handleQueryData(ctx context.Context, user *models.SignedInUser
 			if ok && idToken != "" {
 				req.Headers["X-ID-Token"] = idToken
 			}
+			middlewares = append(middlewares, httpclientprovider.ForwardedOAuthIdentityMiddleware(token))
 		}
 	}
 
@@ -162,6 +172,8 @@ func (s *Service) handleQueryData(ctx context.Context, user *models.SignedInUser
 		req.Queries = append(req.Queries, q.query)
 	}
 
+	ctx = httpclient.WithContextualMiddleware(ctx, middlewares...)
+
 	return s.pluginClient.QueryData(ctx, req)
 }