[query] Support additional/custom HTTP headers (#2056)

* Added new FlagList type and used it in jaeger-query

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added tests to confirm FlagList will be treated as stringSlice by viper

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added additionalheaders logic.  Added tests

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Changed http handler to take a header struct

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Updated help text on param

Signed-off-by: Joe Elliott <number101010@gmail.com>

* FlagList => StringSlice

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Additional tests to further specify string slice behavior

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added header parsing method and tests

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added support for commas in params

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Updated flags to use expected format

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Moved stringSliceAsHeader() to flags

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added test for empty string slice

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Moved additional headers to server handler

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added handler and tests

Signed-off-by: Joe Elliott <number101010@gmail.com>

* removed unnecessary whitespace

Signed-off-by: Joe Elliott <number101010@gmail.com>

* checked out http_handler_test to master

Signed-off-by: Joe Elliott <number101010@gmail.com>

* checked out http_handler to master

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Update cmd/query/app/additional_headers_test.go

Co-Authored-By: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Joe Elliott <number101010@gmail.com>

* Update cmd/query/app/flags_test.go

Co-Authored-By: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Joe Elliott <number101010@gmail.com>

* Update cmd/query/app/flags_test.go

Co-Authored-By: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Joe Elliott <number101010@gmail.com>

* Made stringSliceAsHeader return err

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Added test for bad header params

Signed-off-by: Joe Elliott <number101010@gmail.com>

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>

cmd/all-in-one/main.go

@@ -129,7 +129,7 @@ by default uses only in-memory database.`,
 			tchanBuilder := agentTchanRep.NewBuilder().InitFromViper(v, logger)
 			grpcBuilder := agentGrpcRep.NewConnBuilder().InitFromViper(v)
 			cOpts := new(collector.CollectorOptions).InitFromViper(v)
-			qOpts := new(queryApp.QueryOptions).InitFromViper(v)
+			qOpts := new(queryApp.QueryOptions).InitFromViper(v, logger)
 
 			collectorSrv := startCollector(cOpts, spanWriter, logger, metricsFactory, strategyStore, svc.HC())
 			startAgent(aOpts, repOpts, tchanBuilder, grpcBuilder, cOpts, logger, metricsFactory)

cmd/query/app/additional_headers_handler.go

@@ -0,0 +1,30 @@
+// Copyright (c) 2020 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"net/http"
+)
+
+func additionalHeadersHandler(h http.Handler, additionalHeaders http.Header) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		header := w.Header()
+		for key, values := range additionalHeaders {
+			header[key] = values
+		}
+
+		h.ServeHTTP(w, r)
+	})
+}

cmd/query/app/additional_headers_test.go

@@ -0,0 +1,49 @@
+// Copyright (c) 2020 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAdditionalHeadersHandler(t *testing.T) {
+	additionalHeaders := http.Header{}
+	additionalHeaders.Add("Access-Control-Allow-Origin", "https://mozilla.org")
+	additionalHeaders.Add("Access-Control-Expose-Headers", "X-My-Custom-Header")
+	additionalHeaders.Add("Access-Control-Expose-Headers", "X-Another-Custom-Header")
+	additionalHeaders.Add("Access-Control-Request-Headers", "field1, field2")
+
+	emptyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte{})
+	})
+
+	handler := additionalHeadersHandler(emptyHandler, additionalHeaders)
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	req, err := http.NewRequest("GET", server.URL, nil)
+	assert.NoError(t, err)
+
+	resp, err := server.Client().Do(req)
+	assert.NoError(t, err)
+
+	for k, v := range additionalHeaders {
+		assert.Equal(t, v, resp.Header[k])
+	}
+}

cmd/query/app/flags.go

@@ -16,19 +16,28 @@
 package app
 
 import (
+	"bufio"
 	"flag"
+	"fmt"
+	"io"
+	"net/http"
+	"net/textproto"
+	"strings"
 
 	"github.com/spf13/viper"
+	"go.uber.org/zap"
 
+	"github.com/jaegertracing/jaeger/pkg/config"
 	"github.com/jaegertracing/jaeger/ports"
 )
 
 const (
-	queryPort             = "query.port"
-	queryBasePath         = "query.base-path"
-	queryStaticFiles      = "query.static-files"
-	queryUIConfig         = "query.ui-config"
-	queryTokenPropagation = "query.bearer-token-propagation"
+	queryPort              = "query.port"
+	queryBasePath          = "query.base-path"
+	queryStaticFiles       = "query.static-files"
+	queryUIConfig          = "query.ui-config"
+	queryTokenPropagation  = "query.bearer-token-propagation"
+	queryAdditionalHeaders = "query.additional-headers"
 )
 
 // QueryOptions holds configuration for query service
@@ -43,24 +52,54 @@ type QueryOptions struct {
 	UIConfig string
 	// BearerTokenPropagation activate/deactivate bearer token propagation to storage
 	BearerTokenPropagation bool
+	// AdditionalHeaders
+	AdditionalHeaders http.Header
 }
 
 // AddFlags adds flags for QueryOptions
 func AddFlags(flagSet *flag.FlagSet) {
+	flagSet.Var(&config.StringSlice{}, queryAdditionalHeaders, `Additional HTTP response headers.  Can be specified multiple times.  Format: "Key: Value"`)
 	flagSet.Int(queryPort, ports.QueryHTTP, "The port for the query service")
 	flagSet.String(queryBasePath, "/", "The base path for all HTTP routes, e.g. /jaeger; useful when running behind a reverse proxy")
 	flagSet.String(queryStaticFiles, "", "The directory path override for the static assets for the UI")
 	flagSet.String(queryUIConfig, "", "The path to the UI configuration file in JSON format")
 	flagSet.Bool(queryTokenPropagation, false, "Allow propagation of bearer token to be used by storage plugins")
-
 }
 
 // InitFromViper initializes QueryOptions with properties from viper
-func (qOpts *QueryOptions) InitFromViper(v *viper.Viper) *QueryOptions {
+func (qOpts *QueryOptions) InitFromViper(v *viper.Viper, logger *zap.Logger) *QueryOptions {
 	qOpts.Port = v.GetInt(queryPort)
 	qOpts.BasePath = v.GetString(queryBasePath)
 	qOpts.StaticAssets = v.GetString(queryStaticFiles)
 	qOpts.UIConfig = v.GetString(queryUIConfig)
 	qOpts.BearerTokenPropagation = v.GetBool(queryTokenPropagation)
+
+	stringSlice := v.GetStringSlice(queryAdditionalHeaders)
+	headers, err := stringSliceAsHeader(stringSlice)
+	if err != nil {
+		logger.Error("Failed to parse headers", zap.Strings("slice", stringSlice), zap.Error(err))
+	} else {
+		qOpts.AdditionalHeaders = headers
+	}
 	return qOpts
 }
+
+// stringSliceAsHeader parses a slice of strings and returns a http.Header.
+//  Each string in the slice is expected to be in the format "key: value"
+func stringSliceAsHeader(slice []string) (http.Header, error) {
+	if len(slice) == 0 {
+		return nil, nil
+	}
+
+	allHeaders := strings.Join(slice, "\r\n")
+
+	reader := bufio.NewReader(strings.NewReader(allHeaders))
+	tp := textproto.NewReader(reader)
+
+	header, err := tp.ReadMIMEHeader()
+	if err != nil && err != io.EOF {
+		return nil, fmt.Errorf("failed to parse headers")
+	}
+
+	return http.Header(header), nil
+}

cmd/query/app/flags_test.go

@@ -16,9 +16,11 @@
 package app
 
 import (
+	"net/http"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
 
 	"github.com/jaegertracing/jaeger/pkg/config"
 )
@@ -30,10 +32,52 @@ func TestQueryBuilderFlags(t *testing.T) {
 		"--query.ui-config=some.json",
 		"--query.base-path=/jaeger",
 		"--query.port=80",
+		"--query.additional-headers=access-control-allow-origin:blerg",
+		"--query.additional-headers=whatever:thing",
 	})
-	qOpts := new(QueryOptions).InitFromViper(v)
+	qOpts := new(QueryOptions).InitFromViper(v, zap.NewNop())
 	assert.Equal(t, "/dev/null", qOpts.StaticAssets)
 	assert.Equal(t, "some.json", qOpts.UIConfig)
 	assert.Equal(t, "/jaeger", qOpts.BasePath)
 	assert.Equal(t, 80, qOpts.Port)
+	assert.Equal(t, http.Header{
+		"Access-Control-Allow-Origin": []string{"blerg"},
+		"Whatever":                    []string{"thing"},
+	}, qOpts.AdditionalHeaders)
+}
+
+func TestQueryBuilderBadHeadersFlags(t *testing.T) {
+	v, command := config.Viperize(AddFlags)
+	command.ParseFlags([]string{
+		"--query.additional-headers=malformedheader",
+	})
+	qOpts := new(QueryOptions).InitFromViper(v, zap.NewNop())
+	assert.Nil(t, qOpts.AdditionalHeaders)
+}
+
+func TestStringSliceAsHeader(t *testing.T) {
+	headers := []string{
+		"Access-Control-Allow-Origin: https://mozilla.org",
+		"Access-Control-Expose-Headers: X-My-Custom-Header",
+		"Access-Control-Expose-Headers: X-Another-Custom-Header",
+	}
+
+	parsedHeaders, err := stringSliceAsHeader(headers)
+
+	assert.Equal(t, []string{"https://mozilla.org"}, parsedHeaders["Access-Control-Allow-Origin"])
+	assert.Equal(t, []string{"X-My-Custom-Header", "X-Another-Custom-Header"}, parsedHeaders["Access-Control-Expose-Headers"])
+	assert.NoError(t, err)
+
+	malformedHeaders := append(headers, "this is not a valid header")
+	parsedHeaders, err = stringSliceAsHeader(malformedHeaders)
+	assert.Nil(t, parsedHeaders)
+	assert.Error(t, err)
+
+	parsedHeaders, err = stringSliceAsHeader([]string{})
+	assert.Nil(t, parsedHeaders)
+	assert.NoError(t, err)
+
+	parsedHeaders, err = stringSliceAsHeader(nil)
+	assert.Nil(t, parsedHeaders)
+	assert.NoError(t, err)
 }

cmd/query/app/server.go

@@ -82,8 +82,9 @@ func createHTTPServer(querySvc *querysvc.QueryService, queryOpts *QueryOptions,
 	apiHandler.RegisterRoutes(r)
 	RegisterStaticHandler(r, logger, queryOpts)
 	var handler http.Handler = r
+	handler = additionalHeadersHandler(handler, queryOpts.AdditionalHeaders)
 	if queryOpts.BearerTokenPropagation {
-		handler = bearerTokenPropagationHandler(logger, r)
+		handler = bearerTokenPropagationHandler(logger, handler)
 	}
 	handler = handlers.CompressHandler(handler)
 	recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)

cmd/query/main.go

@@ -85,7 +85,7 @@ func main() {
 			}
 			defer closer.Close()
 			opentracing.SetGlobalTracer(tracer)
-			queryOpts := new(app.QueryOptions).InitFromViper(v)
+			queryOpts := new(app.QueryOptions).InitFromViper(v, logger)
 			// TODO: Need to figure out set enable/disable propagation on storage plugins.
 			v.Set(spanstore.StoragePropagationKey, queryOpts.BearerTokenPropagation)
 			storageFactory.InitFromViper(v)

pkg/config/string_slice.go

@@ -0,0 +1,44 @@
+// Copyright (c) 2020 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package config
+
+import "strings"
+
+// StringSlice implements the pflag.Value interface and allows for parsing multiple
+//  config values with the same name
+// It purposefully mimics pFlag.stringSliceValue (https://github.com/spf13/pflag/blob/master/string_slice.go)
+//  in order to be treated like a string slice by both viper and pflag cleanly
+type StringSlice []string
+
+// String implements pflag.Value
+func (l *StringSlice) String() string {
+	if len(*l) == 0 {
+		return "[]"
+	}
+
+	return `["` + strings.Join(*l, `","`) + `"]`
+}
+
+// Set implements pflag.Value
+func (l *StringSlice) Set(value string) error {
+	*l = append(*l, value)
+	return nil
+}
+
+// Type implements pflag.Value
+func (l *StringSlice) Type() string {
+	// this type string needs to match pflag.stringSliceValue's Type
+	return "stringSlice"
+}

pkg/config/string_slice_test.go

@@ -0,0 +1,62 @@
+// Copyright (c) 2020 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package config
+
+import (
+	"flag"
+	"testing"
+
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestStringSlice(t *testing.T) {
+	f := &StringSlice{}
+
+	assert.Equal(t, "[]", f.String())
+	assert.Equal(t, "stringSlice", f.Type())
+
+	f.Set("test")
+	assert.Equal(t, `["test"]`, f.String())
+
+	f.Set("test2")
+	assert.Equal(t, `["test","test2"]`, f.String())
+
+	f.Set("test3,test4")
+	assert.Equal(t, `["test","test2","test3,test4"]`, f.String())
+}
+
+func TestStringSliceTreatedAsStringSlice(t *testing.T) {
+	f := &StringSlice{}
+
+	// create and add flags/values to a go flag set
+	flagset := flag.NewFlagSet("test", flag.ContinueOnError)
+	flagset.Var(f, "test", "test")
+
+	err := flagset.Set("test", "asdf")
+	assert.NoError(t, err)
+	err = flagset.Set("test", "blerg")
+	assert.NoError(t, err)
+	err = flagset.Set("test", "other,thing")
+	assert.NoError(t, err)
+
+	// add go flag set to pflag
+	pflagset := pflag.FlagSet{}
+	pflagset.AddGoFlagSet(flagset)
+	actual, err := pflagset.GetStringSlice("test")
+	assert.NoError(t, err)
+
+	assert.Equal(t, []string{"asdf", "blerg", "other,thing"}, actual)
+}