Auth: Rotate auth tokens at the end of requests (#21347)

By rotating the auth tokens at the end of the request we ensure
that there is minimum delay between a new token being generated
and the client receiving it.
Adds auth token slow load test which uses random latency for all 
tsdb queries..
Cleans up datasource proxy response handling.
DefaultHandler in middleware tests should write a response, the 
responseWriter BeforeFuncs wont get executed unless a response
is written.

Fixes #18644 

Co-authored-by: Marcus Efraimsson <marcus.efraimsson@gmail.com>

devenv/docker/loadtest/README.md

@@ -35,6 +35,12 @@ Run load test for 10 virtual users:
 $ ./run.sh -v 10
 ```
 
+Run auth token slow test (random query latency between 1 and 30 seconds):
+
+```bash
+$ ./run.sh -c auth_token_slow_test -s 30
+```
+
 Run auth proxy test:
 
 ```bash

devenv/docker/loadtest/auth_token_slow_test.js

@@ -0,0 +1,73 @@
+import { sleep, check, group } from 'k6';
+import { createClient, createBasicAuthClient } from './modules/client.js';
+import { createTestOrgIfNotExists, createTestdataDatasourceIfNotExists } from './modules/util.js';
+
+export let options = {
+  noCookiesReset: true
+};
+
+let endpoint = __ENV.URL || 'http://localhost:3000';
+const slowQuery = (__ENV.SLOW_QUERY && __ENV.SLOW_QUERY.length > 0) ? parseInt(__ENV.SLOW_QUERY, 10) : 5;
+const client = createClient(endpoint);
+
+export const setup = () => {
+  const basicAuthClient = createBasicAuthClient(endpoint, 'admin', 'admin');
+  const orgId = createTestOrgIfNotExists(basicAuthClient);
+  const datasourceId = createTestdataDatasourceIfNotExists(basicAuthClient);
+  client.withOrgId(orgId);
+  return {
+    orgId: orgId,
+    datasourceId: datasourceId,
+  };
+}
+
+export default (data) => {
+  group(`user auth token slow test (queries between 1 and ${slowQuery} seconds)`, () => {
+    if (__ITER === 0) {
+      group("user authenticates thru ui with username and password", () => {
+        let res = client.ui.login('admin', 'admin');
+
+        check(res, {
+          'response status is 200': (r) => r.status === 200,
+          'response has cookie \'grafana_session\' with 32 characters': (r) => r.cookies.grafana_session[0].value.length === 32,
+        });
+      });
+    }
+
+    if (__ITER !== 0) {
+      group('batch tsdb requests', () => {
+        const batchCount = 20;
+        const requests = [];
+        const payload = {
+          from: '1547765247624',
+          to: '1547768847624',
+          queries: [{
+            refId: 'A',
+            scenarioId: 'slow_query',
+            stringInput: `${Math.floor(Math.random() * slowQuery) + 1}s`,
+            intervalMs: 10000,
+            maxDataPoints: 433,
+            datasourceId: data.datasourceId,
+          }]
+        };
+
+        requests.push({ method: 'GET', url: '/api/annotations?dashboardId=2074&from=1548078832772&to=1548082432772' });
+
+        for (let n = 0; n < batchCount; n++) {
+          requests.push({ method: 'POST', url: '/api/tsdb/query', body: payload });
+        }
+
+        let responses = client.batch(requests);
+        for (let n = 0; n < batchCount; n++) {
+          check(responses[n], {
+            'response status is 200': (r) => r.status === 200,
+          });
+        }
+      });
+    }
+  });
+
+  sleep(5)
+}
+
+export const teardown = (data) => {}

devenv/docker/loadtest/modules/client.js

@@ -144,6 +144,7 @@ export const BaseClient = class BaseClient {
       let params = requests[n].params || {};
       params.headers = params.headers || {};
       params.headers['Content-Type'] = 'application/json';
+      params.timeout = 120000;
       this.beforeRequest(params);
       this.onBeforeRequest(params);
       requests[n].params = params;

devenv/docker/loadtest/run.sh

@@ -7,8 +7,9 @@ run() {
   url='http://localhost:3000'
   vus='2'
   testcase='auth_token_test'
+  slowQuery=''
 
-  while getopts ":d:u:v:c:" o; do
+  while getopts ":d:u:v:c:s:" o; do
     case "${o}" in
 				d)
             duration=${OPTARG}
@@ -22,11 +23,14 @@ run() {
         c)
             testcase=${OPTARG}
             ;;
+        s)
+            slowQuery=${OPTARG}
+            ;;
     esac
 	done
 	shift $((OPTIND-1))
 
-  docker run -t --network=host -v $PWD:/src -e URL=$url --rm -i loadimpact/k6:master run --vus $vus --duration $duration src/$testcase.js
+  docker run -t --network=host -v $PWD:/src -e URL=$url -e SLOW_QUERY=$slowQuery --rm -i loadimpact/k6:master run --vus $vus --duration $duration src/$testcase.js
 }
 
 run "$@"

