Auth: Use sessionStorage instead of cookie for automatic redirection (#92759)

* WIP: working as expected, has to be tested

* Rename query param, small changes

* Remove unused code

* Address feedback

* Cleanup

* Use the feature toggle to control the behaviour

* Use the toggle on the FE too

* Prevent the extra redirect/reload

 Co-authored-by: Josh Hunt <joshhunt@users.noreply.github.com>

* Return to login if user is not authenticated

* Add tracking issue

* Align BE redirect constructor to locationSvc

pkg/api/common_test.go

@@ -191,6 +191,7 @@ func getContextHandler(t *testing.T, cfg *setting.Cfg) *contexthandler.ContextHa
 		cfg,
 		tracing.InitializeTracerForTest(),
 		&authntest.FakeService{ExpectedIdentity: &authn.Identity{ID: "0", Type: claims.TypeAnonymous, SessionToken: &usertoken.UserToken{}}},
+		featuremgmt.WithFeatures(),
 	)
 }
 

pkg/api/login.go

@@ -137,10 +137,15 @@ func (hs *HTTPServer) LoginView(c *contextmodel.ReqContext) {
 			}
 		}
 
+		if !c.UseSessionStorageRedirect {
 			c.Redirect(hs.GetRedirectURL(c))
 			return
 		}
 
+		c.Redirect(hs.Cfg.AppSubURL + "/")
+		return
+	}
+
 	c.HTML(http.StatusOK, getViewIndex(), viewData)
 }
 
@@ -211,7 +216,7 @@ func (hs *HTTPServer) LoginPost(c *contextmodel.ReqContext) response.Response {
 	}
 
 	metrics.MApiLoginPost.Inc()
-	return authn.HandleLoginResponse(c.Req, c.Resp, hs.Cfg, identity, hs.ValidateRedirectTo)
+	return authn.HandleLoginResponse(c.Req, c.Resp, hs.Cfg, identity, hs.ValidateRedirectTo, hs.Features)
 }
 
 func (hs *HTTPServer) loginUserWithUser(user *user.User, c *contextmodel.ReqContext) error {

pkg/api/login_oauth.go

@@ -57,5 +57,5 @@ func (hs *HTTPServer) OAuthLogin(reqCtx *contextmodel.ReqContext) {
 	}
 
 	metrics.MApiLoginOAuth.Inc()
-	authn.HandleLoginRedirect(reqCtx.Req, reqCtx.Resp, hs.Cfg, identity, hs.ValidateRedirectTo)
+	authn.HandleLoginRedirect(reqCtx.Req, reqCtx.Resp, hs.Cfg, identity, hs.ValidateRedirectTo, hs.Features)
 }

pkg/api/user_token.go

@@ -84,7 +84,11 @@ func (hs *HTTPServer) RotateUserAuthTokenRedirect(c *contextmodel.ReqContext) re
 		return response.Redirect(hs.Cfg.AppSubURL + "/login")
 	}
 
+	if !c.UseSessionStorageRedirect {
 		return response.Redirect(hs.GetRedirectURL(c))
+	}
+
+	return response.Redirect(hs.Cfg.AppSubURL + "/")
 }
 
 // swagger:route POST /user/auth-tokens/rotate

pkg/middleware/auth.go

@@ -44,14 +44,26 @@ func notAuthorized(c *contextmodel.ReqContext) {
 		return
 	}
 
