dsproxy: allow multiple access tokens per datasource

Changes the cache key for tokens to cache on datasource id +
route path + http method instead of just datasource id.

pkg/api/pluginproxy/ds_proxy.go

@@ -25,12 +25,8 @@ import (
 )
 
 var (
-	logger = log.New("data-proxy-log")
-	client = &http.Client{
-		Timeout:   time.Second * 30,
-		Transport: &http.Transport{Proxy: http.ProxyFromEnvironment},
-	}
-	tokenCache = map[int64]*jwtToken{}
+	logger     = log.New("data-proxy-log")
+	tokenCache = map[string]*jwtToken{}
 )
 
 type jwtToken struct {
@@ -40,12 +36,17 @@ type jwtToken struct {
 }
 
 type DataSourceProxy struct {
-	ds        *m.DataSource
-	ctx       *m.ReqContext
-	targetUrl *url.URL
-	proxyPath string
-	route     *plugins.AppPluginRoute
-	plugin    *plugins.DataSourcePlugin
+	ds         *m.DataSource
+	ctx        *m.ReqContext
+	targetUrl  *url.URL
+	proxyPath  string
+	route      *plugins.AppPluginRoute
+	plugin     *plugins.DataSourcePlugin
+	httpClient HttpClient
+}
+
+type HttpClient interface {
+	Do(req *http.Request) (*http.Response, error)
 }
 
 func NewDataSourceProxy(ds *m.DataSource, plugin *plugins.DataSourcePlugin, ctx *m.ReqContext, proxyPath string) *DataSourceProxy {
@@ -57,6 +58,10 @@ func NewDataSourceProxy(ds *m.DataSource, plugin *plugins.DataSourcePlugin, ctx
 		ctx:       ctx,
 		proxyPath: proxyPath,
 		targetUrl: targetURL,
+		httpClient: &http.Client{
+			Timeout:   time.Second * 30,
+			Transport: &http.Transport{Proxy: http.ProxyFromEnvironment},
+		},
 	}
 }
 
@@ -311,7 +316,7 @@ func (proxy *DataSourceProxy) applyRoute(req *http.Request) {
 }
 
 func (proxy *DataSourceProxy) getAccessToken(data templateData) (string, error) {
-	if cachedToken, found := tokenCache[proxy.ds.Id]; found {
+	if cachedToken, found := tokenCache[proxy.getAccessTokenCacheKey()]; found {
 		if cachedToken.ExpiresOn.After(time.Now().Add(time.Second * 10)) {
 			logger.Info("Using token from cache")
 			return cachedToken.AccessToken, nil
@@ -336,7 +341,7 @@ func (proxy *DataSourceProxy) getAccessToken(data templateData) (string, error)
 	getTokenReq.Header.Add("Content-Type", "application/x-www-form-urlencoded")
 	getTokenReq.Header.Add("Content-Length", strconv.Itoa(len(params.Encode())))
 
-	resp, err := client.Do(getTokenReq)
+	resp, err := proxy.httpClient.Do(getTokenReq)
 	if err != nil {
 		return "", err
 	}
@@ -350,12 +355,16 @@ func (proxy *DataSourceProxy) getAccessToken(data templateData) (string, error)
 
 	expiresOnEpoch, _ := strconv.ParseInt(token.ExpiresOnString, 10, 64)
 	token.ExpiresOn = time.Unix(expiresOnEpoch, 0)
-	tokenCache[proxy.ds.Id] = &token
+	tokenCache[proxy.getAccessTokenCacheKey()] = &token
 
 	logger.Info("Got new access token", "ExpiresOn", token.ExpiresOn)
 	return token.AccessToken, nil
 }
 
+func (proxy *DataSourceProxy) getAccessTokenCacheKey() string {
+	return fmt.Sprintf("%v_%v_%v", proxy.ds.Id, proxy.route.Path, proxy.route.Method)
+}
+
 func interpolateString(text string, data templateData) (string, error) {
 	t, err := template.New("content").Parse(text)
 	if err != nil {

pkg/api/pluginproxy/ds_proxy_test.go

@@ -1,9 +1,13 @@
 package pluginproxy
 
 import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
 	"net/http"
 	"net/url"
 	"testing"
+	"time"
 
 	macaron "gopkg.in/macaron.v1"
 
@@ -100,6 +104,109 @@ func TestDSRouteRule(t *testing.T) {
 			})
 		})
 
+		Convey("Plugin with multiple routes for token auth", func() {
+			plugin := &plugins.DataSourcePlugin{
+				Routes: []*plugins.AppPluginRoute{
+					{
+						Path: "pathwithtoken1",
+						Url:  "https://api.nr1.io/some/path",
+						TokenAuth: &plugins.JwtTokenAuth{
+							Url: "https://login.server.com/{{.JsonData.tenantId}}/oauth2/token",
+							Params: map[string]string{
+								"grant_type":    "client_credentials",
+								"client_id":     "{{.JsonData.clientId}}",
+								"client_secret": "{{.SecureJsonData.clientSecret}}",
+								"resource":      "https://api.nr1.io",
+							},
+						},
+					},
+					{
+						Path: "pathwithtoken2",
+						Url:  "https://api.nr2.io/some/path",
+						TokenAuth: &plugins.JwtTokenAuth{
+							Url: "https://login.server.com/{{.JsonData.tenantId}}/oauth2/token",
+							Params: map[string]string{
+								"grant_type":    "client_credentials",
+								"client_id":     "{{.JsonData.clientId}}",
+								"client_secret": "{{.SecureJsonData.clientSecret}}",
+								"resource":      "https://api.nr2.io",
+							},
+						},
+					},
+				},
+			}
+
+			setting.SecretKey = "password"
+			key, _ := util.Encrypt([]byte("123"), "password")
+
+			ds := &m.DataSource{
+				JsonData: simplejson.NewFromAny(map[string]interface{}{
+					"clientId": "asd",
+					"tenantId": "mytenantId",
+				}),
+				SecureJsonData: map[string][]byte{
+					"clientSecret": key,
+				},
+			}
+
+			req, _ := http.NewRequest("GET", "http://localhost/asd", nil)
+			ctx := &m.ReqContext{
+				Context: &macaron.Context{
+					Req: macaron.Request{Request: req},
+				},
+				SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR},
+			}
+
+			Convey("When creating and caching access tokens", func() {
+				var authorizationHeaderCall1 string
+				var authorizationHeaderCall2 string
+
+				Convey("first call should add authorization header with access token", func() {
+					json, err := ioutil.ReadFile("./test-data/access-token-1.json")
+					So(err, ShouldBeNil)
+
+					proxy1 := NewDataSourceProxyWithMock(ds, plugin, ctx, "pathwithtoken1", json)
+					proxy1.route = plugin.Routes[0]
+					proxy1.applyRoute(req)
+
+					authorizationHeaderCall1 = req.Header.Get("Authorization")
+					So(req.URL.String(), ShouldEqual, "https://api.nr1.io/some/path")
+					So(authorizationHeaderCall1, ShouldStartWith, "Bearer eyJ0e")
+
+					Convey("second call to another route should add a different access token", func() {
+						json2, err := ioutil.ReadFile("./test-data/access-token-2.json")
+						So(err, ShouldBeNil)
+
+						req, _ := http.NewRequest("GET", "http://localhost/asd", nil)
+						proxy2 := NewDataSourceProxyWithMock(ds, plugin, ctx, "pathwithtoken2", json2)
+						proxy2.route = plugin.Routes[1]
+						proxy2.applyRoute(req)
+
+						authorizationHeaderCall2 = req.Header.Get("Authorization")
+
+						So(req.URL.String(), ShouldEqual, "https://api.nr2.io/some/path")
+						So(authorizationHeaderCall1, ShouldStartWith, "Bearer eyJ0e")
+						So(authorizationHeaderCall2, ShouldStartWith, "Bearer eyJ0e")
+						So(authorizationHeaderCall2, ShouldNotEqual, authorizationHeaderCall1)
+
+						Convey("third call to first route should add cached access token", func() {
+							req, _ := http.NewRequest("GET", "http://localhost/asd", nil)
+
+							proxy3 := NewDataSourceProxyWithMock(ds, plugin, ctx, "pathwithtoken1", []byte{})
+							proxy3.route = plugin.Routes[0]
+							proxy3.applyRoute(req)
+
+							authorizationHeaderCall3 := req.Header.Get("Authorization")
+							So(req.URL.String(), ShouldEqual, "https://api.nr1.io/some/path")
+							So(authorizationHeaderCall1, ShouldStartWith, "Bearer eyJ0e")
+							So(authorizationHeaderCall3, ShouldStartWith, "Bearer eyJ0e")
+							So(authorizationHeaderCall3, ShouldEqual, authorizationHeaderCall1)
+						})
+					})
+				})
+			})
+		})
+
 		Convey("When proxying graphite", func() {
 			plugin := &plugins.DataSourcePlugin{}
 			ds := &m.DataSource{Url: "htttp://graphite:8080", Type: m.DS_GRAPHITE}
@@ -214,3 +321,36 @@ func TestDSRouteRule(t *testing.T) {
 
 	})
 }
+
+type HttpClientStub struct {
+	fakeBody []byte
+}
+
+func (c *HttpClientStub) Do(req *http.Request) (*http.Response, error) {
+	bodyJSON, _ := simplejson.NewJson(c.fakeBody)
+	_, passedTokenCacheTest := bodyJSON.CheckGet("expires_on")
+	So(passedTokenCacheTest, ShouldBeTrue)
+
+	bodyJSON.Set("expires_on", fmt.Sprint(time.Now().Add(time.Second*60).Unix()))
+	body, _ := bodyJSON.MarshalJSON()
+	resp := &http.Response{
+		Body: ioutil.NopCloser(bytes.NewReader(body)),
+	}
+
+	return resp, nil
+}
+
+func NewDataSourceProxyWithMock(ds *m.DataSource, plugin *plugins.DataSourcePlugin, ctx *m.ReqContext, proxyPath string, fakeBody []byte) *DataSourceProxy {
+	targetURL, _ := url.Parse(ds.Url)
+
+	return &DataSourceProxy{
+		ds:        ds,
+		plugin:    plugin,
+		ctx:       ctx,
+		proxyPath: proxyPath,
+		targetUrl: targetURL,
+		httpClient: &HttpClientStub{
+			fakeBody: fakeBody,
+		},
+	}
+}

pkg/api/pluginproxy/test-data/access-token-1.json

@@ -0,0 +1,9 @@
+{
+  "token_type": "Bearer",
+  "expires_in": "3599",
+  "ext_expires_in": "0",
+  "expires_on": "1528740417",
+  "not_before": "1528736517",
+  "resource": "https://api.nr1.io",
+  "access_token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6ImlCakwxUmNxemhpeTRmcHhJeGRacW9oTTJZayIsImtpZCI6ImlCakwxUmNxemhpeTRmcHhJeGRacW9oTTJZayJ9.eyJhdWQiOiJodHRwczovL2FwaS5sb2dhbmFseXRpY3MuaW8iLCJpc3MiOiJodHRwczovL3N0cy53aW5kb3dzLm5ldC9lN2YzZjY2MS1hOTMzLTRiM2YtODE3Ni01MWM0Zjk4MmVjNDgvIiwiaWF0IjoxNTI4NzM2NTE3LCJuYmYiOjE1Mjg3MzY1MTcsImV4cCI6MTUyODc0MDQxNywiYWlvIjoiWTJkZ1lBaStzaWRsT3NmQ2JicGhLMSsremttN0NBQT0iLCJhcHBpZCI6IjdmMzJkYjdjLTZmNmYtNGU4OC05M2Q5LTlhZTEyNmMwYTU1ZiIsImFwcGlkYWNyIjoiMSIsImlkcCI6Imh0dHBzOi8vc3RzLndpbmRvd3MubmV0L2U3ZjNmNjYxLWE5MzMtNGIzZi04MTc2LTUxYzRmOTgyZWM0OC8iLCJvaWQiOiI1NDQ5ZmJjOS1mYWJhLTRkNjItODE2Yy05ZmMwMzZkMWViN2UiLCJzdWIiOiI1NDQ5ZmJjOS1mYWJhLTRkNjItODE2Yy05ZmMwMzZkMWViN2UiLCJ0aWQiOiJlN2YzZjY2MS1hOTMzLTRiM2YtODE3Ni01MWM0Zjk4MmVjNDgiLCJ1dGkiOiJZQTlQa2lxUy1VV1hMQjhIRnU0U0FBIiwidmVyIjoiMS4wIn0.ga5qudt4LDMKTStAxUmzjyZH8UFBAaFirJqpTdmYny4NtkH6JT2EILvjTjYxlKeTQisvwx9gof0PyicZIab9d6wlMa2xiLzr2nmaOonYClY8fqBaRTgc1xVjrKFw5SCgpx3FnEyJhIWvVPIfaWaogSHcQbIpe4kdk4tz-ccmrx0D1jsziSI4BZcJcX04aJuHZGz9k4mQZ_AA5sQSeQaNuojIng6rYoIifAXFYBZPTbeeeqmiGq8v0IOLeNKbC0POeQCJC_KKBG6Z_MV2KgPxFEzQuX2ZFmRD_wGPteV5TUBxh1kARdqexA3e0zAKSawR9kmrAiZ21lPr4tX2Br_HDg"
+}

pkg/api/pluginproxy/test-data/access-token-2.json

@@ -0,0 +1,9 @@
+{
+  "token_type": "Bearer",
+  "expires_in": "3599",
+  "ext_expires_in": "0",
+  "expires_on": "1528662059",
+  "not_before": "1528658159",
+  "resource": "https://api.nr2.io",
+  "access_token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6ImlCakwxUmNxemhpeTRmcHhJeGRacW9oTTJZayIsImtpZCI6ImlCakwxUmNxemhpeTRmcHhJeGRacW9oTTJZayJ9.eyJhdWQiOiJodHRwczovL21hbmFnZW1lbnQuYXp1cmUuY29tLyIsImlzcyI6Imh0dHBzOi8vc3RzLndpbmRvd3MubmV0L2U3ZjNmNjYxLWE5MzMtNGIzZi04MTc2LTUxYzRmOTgyZWM0OC8iLCJpYXQiOjE1Mjg2NTgxNTksIm5iZiI6MTUyODY1ODE1OSwiZXhwIjoxNTI4NjYyMDU5LCJhaW8iOiJZMmRnWUFpK3NpZGxPc2ZDYmJwaEsxKyt6a203Q0FBPSIsImFwcGlkIjoiODg5YjdlZDgtMWFlZC00ODZlLTk3ODktODE5NzcwYmJiNjFhIiwiYXBwaWRhY3IiOiIxIiwiaWRwIjoiaHR0cHM6Ly9zdHMud2luZG93cy5uZXQvZTdmM2Y2NjEtYTkzMy00YjNmLTgxNzYtNTFjNGY5ODJlYzQ4LyIsIm9pZCI6IjY0YzQxNjMyLTliOWUtNDczNy05MTYwLTBlNjAzZTg3NjljYyIsInN1YiI6IjY0YzQxNjMyLTliOWUtNDczNy05MTYwLTBlNjAzZTg3NjljYyIsInRpZCI6ImU3ZjNmNjYxLWE5MzMtNGIzZi04MTc2LTUxYzRmOTgyZWM0OCIsInV0aSI6IkQ1ODZHSGUySDBPd0ptOU0xeVlKQUEiLCJ2ZXIiOiIxLjAifQ.Pw8c8gpoZptw3lGreQoHQaMVOozSaTE5D38Vm2aCHRB3DvD3N-Qcm1x0ZCakUEV2sJd7jvx4XtPFuW7063T0V1deExL4rzzvIo0ZfMmURf9tCTiKFKYibqf8_PtfPSz0t9eNDEUGmWDh1Wgssb4W_H-wPqgl9VPMT7T6ynkfIm0-ODPZTBzgSHiY8C_L1-DkhsK7XiqbUlSDgx9FpfChZS3ah8QhA8geqnb_HVuSktg7WhpxmogSpK5QdrwSE3jsbItpzOfLJ4iBd2ExzS2C0y8H_Coluk3Y1YA07tAxJ6Y7oBv-XwGqNfZhveOCQOzX-U3dFod3fXXysjB0UB89WQ"
+}