Plugins: Forward HTTP headers to plugin streaming calls (#101960)

* forward HTTP headers to plugin streaming calls

* fixup

* fix test based on feedback

* improve test assertions
This commit is contained in:
Will Browne
2025-03-12 10:37:07 +00:00
committed by GitHub
parent 8de665cbeb
commit 5a2cba7b01
8 changed files with 747 additions and 59 deletions

View File

@ -16,11 +16,68 @@ import (
)
func TestForwardIDMiddleware(t *testing.T) {
pluginContext := backend.PluginContext{
DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
}
t.Run("When not signed in", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
})
t.Run("Should set forwarded id header if present", func(t *testing.T) {
t.Run("And requests are for a datasource", func(t *testing.T) {
pluginContext := backend.PluginContext{
DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
}
t.Run("Should not set forwarded id header if not present for QueryData", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.QueryData(ctx, &backend.QueryDataRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Empty(t, cdt.QueryDataReq.GetHTTPHeaders())
})
t.Run("Should not set forwarded id header if not present for CallResource", func(t *testing.T) {
err := cdt.MiddlewareHandler.CallResource(ctx, &backend.CallResourceRequest{
PluginContext: pluginContext,
}, nopCallResourceSender)
require.NoError(t, err)
require.Empty(t, cdt.CallResourceReq.GetHTTPHeaders())
})
t.Run("Should not set forwarded id header if not present for CheckHealth", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.CheckHealth(ctx, &backend.CheckHealthRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Empty(t, cdt.CheckHealthReq.GetHTTPHeaders())
})
t.Run("Should not set forwarded id header if not present for SubscribeStream", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.SubscribeStream(ctx, &backend.SubscribeStreamRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Empty(t, cdt.SubscribeStreamReq.GetHTTPHeaders())
})
t.Run("Should not set forwarded id header if not present for PublishStream", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.PublishStream(ctx, &backend.PublishStreamRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Empty(t, cdt.PublishStreamReq.GetHTTPHeaders())
})
t.Run("Should not set forwarded id header if not present for RunStream", func(t *testing.T) {
err := cdt.MiddlewareHandler.RunStream(ctx, &backend.RunStreamRequest{
PluginContext: pluginContext,
}, &backend.StreamSender{})
require.NoError(t, err)
require.Empty(t, cdt.RunStreamReq.GetHTTPHeaders())
})
})
})
t.Run("When signed in", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
@ -28,47 +85,154 @@ func TestForwardIDMiddleware(t *testing.T) {
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
err := cdt.MiddlewareHandler.CallResource(ctx, &backend.CallResourceRequest{
PluginContext: pluginContext,
}, nopCallResourceSender)
require.NoError(t, err)
t.Run("And requests are for a datasource", func(t *testing.T) {
pluginContext := backend.PluginContext{
DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
}
require.Equal(t, "some-token", cdt.CallResourceReq.Headers[forwardIDHeaderName][0])
})
t.Run("Should set forwarded id header if present for QueryData", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.QueryData(ctx, &backend.QueryDataRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.QueryDataReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should not set forwarded id header if not present", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
t.Run("Should set forwarded id header if present for CallResource", func(t *testing.T) {
err := cdt.MiddlewareHandler.CallResource(ctx, &backend.CallResourceRequest{
PluginContext: pluginContext,
}, nopCallResourceSender)
require.NoError(t, err)
require.Equal(t, "some-token", cdt.CallResourceReq.GetHTTPHeader(forwardIDHeaderName))
})
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{},
t.Run("Should set forwarded id header if present for CheckHealth", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.CheckHealth(ctx, &backend.CheckHealthRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.CheckHealthReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header if present for SubscribeStream", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.SubscribeStream(ctx, &backend.SubscribeStreamRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.SubscribeStreamReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header if present for PublishStream", func(t *testing.T) {
_, err := cdt.MiddlewareHandler.PublishStream(ctx, &backend.PublishStreamRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.PublishStreamReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header if present for RunStream", func(t *testing.T) {
err := cdt.MiddlewareHandler.RunStream(ctx, &backend.RunStreamRequest{
PluginContext: pluginContext,
}, &backend.StreamSender{})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.RunStreamReq.GetHTTPHeader(forwardIDHeaderName))
})
})
err := cdt.MiddlewareHandler.CallResource(ctx, &backend.CallResourceRequest{
PluginContext: pluginContext,
}, nopCallResourceSender)
require.NoError(t, err)
t.Run("And requests are for an app", func(t *testing.T) {
pluginContext := backend.PluginContext{
AppInstanceSettings: &backend.AppInstanceSettings{},
}
require.Len(t, cdt.CallResourceReq.Headers[forwardIDHeaderName], 0)
})
t.Run("Should set forwarded id header to app plugin if present for QueryData", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
pluginContext = backend.PluginContext{
AppInstanceSettings: &backend.AppInstanceSettings{},
}
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
t.Run("Should set forwarded id header to app plugin if present", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
_, err := cdt.MiddlewareHandler.QueryData(ctx, &backend.QueryDataRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.QueryDataReq.GetHTTPHeader(forwardIDHeaderName))
})
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
t.Run("Should set forwarded id header to app plugin if present for CallResource", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
err := cdt.MiddlewareHandler.CallResource(ctx, &backend.CallResourceRequest{
PluginContext: pluginContext,
}, nopCallResourceSender)
require.NoError(t, err)
require.Equal(t, "some-token", cdt.CallResourceReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header to app plugin if present for CheckHealth", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
_, err := cdt.MiddlewareHandler.CheckHealth(ctx, &backend.CheckHealthRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.CheckHealthReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header to app plugin if present for SubscribeStream", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
_, err := cdt.MiddlewareHandler.SubscribeStream(ctx, &backend.SubscribeStreamRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.SubscribeStreamReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header to app plugin if present for PublishStream", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
_, err := cdt.MiddlewareHandler.PublishStream(ctx, &backend.PublishStreamRequest{
PluginContext: pluginContext,
})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.PublishStreamReq.GetHTTPHeader(forwardIDHeaderName))
})
t.Run("Should set forwarded id header to app plugin if present for RunStream", func(t *testing.T) {
cdt := handlertest.NewHandlerMiddlewareTest(t, handlertest.WithMiddlewares(NewForwardIDMiddleware()))
ctx := context.WithValue(context.Background(), ctxkey.Key{}, &contextmodel.ReqContext{
Context: &web.Context{Req: &http.Request{}},
SignedInUser: &user.SignedInUser{IDToken: "some-token"},
})
err := cdt.MiddlewareHandler.RunStream(ctx, &backend.RunStreamRequest{
PluginContext: pluginContext,
}, &backend.StreamSender{})
require.NoError(t, err)
require.Equal(t, "some-token", cdt.RunStreamReq.GetHTTPHeader(forwardIDHeaderName))
})
})
err := cdt.MiddlewareHandler.CallResource(ctx, &backend.CallResourceRequest{
PluginContext: pluginContext,
}, nopCallResourceSender)
require.NoError(t, err)
require.Equal(t, "some-token", cdt.CallResourceReq.Headers[forwardIDHeaderName][0])
})
}