+	if !c.UseSessionStorageRedirect {
 		writeRedirectCookie(c)
+	}
 
 	if errors.Is(c.LookupTokenErr, authn.ErrTokenNeedsRotation) {
+		if !c.UseSessionStorageRedirect {
 			c.Redirect(setting.AppSubUrl + "/user/auth-tokens/rotate")
 			return
 		}
 
+		c.Redirect(setting.AppSubUrl + "/user/auth-tokens/rotate" + getRedirectToQueryParam(c))
+		return
+	}
+
+	if !c.UseSessionStorageRedirect {
 		c.Redirect(setting.AppSubUrl + "/login")
+		return
+	}
+
+	c.Redirect(setting.AppSubUrl + "/login" + getRedirectToQueryParam(c))
 }
 
 func tokenRevoked(c *contextmodel.ReqContext, err *auth.TokenRevokedError) {
@@ -66,8 +78,13 @@ func tokenRevoked(c *contextmodel.ReqContext, err *auth.TokenRevokedError) {
 		return
 	}
 
+	if !c.UseSessionStorageRedirect {
 		writeRedirectCookie(c)
 		c.Redirect(setting.AppSubUrl + "/login")
+		return
+	}
+
+	c.Redirect(setting.AppSubUrl + "/login" + getRedirectToQueryParam(c))
 }
 
 func writeRedirectCookie(c *contextmodel.ReqContext) {
@@ -85,6 +102,21 @@ func writeRedirectCookie(c *contextmodel.ReqContext) {
 	cookies.WriteCookie(c.Resp, "redirect_to", url.QueryEscape(redirectTo), 0, nil)
 }
 
+func getRedirectToQueryParam(c *contextmodel.ReqContext) string {
+	redirectTo := c.Req.RequestURI
+	if setting.AppSubUrl != "" && strings.HasPrefix(redirectTo, setting.AppSubUrl) {
+		redirectTo = strings.TrimPrefix(redirectTo, setting.AppSubUrl)
+	}
+
+	if redirectTo == "/" {
+		return ""
+	}
+
+	// remove any forceLogin=true params
+	redirectTo = removeForceLoginParams(redirectTo)
+	return "?redirectTo=" + url.QueryEscape(redirectTo)
+}
+
 var forceLoginParamsRegexp = regexp.MustCompile(`&?forceLogin=true`)
 
 func removeForceLoginParams(str string) string {

pkg/middleware/auth_test.go

@@ -33,7 +33,7 @@ func setupAuthMiddlewareTest(t *testing.T, identity *authn.Identity, authErr err
 	return contexthandler.ProvideService(setting.NewCfg(), tracing.InitializeTracerForTest(), &authntest.FakeService{
 		ExpectedErr:      authErr,
 		ExpectedIdentity: identity,
-	})
+	}, featuremgmt.WithFeatures())
 }
 
 func TestAuth_Middleware(t *testing.T) {

pkg/middleware/middleware_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/authn/authntest"
 	"github.com/grafana/grafana/pkg/services/contexthandler"
 	contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/navtree"
 	"github.com/grafana/grafana/pkg/services/user/usertest"
 	"github.com/grafana/grafana/pkg/setting"
@@ -290,5 +291,5 @@ func getContextHandler(t *testing.T, cfg *setting.Cfg, authnService authn.Servic
 	t.Helper()
 
 	tracer := tracing.InitializeTracerForTest()
-	return contexthandler.ProvideService(cfg, tracer, authnService)
+	return contexthandler.ProvideService(cfg, tracer, authnService, featuremgmt.WithFeatures())
 }

pkg/services/accesscontrol/middleware.go

@@ -98,11 +98,16 @@ func deny(c *contextmodel.ReqContext, evaluator Evaluator, err error) {
 
 	if !c.IsApiRequest() {
 		// TODO(emil): I'd like to show a message after this redirect, not sure how that can be done?
+		if !c.UseSessionStorageRedirect {
 			writeRedirectCookie(c)
 			c.Redirect(setting.AppSubUrl + "/")
 			return
 		}
 
+		c.Redirect(setting.AppSubUrl + "/" + getRedirectToQueryParam(c))
+		return
+	}
+
 	message := ""
 	if evaluator != nil {
 		message = evaluator.String()
@@ -125,13 +130,25 @@ func unauthorized(c *contextmodel.ReqContext) {
 		return
 	}
 
+	if !c.UseSessionStorageRedirect {
 		writeRedirectCookie(c)
+	}
+
 	if errors.Is(c.LookupTokenErr, authn.ErrTokenNeedsRotation) {
+		if !c.UseSessionStorageRedirect {
 			c.Redirect(setting.AppSubUrl + "/user/auth-tokens/rotate")
 			return
 		}
+		c.Redirect(setting.AppSubUrl + "/user/auth-tokens/rotate" + getRedirectToQueryParam(c))
+		return
+	}
 
+	if !c.UseSessionStorageRedirect {
 		c.Redirect(setting.AppSubUrl + "/login")
+		return
+	}
+
+	c.Redirect(setting.AppSubUrl + "/login" + getRedirectToQueryParam(c))
 }
 
 func tokenRevoked(c *contextmodel.ReqContext, err *usertoken.TokenRevokedError) {
@@ -146,8 +163,13 @@ func tokenRevoked(c *contextmodel.ReqContext, err *usertoken.TokenRevokedError)
 		return
 	}
 
+	if !c.UseSessionStorageRedirect {
 		writeRedirectCookie(c)
 		c.Redirect(setting.AppSubUrl + "/login")
+		return
+	}
+
+	c.Redirect(setting.AppSubUrl + "/login" + getRedirectToQueryParam(c))
 }
 
 func writeRedirectCookie(c *contextmodel.ReqContext) {
@@ -162,6 +184,21 @@ func writeRedirectCookie(c *contextmodel.ReqContext) {
 	cookies.WriteCookie(c.Resp, "redirect_to", url.QueryEscape(redirectTo), 0, nil)
 }
 
+func getRedirectToQueryParam(c *contextmodel.ReqContext) string {
+	redirectTo := c.Req.RequestURI
+	if setting.AppSubUrl != "" && strings.HasPrefix(redirectTo, setting.AppSubUrl) {
+		redirectTo = strings.TrimPrefix(redirectTo, setting.AppSubUrl)
+	}
+
+	if redirectTo == "/" {
+		return ""
+	}
+
+	// remove any forceLogin=true params
+	redirectTo = removeForceLoginParams(redirectTo)
+	return "?redirectTo=" + url.QueryEscape(redirectTo)
+}
+
 var forceLoginParamsRegexp = regexp.MustCompile(`&?forceLogin=true`)
 
 func removeForceLoginParams(str string) string {

pkg/services/authn/authn.go

@@ -13,6 +13,7 @@ import (
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"github.com/grafana/grafana/pkg/middleware/cookies"
 	"github.com/grafana/grafana/pkg/models/usertoken"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/setting"
 )
@@ -230,24 +231,30 @@ func ClientWithPrefix(name string) string {
 type RedirectValidator func(url string) error
 
 // HandleLoginResponse is a utility function to perform common operations after a successful login and returns response.NormalResponse
-func HandleLoginResponse(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator) *response.NormalResponse {
+func HandleLoginResponse(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator, features featuremgmt.FeatureToggles) *response.NormalResponse {
 	result := map[string]any{"message": "Logged in"}
-	result["redirectUrl"] = handleLogin(r, w, cfg, identity, validator)
+	result["redirectUrl"] = handleLogin(r, w, cfg, identity, validator, features)
 	return response.JSON(http.StatusOK, result)
 }
 
 // HandleLoginRedirect is a utility function to perform common operations after a successful login and redirects
-func HandleLoginRedirect(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator) {
-	redirectURL := handleLogin(r, w, cfg, identity, validator)
+func HandleLoginRedirect(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator, features featuremgmt.FeatureToggles) {
+	redirectURL := handleLogin(r, w, cfg, identity, validator, features)
 	http.Redirect(w, r, redirectURL, http.StatusFound)
 }
 
 // HandleLoginRedirectResponse is a utility function to perform common operations after a successful login and return a response.RedirectResponse
-func HandleLoginRedirectResponse(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator) *response.RedirectResponse {
-	return response.Redirect(handleLogin(r, w, cfg, identity, validator))
+func HandleLoginRedirectResponse(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator, features featuremgmt.FeatureToggles) *response.RedirectResponse {
+	return response.Redirect(handleLogin(r, w, cfg, identity, validator, features))
 }
 
-func handleLogin(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator) string {
+func handleLogin(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, identity *Identity, validator RedirectValidator, features featuremgmt.FeatureToggles) string {
+	WriteSessionCookie(w, cfg, identity.SessionToken)
+
+	if features.IsEnabledGlobally(featuremgmt.FlagUseSessionStorageForRedirection) {
+		return cfg.AppSubURL + "/"
+	}
+
 	redirectURL := cfg.AppSubURL + "/"
 	if redirectTo := getRedirectURL(r); len(redirectTo) > 0 {
 		if validator(redirectTo) == nil {
@@ -256,7 +263,6 @@ func handleLogin(r *http.Request, w http.ResponseWriter, cfg *setting.Cfg, ident
 		cookies.DeleteCookie(w, "redirect_to", cookieOptions(cfg))
 	}
 
-	WriteSessionCookie(w, cfg, identity.SessionToken)
 	return redirectURL
 }
 

pkg/services/contexthandler/contexthandler.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/grafana/authlib/claims"
 	authnClients "github.com/grafana/grafana/pkg/services/authn/clients"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 
 	"github.com/grafana/grafana/pkg/api/response"
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
@@ -25,12 +26,13 @@ import (
 	"github.com/grafana/grafana/pkg/web"
 )
 
-func ProvideService(cfg *setting.Cfg, tracer tracing.Tracer, authenticator authn.Authenticator,
+func ProvideService(cfg *setting.Cfg, tracer tracing.Tracer, authenticator authn.Authenticator, features featuremgmt.FeatureToggles,
 ) *ContextHandler {
 	return &ContextHandler{
 		Cfg:           cfg,
 		tracer:        tracer,
 		authenticator: authenticator,
+		features:      features,
 	}
 }
 
@@ -39,6 +41,7 @@ type ContextHandler struct {
 	Cfg           *setting.Cfg
 	tracer        tracing.Tracer
 	authenticator authn.Authenticator
+	features      featuremgmt.FeatureToggles
 }
 
 type reqContextKey = ctxkey.Key
@@ -96,6 +99,7 @@ func (h *ContextHandler) Middleware(next http.Handler) http.Handler {
 			AllowAnonymous:            false,
 			SkipDSCache:               false,
 			Logger:                    log.New("context"),
+			UseSessionStorageRedirect: h.features.IsEnabledGlobally(featuremgmt.FlagUseSessionStorageForRedirection),
 		}
 
 		// inject ReqContext in the context

pkg/services/contexthandler/contexthandler_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/authn/authntest"
 	"github.com/grafana/grafana/pkg/services/contexthandler"
 	contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
@@ -28,6 +29,7 @@ func TestContextHandler(t *testing.T) {
 			setting.NewCfg(),
 			tracing.InitializeTracerForTest(),
 			&authntest.FakeService{ExpectedErr: errors.New("some error")},
+			featuremgmt.WithFeatures(),
 		)
 
 		server := webtest.NewServer(t, routing.NewRouteRegister())
@@ -49,6 +51,7 @@ func TestContextHandler(t *testing.T) {
 			setting.NewCfg(),
 			tracing.InitializeTracerForTest(),
 			&authntest.FakeService{ExpectedIdentity: id},
+			featuremgmt.WithFeatures(),
 		)
 
 		server := webtest.NewServer(t, routing.NewRouteRegister())
@@ -74,6 +77,7 @@ func TestContextHandler(t *testing.T) {
 			setting.NewCfg(),
 			tracing.InitializeTracerForTest(),
 			&authntest.FakeService{ExpectedIdentity: identity},
+			featuremgmt.WithFeatures(),
 		)
 
 		server := webtest.NewServer(t, routing.NewRouteRegister())
@@ -95,6 +99,7 @@ func TestContextHandler(t *testing.T) {
 			setting.NewCfg(),
 			tracing.InitializeTracerForTest(),
 			&authntest.FakeService{ExpectedIdentity: identity},
+			featuremgmt.WithFeatures(),
 		)
 
 		server := webtest.NewServer(t, routing.NewRouteRegister())
@@ -125,6 +130,7 @@ func TestContextHandler(t *testing.T) {
 			cfg,
 			tracing.InitializeTracerForTest(),
 			&authntest.FakeService{ExpectedIdentity: &authn.Identity{}},
+			featuremgmt.WithFeatures(),
 		)
 
 		server := webtest.NewServer(t, routing.NewRouteRegister())
@@ -153,6 +159,7 @@ func TestContextHandler(t *testing.T) {
 				cfg,
 				tracing.InitializeTracerForTest(),
 				&authntest.FakeService{ExpectedIdentity: &authn.Identity{ID: i, Type: typ}},
+				featuremgmt.WithFeatures(),
 			)
 
 			server := webtest.NewServer(t, routing.NewRouteRegister())

pkg/services/contexthandler/model/model.go

@@ -35,6 +35,11 @@ type ReqContext struct {
 
 	PerfmonTimer   prometheus.Summary
 	LookupTokenErr error
+
+	// FIXME: Remove this temporary flag after the rollout of FlagUseSessionStorageForRedirection feature flag
+	// Tracking issue for cleaning up this flag: https://github.com/grafana/identity-access-team/issues/908
+	// UseSessionStorageRedirect is introduced to simplify the rollout of the new redirect logic
+	UseSessionStorageRedirect bool
 }
 
 // Handle handles and logs error by given status.

public/app/app.ts

@@ -62,7 +62,7 @@ import { setMonacoEnv } from './core/monacoEnv';
 import { interceptLinkClicks } from './core/navigation/patch/interceptLinkClicks';
 import { NewFrontendAssetsChecker } from './core/services/NewFrontendAssetsChecker';
 import { backendSrv } from './core/services/backend_srv';
-import { contextSrv } from './core/services/context_srv';
+import { contextSrv, RedirectToUrlKey } from './core/services/context_srv';
 import { Echo } from './core/services/echo/Echo';
 import { reportPerformance } from './core/services/echo/EchoSrv';
 import { PerformanceBackend } from './core/services/echo/backends/PerformanceBackend';
@@ -196,6 +196,10 @@ export class GrafanaApp {
         getPanelPluginFromCache: syncGetPanelPlugin,
       });
 
+      if (config.featureToggles.useSessionStorageForRedirection) {
+        handleRedirectTo();
+      }
+
       locationUtil.initialize({
         config,
         getTimeRangeForUrl: getTimeSrv().timeRangeForUrl,
@@ -382,4 +386,30 @@ function reportMetricPerformanceMark(metricName: string, prefix = '', suffix = '
   }
 }
 
+function handleRedirectTo(): void {
+  const queryParams = locationService.getSearch();
+  const redirectToParamKey = 'redirectTo';
+
+  if (queryParams.has(redirectToParamKey) && window.location.pathname !== '/') {
+    const rawRedirectTo = queryParams.get(redirectToParamKey)!;
+    window.sessionStorage.setItem(RedirectToUrlKey, encodeURIComponent(rawRedirectTo));
+    queryParams.delete(redirectToParamKey);
+    window.history.replaceState({}, '', `${window.location.pathname}${queryParams.size > 0 ? `?${queryParams}` : ''}`);
+    return;
+  }
+
+  if (!contextSrv.user.isSignedIn) {
+    locationService.replace('/login');
+    return;
+  }
+
+  const redirectTo = window.sessionStorage.getItem(RedirectToUrlKey);
+  if (!redirectTo) {
+    return;
+  }
+
+  window.sessionStorage.removeItem(RedirectToUrlKey);
+  locationService.replace(decodeURIComponent(redirectTo));
+}
+
 export default new GrafanaApp();

public/app/core/components/Login/LoginCtrl.tsx

@@ -118,7 +118,11 @@ export class LoginCtrl extends PureComponent<Props, State> {
   };
 
   toGrafana = () => {
-    // Use window.location.href to force page reload
+    if (config.featureToggles.useSessionStorageForRedirection) {
+      window.location.assign(config.appSubUrl + '/');
+      return;
+    }
+
     if (this.result?.redirectUrl) {
       if (config.appSubUrl !== '' && !this.result.redirectUrl.startsWith(config.appSubUrl)) {
         window.location.assign(config.appSubUrl + this.result.redirectUrl);

public/app/core/services/context_srv.ts

@@ -19,6 +19,7 @@ import config from '../../core/config';
 // When set to auto, the interval will be based on the query range
 // NOTE: this is defined here rather than TimeSrv so we avoid circular dependencies
 export const AutoRefreshInterval = 'auto';
+export const RedirectToUrlKey = 'redirectTo';
 
 export class User implements Omit<CurrentUserInternal, 'lightTheme'> {
   isSignedIn: boolean;
@@ -119,6 +120,12 @@ export class ContextSrv {
    * Indicate the user has been logged out
    */
   setLoggedOut() {
+    if (config.featureToggles.useSessionStorageForRedirection) {
+      window.sessionStorage.setItem(
+        RedirectToUrlKey,
+        encodeURIComponent(window.location.href.substring(window.location.origin.length))
+      );
+    }
     this.cancelTokenRotationJob();
     this.user.isSignedIn = false;
     this.isSignedIn = false;