pkg/api/pluginproxy/ds_proxy.go

@@ -40,6 +40,19 @@ type DataSourceProxy struct {
 	cfg       *setting.Cfg
 }
 
+type handleResponseTransport struct {
+	transport http.RoundTripper
+}
+
+func (t *handleResponseTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	res, err := t.transport.RoundTrip(req)
+	if err != nil {
+		return nil, err
+	}
+	res.Header.Del("Set-Cookie")
+	return res, nil
+}
+
 type httpClient interface {
 	Do(req *http.Request) (*http.Response, error)
 }
@@ -75,13 +88,16 @@ func (proxy *DataSourceProxy) HandleRequest() {
 		FlushInterval: time.Millisecond * 200,
 	}
 
-	var err error
-	reverseProxy.Transport, err = proxy.ds.GetHttpTransport()
+	transport, err := proxy.ds.GetHttpTransport()
 	if err != nil {
 		proxy.ctx.JsonApiErr(400, "Unable to load TLS certificate", err)
 		return
 	}
 
+	reverseProxy.Transport = &handleResponseTransport{
+		transport: transport,
+	}
+
 	proxy.logRequest()
 
 	span, ctx := opentracing.StartSpanFromContext(proxy.ctx.Req.Context(), "datasource reverse proxy")
@@ -103,14 +119,7 @@ func (proxy *DataSourceProxy) HandleRequest() {
 		logger.Error("Failed to inject span context instance", "err", err)
 	}
 
-	originalSetCookie := proxy.ctx.Resp.Header().Get("Set-Cookie")
-
 	reverseProxy.ServeHTTP(proxy.ctx.Resp, proxy.ctx.Req.Request)
-	proxy.ctx.Resp.Header().Del("Set-Cookie")
-
-	if originalSetCookie != "" {
-		proxy.ctx.Resp.Header().Set("Set-Cookie", originalSetCookie)
-	}
 }
 
 func (proxy *DataSourceProxy) addTraceFromHeaderValue(span opentracing.Span, headerName string, tagName string) {

pkg/middleware/middleware.go

@@ -226,15 +226,19 @@ func initContextWithToken(authTokenService models.UserTokenService, ctx *models.
 	ctx.IsSignedIn = true
 	ctx.UserToken = token
 
-	rotated, err := authTokenService.TryRotateToken(ctx.Req.Context(), token, ctx.RemoteAddr(), ctx.Req.UserAgent())
-	if err != nil {
-		ctx.Logger.Error("Failed to rotate token", "error", err)
-		return true
-	}
+	// Rotate the token just before we write response headers to ensure there is no delay between
+	// the new token being generated and the client receiving it.
+	ctx.Resp.Before(func(w macaron.ResponseWriter) {
+		rotated, err := authTokenService.TryRotateToken(ctx.Req.Context(), token, ctx.RemoteAddr(), ctx.Req.UserAgent())
+		if err != nil {
+			ctx.Logger.Error("Failed to rotate token", "error", err)
+			return
+		}
 
-	if rotated {
-		WriteSessionCookie(ctx, token.UnhashedToken, setting.LoginMaxLifetimeDays)
-	}
+		if rotated {
+			WriteSessionCookie(ctx, token.UnhashedToken, setting.LoginMaxLifetimeDays)
+		}
+	})
 
 	return true
 }

pkg/middleware/middleware_test.go

@@ -561,6 +561,8 @@ func middlewareScenario(t *testing.T, desc string, fn scenarioFunc) {
 			sc.context = c
 			if sc.handlerFunc != nil {
 				sc.handlerFunc(sc.context)
+			} else {
+				c.JsonOK("OK")
 			}
 		}