From 628d228ca63e7ec4cee75f26c9364d3c371a14f3 Mon Sep 17 00:00:00 2001
From: Ruben Vargas <ruben.vp8510@gmail.com>
Date: Wed, 31 Jul 2019 23:06:11 -0500
Subject: [PATCH] e2e test for elasticsearch token propagation

Signed-off-by: Ruben Vargas <ruben.vp8510@gmail.com>
---
 Makefile                                      |   5 +
 .../integration/token_propagation_test.go     | 107 ++++++++++++++++++
 scripts/travis/es-integration-test.sh         |  14 +++
 3 files changed, 126 insertions(+)
 create mode 100644 plugin/storage/integration/token_propagation_test.go

diff --git a/Makefile b/Makefile
index c5f34433a11..0a81d8a2364 100644
--- a/Makefile
+++ b/Makefile
@@ -102,6 +102,11 @@ index-cleaner-integration-test: docker-images-elastic
 	go clean -testcache
 	bash -c "set -e; set -o pipefail; $(GOTEST) -tags index_cleaner $(STORAGE_PKGS) | $(COLORIZE)"
 
+.PHONY: token-propagation-integration-test
+token-propagation-integration-test:
+	go clean -testcache
+	bash -c "set -e; set -o pipefail; $(GOTEST) -tags token_propagation -run TestBearTokenPropagation $(STORAGE_PKGS) | $(COLORIZE)"
+
 all-pkgs:
 	@echo $(ALL_PKGS) | tr ' ' '\n' | sort
 
diff --git a/plugin/storage/integration/token_propagation_test.go b/plugin/storage/integration/token_propagation_test.go
new file mode 100644
index 00000000000..f71045362ec
--- /dev/null
+++ b/plugin/storage/integration/token_propagation_test.go
@@ -0,0 +1,107 @@
+// Copyright (c) 2019 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.
+
+// +build token_propagation
+
+package integration
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/olivere/elastic.v5"
+)
+
+const bearerToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhZG1pbiIsIm5hbWUiOiJKb2huIERvZSIsImlhdCI"
+const bearerHeader = "Bearer " + bearerToken
+const queryUrl = "http://127.0.0.1:16686/api/services"
+
+type esTokenPropagationTestHandler struct {
+	test *testing.T
+}
+
+func (h *esTokenPropagationTestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	authValue := r.Header.Get("Authorization")
+	if authValue != "" {
+		headerValue := strings.Split(authValue, " ")
+		token := ""
+		if len(headerValue) == 2 {
+			if headerValue[0] == "Bearer" {
+				token = headerValue[1]
+			}
+		} else if len(headerValue) == 1 {
+			token = authValue
+		}
+		assert.Equal(h.test, bearerToken, token)
+		if token == bearerToken {
+			// Return empty results, we don't care about the result here.
+			// we just need to make sure the token was propagated to the storage and the query-service returns 200
+			ret := new(elastic.SearchResult)
+			json_ret, _ := json.Marshal(ret)
+			w.Header().Add("Content-Type", "application/json; charset=UTF-8")
+			w.Write(json_ret)
+			return
+		}
+	}
+	// No token, return error!
+	http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
+}
+
+func createElasticSearchMock(srv *http.Server, t *testing.T) {
+	srv.Handler = &esTokenPropagationTestHandler{
+		test: t,
+	}
+	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
+		assert.Error(t, err)
+	}
+}
+
+func TestBearTokenPropagation(t *testing.T) {
+	testCases := []struct {
+		name        string
+		headerValue string
+		headerName  string
+	}{
+		{name: "Bearer token", headerName: "Authorization", headerValue: bearerHeader},
+		{name: "Raw Bearer token", headerName: "Authorization", headerValue: bearerToken},
+		{name: "X-Forwarded-Access-Token", headerName: "X-Forwarded-Access-Token", headerValue: bearerHeader},
+	}
+
+	// Run elastic search mocked server in background..
+	// is not a full server, just mocked the necessary stuff for this test.
+	srv := &http.Server{Addr: ":9200"}
+	defer srv.Shutdown(context.Background())
+
+	go createElasticSearchMock(srv, t)
+
+	// Test cases.
+	for _, testCase := range testCases {
+		// Ask for services query, this should return 200
+		req, err := http.NewRequest("GET", queryUrl, nil)
+		require.NoError(t, err)
+		req.Header.Add(testCase.headerName, testCase.headerValue)
+		client := &http.Client{}
+		resp, err := client.Do(req)
+		assert.NoError(t, err)
+		assert.NotNil(t, resp)
+		if err == nil && resp != nil {
+			assert.Equal(t, resp.StatusCode, http.StatusOK)
+		}
+	}
+}
diff --git a/scripts/travis/es-integration-test.sh b/scripts/travis/es-integration-test.sh
index ca6cd54a10d..4798e181f31 100755
--- a/scripts/travis/es-integration-test.sh
+++ b/scripts/travis/es-integration-test.sh
@@ -9,3 +9,17 @@ STORAGE=elasticsearch make storage-integration-test
 make index-cleaner-integration-test
 
 docker kill $CID
+
+echo "Executing token propatagion test"
+
+# Mock UI, needed only for build query service.
+make build-crossdock-ui-placeholder
+GOOS=linux make build-query
+
+SPAN_STORAGE_TYPE=elasticsearch ./cmd/query/query-linux --es.server-urls=http://127.0.0.1:9200 --es.tls=false --query.bearer-token-propagation=true &
+
+PID=$(echo $!)
+
+make token-propagation-integration-test
+
+kill -9 ${PID}
\ No newline at end of file