From 46d65739684cf7bfd14c5d34fe09613d85ef3b3d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=A1bor=20Farkas?= Date: Thu, 2 Jun 2022 11:52:27 +0200 Subject: [PATCH] loki: add cookie-handling functionality (#49978) --- pkg/tsdb/loki/api.go | 30 +++---- pkg/tsdb/loki/api_mock.go | 2 +- pkg/tsdb/loki/{oauth_test.go => auth_test.go} | 69 +++++++++++----- pkg/tsdb/loki/loki.go | 82 +++++++++++-------- .../app/plugins/datasource/loki/datasource.ts | 4 +- 5 files changed, 111 insertions(+), 76 deletions(-) rename pkg/tsdb/loki/{oauth_test.go => auth_test.go} (63%) diff --git a/pkg/tsdb/loki/api.go b/pkg/tsdb/loki/api.go index 2bbe18aba30..a6544f2f4c1 100644 --- a/pkg/tsdb/loki/api.go +++ b/pkg/tsdb/loki/api.go @@ -17,23 +17,23 @@ import ( ) type LokiAPI struct { - client *http.Client - url string - log log.Logger - oauthToken string + client *http.Client + url string + log log.Logger + headers map[string]string } -func newLokiAPI(client *http.Client, url string, log log.Logger, oauthToken string) *LokiAPI { - return &LokiAPI{client: client, url: url, log: log, oauthToken: oauthToken} +func newLokiAPI(client *http.Client, url string, log log.Logger, headers map[string]string) *LokiAPI { + return &LokiAPI{client: client, url: url, log: log, headers: headers} } -func addOauthHeader(req *http.Request, oauthToken string) { - if oauthToken != "" { - req.Header.Set("Authorization", oauthToken) +func addHeaders(req *http.Request, headers map[string]string) { + for name, value := range headers { + req.Header.Set(name, value) } } -func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery, oauthToken string) (*http.Request, error) { +func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery, headers map[string]string) (*http.Request, error) { qs := url.Values{} qs.Set("query", query.Expr) @@ -86,7 +86,7 @@ func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery, oau return nil, err } - addOauthHeader(req, oauthToken) + addHeaders(req, headers) if query.VolumeQuery { req.Header.Set("X-Query-Tags", "Source=logvolhist") @@ -139,7 +139,7 @@ func makeLokiError(body io.ReadCloser) error { } func (api *LokiAPI) DataQuery(ctx context.Context, query lokiQuery) (data.Frames, error) { - req, err := makeDataRequest(ctx, api.url, query, api.oauthToken) + req, err := makeDataRequest(ctx, api.url, query, api.headers) if err != nil { return nil, err } @@ -169,7 +169,7 @@ func (api *LokiAPI) DataQuery(ctx context.Context, query lokiQuery) (data.Frames return res.Frames, nil } -func makeRawRequest(ctx context.Context, lokiDsUrl string, resourceURL string, oauthToken string) (*http.Request, error) { +func makeRawRequest(ctx context.Context, lokiDsUrl string, resourceURL string, headers map[string]string) (*http.Request, error) { lokiUrl, err := url.Parse(lokiDsUrl) if err != nil { return nil, err @@ -186,13 +186,13 @@ func makeRawRequest(ctx context.Context, lokiDsUrl string, resourceURL string, o return nil, err } - addOauthHeader(req, oauthToken) + addHeaders(req, headers) return req, nil } func (api *LokiAPI) RawQuery(ctx context.Context, resourceURL string) ([]byte, error) { - req, err := makeRawRequest(ctx, api.url, resourceURL, api.oauthToken) + req, err := makeRawRequest(ctx, api.url, resourceURL, api.headers) if err != nil { return nil, err } diff --git a/pkg/tsdb/loki/api_mock.go b/pkg/tsdb/loki/api_mock.go index a48633677e7..ff74a1b9392 100644 --- a/pkg/tsdb/loki/api_mock.go +++ b/pkg/tsdb/loki/api_mock.go @@ -37,5 +37,5 @@ func makeMockedAPI(statusCode int, contentType string, responseBytes []byte, req Transport: &mockedRoundTripper{statusCode: statusCode, contentType: contentType, responseBytes: responseBytes, requestCallback: requestCallback}, } - return newLokiAPI(&client, "http://localhost:9999", log.New("test"), "") + return newLokiAPI(&client, "http://localhost:9999", log.New("test"), nil) } diff --git a/pkg/tsdb/loki/oauth_test.go b/pkg/tsdb/loki/auth_test.go similarity index 63% rename from pkg/tsdb/loki/oauth_test.go rename to pkg/tsdb/loki/auth_test.go index 641c29cc59c..934e3882331 100644 --- a/pkg/tsdb/loki/oauth_test.go +++ b/pkg/tsdb/loki/auth_test.go @@ -36,32 +36,32 @@ func (s *mockedCallResourceResponseSenderForOauth) Send(resp *backend.CallResour return nil } -func makeMockedDsInfoForOauth(oauthPassThru bool, body []byte, requestCallback func(req *http.Request)) datasourceInfo { +func makeMockedDsInfoForOauth(body []byte, requestCallback func(req *http.Request)) datasourceInfo { client := http.Client{ Transport: &mockedRoundTripperForOauth{requestCallback: requestCallback, body: body}, } return datasourceInfo{ - HTTPClient: &client, - OauthPassThru: oauthPassThru, + HTTPClient: &client, } } func TestOauthForwardIdentity(t *testing.T) { tt := []struct { - name string - oauthPassThru bool - headerGiven bool - headerSent bool + name string + auth bool + cookie bool }{ - {name: "when enabled and headers exist => add headers", oauthPassThru: true, headerGiven: true, headerSent: true}, - {name: "when disabled and headers exist => do not add headers", oauthPassThru: false, headerGiven: true, headerSent: false}, - {name: "when enabled and no headers exist => do not add headers", oauthPassThru: true, headerGiven: false, headerSent: false}, - {name: "when disabled and no headers exist => do not add headers", oauthPassThru: false, headerGiven: false, headerSent: false}, + {name: "when auth header exists => add auth header", auth: true, cookie: false}, + {name: "when cookie header exists => add cookie header", auth: false, cookie: true}, + {name: "when cookie&auth headers exist => add cookie&auth headers", auth: true, cookie: true}, + {name: "when no header exists => do not add headers", auth: false, cookie: false}, } authName := "Authorization" authValue := "auth" + cookieName := "Cookie" + cookieValue := "a=1" for _, test := range tt { t.Run("QueryData: "+test.name, func(t *testing.T) { @@ -83,12 +83,22 @@ func TestOauthForwardIdentity(t *testing.T) { `) clientUsed := false - dsInfo := makeMockedDsInfoForOauth(test.oauthPassThru, response, func(req *http.Request) { + dsInfo := makeMockedDsInfoForOauth(response, func(req *http.Request) { clientUsed = true - if test.headerSent { - require.Equal(t, authValue, req.Header.Get(authName)) + // we need to check for "header does not exist", + // and the only way i can find is to get the values + // as an array + authValues := req.Header.Values(authName) + cookieValues := req.Header.Values(cookieName) + if test.auth { + require.Equal(t, []string{authValue}, authValues) } else { - require.Equal(t, "", req.Header.Get(authName)) + require.Len(t, authValues, 0) + } + if test.cookie { + require.Equal(t, []string{cookieValue}, cookieValues) + } else { + require.Len(t, cookieValues, 0) } }) @@ -102,10 +112,14 @@ func TestOauthForwardIdentity(t *testing.T) { }, } - if test.headerGiven { + if test.auth { req.Headers[authName] = authValue } + if test.cookie { + req.Headers[cookieName] = cookieValue + } + tracer, err := tracing.InitializeTracerForTest() require.NoError(t, err) @@ -128,12 +142,22 @@ func TestOauthForwardIdentity(t *testing.T) { response := []byte("mocked resource response") clientUsed := false - dsInfo := makeMockedDsInfoForOauth(test.oauthPassThru, response, func(req *http.Request) { + dsInfo := makeMockedDsInfoForOauth(response, func(req *http.Request) { clientUsed = true - if test.headerSent { - require.Equal(t, authValue, req.Header.Get(authName)) + authValues := req.Header.Values(authName) + cookieValues := req.Header.Values(cookieName) + // we need to check for "header does not exist", + // and the only way i can find is to get the values + // as an array + if test.auth { + require.Equal(t, []string{authValue}, authValues) } else { - require.Equal(t, "", req.Header.Get(authName)) + require.Len(t, authValues, 0) + } + if test.cookie { + require.Equal(t, []string{cookieValue}, cookieValues) + } else { + require.Len(t, cookieValues, 0) } }) @@ -143,9 +167,12 @@ func TestOauthForwardIdentity(t *testing.T) { URL: "labels?", } - if test.headerGiven { + if test.auth { req.Headers[authName] = []string{authValue} } + if test.cookie { + req.Headers[cookieName] = []string{cookieValue} + } sender := &mockedCallResourceResponseSenderForOauth{} diff --git a/pkg/tsdb/loki/loki.go b/pkg/tsdb/loki/loki.go index 4a1a59e832d..35a09e296c7 100644 --- a/pkg/tsdb/loki/loki.go +++ b/pkg/tsdb/loki/loki.go @@ -44,9 +44,8 @@ var ( ) type datasourceInfo struct { - HTTPClient *http.Client - URL string - OauthPassThru bool + HTTPClient *http.Client + URL string // open streams streams map[string]data.FrameJSONCache @@ -65,10 +64,6 @@ type QueryJSONModel struct { VolumeQuery bool `json:"volumeQuery"` } -type DataSourceJSONModel struct { - OauthPassThru bool `json:"oauthPassThru"` -} - func parseQueryModel(raw json.RawMessage) (*QueryJSONModel, error) { model := &QueryJSONModel{} err := json.Unmarshal(raw, model) @@ -87,42 +82,29 @@ func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.Inst return nil, err } - jsonModel := DataSourceJSONModel{} - err = json.Unmarshal(settings.JSONData, &jsonModel) - if err != nil { - return nil, err - } - model := &datasourceInfo{ - HTTPClient: client, - URL: settings.URL, - OauthPassThru: jsonModel.OauthPassThru, - streams: make(map[string]data.FrameJSONCache), + HTTPClient: client, + URL: settings.URL, + streams: make(map[string]data.FrameJSONCache), } return model, nil } } -func getOauthTokenForQueryData(dsInfo *datasourceInfo, headers map[string]string) string { - if !dsInfo.OauthPassThru { +// in the CallResource API, request-headers are in a map where the value is an array-of-strings, +// so we need a helper function that can extract a single string-value from an array-of-strings. +// i only deal with two cases: +// - zero-length array +// - first-item of the array +// i do not handle the case where there are multiple items in the array, i do not know +// if that can even happen ever, for the headers that we are interested in. +func arrayHeaderFirstValue(values []string) string { + if len(values) == 0 { return "" } - return headers["Authorization"] -} - -func getOauthTokenForCallResource(dsInfo *datasourceInfo, headers map[string][]string) string { - if !dsInfo.OauthPassThru { - return "" - } - - accessValues := headers["Authorization"] - - if len(accessValues) == 0 { - return "" - } - - return accessValues[0] + // NOTE: we assume there never is a second item in the http-header-values-array + return values[0] } func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error { @@ -134,6 +116,20 @@ func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceReq return callResource(ctx, req, sender, dsInfo, s.plog) } +func getAuthHeadersForCallResource(headers map[string][]string) map[string]string { + data := make(map[string]string) + + if auth := arrayHeaderFirstValue(headers["Authorization"]); auth != "" { + data["Authorization"] = auth + } + + if cookie := arrayHeaderFirstValue(headers["Cookie"]); cookie != "" { + data["Cookie"] = cookie + } + + return data +} + func callResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender, dsInfo *datasourceInfo, plog log.Logger) error { url := req.URL @@ -148,7 +144,7 @@ func callResource(ctx context.Context, req *backend.CallResourceRequest, sender } lokiURL := fmt.Sprintf("/loki/api/v1/%s", url) - api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, getOauthTokenForCallResource(dsInfo, req.Headers)) + api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, getAuthHeadersForCallResource(req.Headers)) bytes, err := api.RawQuery(ctx, lokiURL) if err != nil { @@ -174,10 +170,24 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest) return queryData(ctx, req, dsInfo, s.plog, s.tracer) } +func getAuthHeadersForQueryData(headers map[string]string) map[string]string { + data := make(map[string]string) + + if auth := headers["Authorization"]; auth != "" { + data["Authorization"] = auth + } + + if cookie := headers["Cookie"]; cookie != "" { + data["Cookie"] = cookie + } + + return data +} + func queryData(ctx context.Context, req *backend.QueryDataRequest, dsInfo *datasourceInfo, plog log.Logger, tracer tracing.Tracer) (*backend.QueryDataResponse, error) { result := backend.NewQueryDataResponse() - api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, getOauthTokenForQueryData(dsInfo, req.Headers)) + api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, getAuthHeadersForQueryData(req.Headers)) queries, err := parseQuery(req) if err != nil { diff --git a/public/app/plugins/datasource/loki/datasource.ts b/public/app/plugins/datasource/loki/datasource.ts index f051816185e..ca8b020e0be 100644 --- a/public/app/plugins/datasource/loki/datasource.ts +++ b/public/app/plugins/datasource/loki/datasource.ts @@ -117,9 +117,7 @@ export class LokiDatasource this.languageProvider = new LanguageProvider(this); const settingsData = instanceSettings.jsonData || {}; this.maxLines = parseInt(settingsData.maxLines ?? '0', 10) || DEFAULT_MAX_LINES; - const keepCookiesUsed = (settingsData.keepCookies ?? []).length > 0; - // only use backend-mode when keep-cookies is not used - this.useBackendMode = !keepCookiesUsed && (config.featureToggles.lokiBackendMode ?? false); + this.useBackendMode = config.featureToggles.lokiBackendMode ?? false; } _request(apiUrl: string, data?: any, options?: Partial): Observable> {