Move middleware context handler logic to service (#29605)

* middleware: Move context handler to own service

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

Co-authored-by: Emil Tullsted <sakjur@users.noreply.github.com>
Co-authored-by: Will Browne <wbrowne@users.noreply.github.com>

pkg/services/auth/auth_token.go

@@ -18,8 +18,14 @@ import (
 	"github.com/grafana/grafana/pkg/util"
 )
 
+const ServiceName = "UserAuthTokenService"
+
 func init() {
-	registry.RegisterService(&UserAuthTokenService{})
+	registry.Register(&registry.Descriptor{
+		Name:         ServiceName,
+		Instance:     &UserAuthTokenService{},
+		InitPriority: registry.Medium,
+	})
 }
 
 var getTime = time.Now

pkg/services/auth/testing.go

@@ -57,8 +57,13 @@ func NewFakeUserAuthTokenService() *FakeUserAuthTokenService {
 	}
 }
 
-func (s *FakeUserAuthTokenService) CreateToken(ctx context.Context, userId int64, clientIP net.IP,
-	userAgent string) (*models.UserToken, error) {
+// Init initializes the service.
+// Required for dependency injection.
+func (s *FakeUserAuthTokenService) Init() error {
+	return nil
+}
+
+func (s *FakeUserAuthTokenService) CreateToken(ctx context.Context, userId int64, clientIP net.IP, userAgent string) (*models.UserToken, error) {
 	return s.CreateTokenProvider(context.Background(), userId, clientIP, userAgent)
 }
 

pkg/services/contexthandler/auth_proxy_test.go

@@ -0,0 +1,135 @@
+package contexthandler
+
+import (
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/remotecache"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/registry"
+	"github.com/grafana/grafana/pkg/services/auth"
+	"github.com/grafana/grafana/pkg/services/contexthandler/authproxy"
+	"github.com/grafana/grafana/pkg/services/rendering"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/require"
+	macaron "gopkg.in/macaron.v1"
+)
+
+// Test initContextWithAuthProxy with a cached user ID that is no longer valid.
+//
+// In this case, the cache entry should be ignored/cleared and another attempt should be done to sign the user
+// in without cache.
+func TestInitContextWithAuthProxy_CachedInvalidUserID(t *testing.T) {
+	const name = "markelog"
+	const userID = int64(1)
+	const orgID = int64(4)
+
+	upsertHandler := func(cmd *models.UpsertUserCommand) error {
+		require.Equal(t, name, cmd.ExternalUser.Login)
+		cmd.Result = &models.User{Id: userID}
+		return nil
+	}
+	getUserHandler := func(cmd *models.GetSignedInUserQuery) error {
+		// Simulate that the cached user ID is stale
+		if cmd.UserId != userID {
+			return models.ErrUserNotFound
+		}
+
+		cmd.Result = &models.SignedInUser{
+			UserId: userID,
+			OrgId:  orgID,
+		}
+		return nil
+	}
+	bus.AddHandler("", upsertHandler)
+	bus.AddHandler("", getUserHandler)
+	t.Cleanup(func() {
+		bus.ClearBusHandlers()
+	})
+
+	svc := getContextHandler(t)
+
+	req, err := http.NewRequest("POST", "http://example.com", nil)
+	require.NoError(t, err)
+	ctx := &models.ReqContext{
+		Context: &macaron.Context{
+			Req: macaron.Request{
+				Request: req,
+			},
+			Data: map[string]interface{}{},
+		},
+		Logger: log.New("Test"),
+	}
+	req.Header.Set(svc.Cfg.AuthProxyHeaderName, name)
+	key := fmt.Sprintf(authproxy.CachePrefix, authproxy.HashCacheKey(name))
+
+	t.Logf("Injecting stale user ID in cache with key %q", key)
+	err = svc.RemoteCache.Set(key, int64(33), 0)
+	require.NoError(t, err)
+
+	authEnabled := svc.initContextWithAuthProxy(ctx, orgID)
+	require.True(t, authEnabled)
+
+	require.Equal(t, userID, ctx.SignedInUser.UserId)
+	require.True(t, ctx.IsSignedIn)
+
+	i, err := svc.RemoteCache.Get(key)
+	require.NoError(t, err)
+	require.Equal(t, userID, i.(int64))
+}
+
+type fakeRenderService struct {
+	rendering.Service
+}
+
+func (s *fakeRenderService) Init() error {
+	return nil
+}
+
+func getContextHandler(t *testing.T) *ContextHandler {
+	t.Helper()
+
+	sqlStore := sqlstore.InitTestDB(t)
+	remoteCacheSvc := &remotecache.RemoteCache{}
+
+	cfg := setting.NewCfg()
+	cfg.RemoteCacheOptions = &setting.RemoteCacheOptions{
+		Name: "database",
+	}
+	cfg.AuthProxyHeaderName = "X-Killa"
+	cfg.AuthProxyEnabled = true
+	cfg.AuthProxyHeaderProperty = "username"
+	userAuthTokenSvc := auth.NewFakeUserAuthTokenService()
+	renderSvc := &fakeRenderService{}
+	svc := &ContextHandler{}
+
+	err := registry.BuildServiceGraph([]interface{}{cfg}, []*registry.Descriptor{
+		{
+			Name:     sqlstore.ServiceName,
+			Instance: sqlStore,
+		},
+		{
+			Name:     remotecache.ServiceName,
+			Instance: remoteCacheSvc,
+		},
+		{
+			Name:     auth.ServiceName,
+			Instance: userAuthTokenSvc,
+		},
+		{
+			Name:     rendering.ServiceName,
+			Instance: renderSvc,
+		},
+		{
+			Name:     ServiceName,
+			Instance: svc,
+		},
+	})
+	require.NoError(t, err)
+
+	return svc
+}

pkg/services/contexthandler/authproxy/authproxy.go

@@ -0,0 +1,351 @@
+package authproxy
+
+import (
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"hash/fnv"
+	"net"
+	"net/mail"
+	"path"
+	"reflect"
+	"strings"
+	"time"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/remotecache"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/ldap"
+	"github.com/grafana/grafana/pkg/services/multildap"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/util"
+)
+
+const (
+
+	// CachePrefix is a prefix for the cache key
+	CachePrefix = "auth-proxy-sync-ttl:%s"
+)
+
+// getLDAPConfig gets LDAP config
+var getLDAPConfig = ldap.GetConfig
+
+// isLDAPEnabled checks if LDAP is enabled
+var isLDAPEnabled = func(cfg *setting.Cfg) bool {
+	if cfg != nil {
+		return cfg.LDAPEnabled
+	}
+
+	return setting.LDAPEnabled
+}
+
+// newLDAP creates multiple LDAP instance
+var newLDAP = multildap.New
+
+// supportedHeaders states the supported headers configuration fields
+var supportedHeaderFields = []string{"Name", "Email", "Login", "Groups"}
+
+// AuthProxy struct
+type AuthProxy struct {
+	cfg         *setting.Cfg
+	remoteCache *remotecache.RemoteCache
+	ctx         *models.ReqContext
+	orgID       int64
+	header      string
+}
+
+// Error auth proxy specific error
+type Error struct {
+	Message      string
+	DetailsError error
+}
+
+// newError returns an Error.
+func newError(message string, err error) Error {
+	return Error{
+		Message:      message,
+		DetailsError: err,
+	}
+}
+
+// Error returns the error message.
+func (err Error) Error() string {
+	return err.Message
+}
+
+// Options for the AuthProxy
+type Options struct {
+	RemoteCache *remotecache.RemoteCache
+	Ctx         *models.ReqContext
+	OrgID       int64
+}
+
+// New instance of the AuthProxy
+func New(cfg *setting.Cfg, options *Options) *AuthProxy {
+	header := options.Ctx.Req.Header.Get(cfg.AuthProxyHeaderName)
+	return &AuthProxy{
+		remoteCache: options.RemoteCache,
+		cfg:         cfg,
+		ctx:         options.Ctx,
+		orgID:       options.OrgID,
+		header:      header,
+	}
+}
+
+// IsEnabled checks if the proxy auth is enabled
+func (auth *AuthProxy) IsEnabled() bool {
+	// Bail if the setting is not enabled
+	return auth.cfg.AuthProxyEnabled
+}
+
+// HasHeader checks if the we have specified header
+func (auth *AuthProxy) HasHeader() bool {
+	return len(auth.header) != 0
+}
+
+// IsAllowedIP returns whether provided IP is allowed.
+func (auth *AuthProxy) IsAllowedIP() error {
+	ip := auth.ctx.Req.RemoteAddr
+
+	if len(strings.TrimSpace(auth.cfg.AuthProxyWhitelist)) == 0 {
+		return nil
+	}
+
+	proxies := strings.Split(auth.cfg.AuthProxyWhitelist, ",")
+	var proxyObjs []*net.IPNet
+	for _, proxy := range proxies {
+		result, err := coerceProxyAddress(proxy)
+		if err != nil {
+			return newError("could not get the network", err)
+		}
+
+		proxyObjs = append(proxyObjs, result)
+	}
+
+	sourceIP, _, err := net.SplitHostPort(ip)
+	if err != nil {
+		return newError("could not parse address", err)
+	}
+	sourceObj := net.ParseIP(sourceIP)
+
+	for _, proxyObj := range proxyObjs {
+		if proxyObj.Contains(sourceObj) {
+			return nil
+		}
+	}
+
+	return newError("proxy authentication required", fmt.Errorf(
+		"request for user (%s) from %s is not from the authentication proxy", auth.header,
+		sourceIP,
+	))
+}
+
+func HashCacheKey(key string) string {
+	hasher := fnv.New128a()
+	// according to the documentation, Hash.Write cannot error, but linter is complaining
+	hasher.Write([]byte(key)) // nolint: errcheck
+	return hex.EncodeToString(hasher.Sum(nil))
+}
+
+// getKey forms a key for the cache based on the headers received as part of the authentication flow.
+// Our configuration supports multiple headers. The main header contains the email or username.
+// And the additional ones that allow us to specify extra attributes: Name, Email or Groups.
+func (auth *AuthProxy) getKey() string {
+	key := strings.TrimSpace(auth.header) // start the key with the main header
+
+	auth.headersIterator(func(_, header string) {
+		key = strings.Join([]string{key, header}, "-") // compose the key with any additional headers
+	})
+
+	hashedKey := HashCacheKey(key)
+	return fmt.Sprintf(CachePrefix, hashedKey)
+}
+
+// Login logs in user ID by whatever means possible.
+func (auth *AuthProxy) Login(logger log.Logger, ignoreCache bool) (int64, error) {
+	if !ignoreCache {
+		// Error here means absent cache - we don't need to handle that
+		id, err := auth.GetUserViaCache(logger)
+		if err == nil && id != 0 {
+			return id, nil
+		}
+	}
+
+	if isLDAPEnabled(auth.cfg) {
+		id, err := auth.LoginViaLDAP()
+		if err != nil {
+			if errors.Is(err, ldap.ErrInvalidCredentials) {
+				return 0, newError("proxy authentication required", ldap.ErrInvalidCredentials)
+			}
+			return 0, newError("failed to get the user", err)
+		}
+
+		return id, nil
+	}
+
+	id, err := auth.LoginViaHeader()
+	if err != nil {
+		return 0, newError("failed to log in as user, specified in auth proxy header", err)
+	}
+
+	return id, nil
+}
+
+// GetUserViaCache gets user ID from cache.
+func (auth *AuthProxy) GetUserViaCache(logger log.Logger) (int64, error) {
+	cacheKey := auth.getKey()
+	logger.Debug("Getting user ID via auth cache", "cacheKey", cacheKey)
+	userID, err := auth.remoteCache.Get(cacheKey)
+	if err != nil {
+		logger.Debug("Failed getting user ID via auth cache", "error", err)
+		return 0, err
+	}
+
+	logger.Debug("Successfully got user ID via auth cache", "id", userID)
+	return userID.(int64), nil
+}
+
+// RemoveUserFromCache removes user from cache.
+func (auth *AuthProxy) RemoveUserFromCache(logger log.Logger) error {
+	cacheKey := auth.getKey()
+	logger.Debug("Removing user from auth cache", "cacheKey", cacheKey)
+	if err := auth.remoteCache.Delete(cacheKey); err != nil {
+		return err
+	}
+
+	logger.Debug("Successfully removed user from auth cache", "cacheKey", cacheKey)
+	return nil
+}
+
+// LoginViaLDAP logs in user via LDAP request
+func (auth *AuthProxy) LoginViaLDAP() (int64, error) {
+	config, err := getLDAPConfig(auth.cfg)
+	if err != nil {
+		return 0, newError("failed to get LDAP config", err)
+	}
+
+	mldap := newLDAP(config.Servers)
+	extUser, _, err := mldap.User(auth.header)
+	if err != nil {
+		return 0, err
+	}
+
+	// Have to sync grafana and LDAP user during log in
+	upsert := &models.UpsertUserCommand{
+		ReqContext:    auth.ctx,
+		SignupAllowed: auth.cfg.LDAPAllowSignup,
+		ExternalUser:  extUser,
+	}
+	if err := bus.Dispatch(upsert); err != nil {
+		return 0, err
+	}
+
+	return upsert.Result.Id, nil
+}
+
+// LoginViaHeader logs in user from the header only
+func (auth *AuthProxy) LoginViaHeader() (int64, error) {
+	extUser := &models.ExternalUserInfo{
+		AuthModule: "authproxy",
+		AuthId:     auth.header,
+	}
+
+	switch auth.cfg.AuthProxyHeaderProperty {
+	case "username":
+		extUser.Login = auth.header
+
+		emailAddr, emailErr := mail.ParseAddress(auth.header) // only set Email if it can be parsed as an email address
+		if emailErr == nil {
+			extUser.Email = emailAddr.Address
+		}
+	case "email":
+		extUser.Email = auth.header
+		extUser.Login = auth.header
+	default:
+		return 0, fmt.Errorf("auth proxy header property invalid")
+	}
+
+	auth.headersIterator(func(field string, header string) {
+		if field == "Groups" {
+			extUser.Groups = util.SplitString(header)
+		} else {
+			reflect.ValueOf(extUser).Elem().FieldByName(field).SetString(header)
+		}
+	})
+
+	upsert := &models.UpsertUserCommand{
+		ReqContext:    auth.ctx,
+		SignupAllowed: auth.cfg.AuthProxyAutoSignUp,
+		ExternalUser:  extUser,
+	}
+
+	err := bus.Dispatch(upsert)
+	if err != nil {
+		return 0, err
+	}
+
+	return upsert.Result.Id, nil
+}
+
+// headersIterator iterates over all non-empty supported additional headers
+func (auth *AuthProxy) headersIterator(fn func(field string, header string)) {
+	for _, field := range supportedHeaderFields {
+		h := auth.cfg.AuthProxyHeaders[field]
+		if h == "" {
+			continue
+		}
+
+		if value := auth.ctx.Req.Header.Get(h); value != "" {
+			fn(field, strings.TrimSpace(value))
+		}
+	}
+}
+
+// GetSignedUser gets full signed in user info.
+func (auth *AuthProxy) GetSignedInUser(userID int64) (*models.SignedInUser, error) {
+	query := &models.GetSignedInUserQuery{
+		OrgId:  auth.orgID,
+		UserId: userID,
+	}
+
+	if err := bus.Dispatch(query); err != nil {
+		return nil, err
+	}
+
+	return query.Result, nil
+}
+
+// Remember user in cache
+func (auth *AuthProxy) Remember(id int64) error {
+	key := auth.getKey()
+
+	// Check if user already in cache
+	userID, _ := auth.remoteCache.Get(key)
+	if userID != nil {
+		return nil
+	}
+
+	expiration := time.Duration(auth.cfg.AuthProxySyncTTL) * time.Minute
+
+	err := auth.remoteCache.Set(key, id, expiration)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// coerceProxyAddress gets network of the presented CIDR notation
+func coerceProxyAddress(proxyAddr string) (*net.IPNet, error) {
+	proxyAddr = strings.TrimSpace(proxyAddr)
+	if !strings.Contains(proxyAddr, "/") {
+		proxyAddr = path.Join(proxyAddr, "32")
+	}
+
+	_, network, err := net.ParseCIDR(proxyAddr)
+	if err != nil {
+		return nil, fmt.Errorf("could not parse the network: %w", err)
+	}
+	return network, nil
+}

pkg/services/contexthandler/authproxy/authproxy_test.go

@@ -0,0 +1,217 @@
+package authproxy
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/remotecache"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/ldap"
+	"github.com/grafana/grafana/pkg/services/multildap"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/macaron.v1"
+)
+
+type fakeMultiLDAP struct {
+	multildap.MultiLDAP
+	ID          int64
+	userCalled  bool
+	loginCalled bool
+}
+
+func (m *fakeMultiLDAP) Login(query *models.LoginUserQuery) (
+	*models.ExternalUserInfo, error,
+) {
+	m.loginCalled = true
+	result := &models.ExternalUserInfo{
+		UserId: m.ID,
+	}
+	return result, nil
+}
+
+func (m *fakeMultiLDAP) User(login string) (
+	*models.ExternalUserInfo,
+	ldap.ServerConfig,
+	error,
+) {
+	m.userCalled = true
+	result := &models.ExternalUserInfo{
+		UserId: m.ID,
+	}
+	return result, ldap.ServerConfig{}, nil
+}
+
+const hdrName = "markelog"
+
+func prepareMiddleware(t *testing.T, remoteCache *remotecache.RemoteCache, cb func(*http.Request, *setting.Cfg)) *AuthProxy {
+	t.Helper()
+
+	cfg := setting.NewCfg()
+	cfg.AuthProxyHeaderName = "X-Killa"
+
+	req, err := http.NewRequest("POST", "http://example.com", nil)
+	require.NoError(t, err)
+	req.Header.Set(cfg.AuthProxyHeaderName, hdrName)
+
+	if cb != nil {
+		cb(req, cfg)
+	}
+
+	ctx := &models.ReqContext{
+		Context: &macaron.Context{
+			Req: macaron.Request{
+				Request: req,
+			},
+		},
+	}
+
+	auth := New(cfg, &Options{
+		RemoteCache: remoteCache,
+		Ctx:         ctx,
+		OrgID:       4,
+	})
+
+	return auth
+}
+
+func TestMiddlewareContext(t *testing.T) {
+	logger := log.New("test")
+	cache := remotecache.NewFakeStore(t)
+
+	t.Run("When the cache only contains the main header with a simple cache key", func(t *testing.T) {
+		const id int64 = 33
+		// Set cache key
+		key := fmt.Sprintf(CachePrefix, HashCacheKey(hdrName))
+		err := cache.Set(key, id, 0)
+		require.NoError(t, err)
+		// Set up the middleware
+		auth := prepareMiddleware(t, cache, nil)
+		assert.Equal(t, key, auth.getKey())
+
+		gotID, err := auth.Login(logger, false)
+		require.NoError(t, err)
+
+		assert.Equal(t, id, gotID)
+	})
+
+	t.Run("When the cache key contains additional headers", func(t *testing.T) {
+		const id int64 = 33
+		const group = "grafana-core-team"
+
+		key := fmt.Sprintf(CachePrefix, HashCacheKey(hdrName+"-"+group))
+		err := cache.Set(key, id, 0)
+		require.NoError(t, err)
+
+		auth := prepareMiddleware(t, cache, func(req *http.Request, cfg *setting.Cfg) {
+			req.Header.Set("X-WEBAUTH-GROUPS", group)
+			cfg.AuthProxyHeaders = map[string]string{"Groups": "X-WEBAUTH-GROUPS"}
+		})
+		assert.Equal(t, "auth-proxy-sync-ttl:14f69b7023baa0ac98c96b31cec07bc0", auth.getKey())
+
+		gotID, err := auth.Login(logger, false)
+		require.NoError(t, err)
+		assert.Equal(t, id, gotID)
+	})
+}
+
+func TestMiddlewareContext_ldap(t *testing.T) {
+	logger := log.New("test")
+
+	t.Run("Logs in via LDAP", func(t *testing.T) {
+		const id int64 = 42
+
+		bus.AddHandler("test", func(cmd *models.UpsertUserCommand) error {
+			cmd.Result = &models.User{
+				Id: id,
+			}
+
+			return nil
+		})
+
+		origIsLDAPEnabled := isLDAPEnabled
+		origGetLDAPConfig := getLDAPConfig
+		origNewLDAP := newLDAP
+		t.Cleanup(func() {
+			newLDAP = origNewLDAP
+			isLDAPEnabled = origIsLDAPEnabled
+			getLDAPConfig = origGetLDAPConfig
+		})
+
+		isLDAPEnabled = func(*setting.Cfg) bool {
+			return true
+		}
+
+		stub := &fakeMultiLDAP{
+			ID: id,
+		}
+
+		getLDAPConfig = func(*setting.Cfg) (*ldap.Config, error) {
+			config := &ldap.Config{
+				Servers: []*ldap.ServerConfig{
+					{
+						SearchBaseDNs: []string{"BaseDNHere"},
+					},
+				},
+			}
+			return config, nil
+		}
+
+		newLDAP = func(servers []*ldap.ServerConfig) multildap.IMultiLDAP {
+			return stub
+		}
+
+		cache := remotecache.NewFakeStore(t)
+
+		auth := prepareMiddleware(t, cache, nil)
+
+		gotID, err := auth.Login(logger, false)
+		require.NoError(t, err)
+
+		assert.Equal(t, id, gotID)
+		assert.True(t, stub.userCalled)
+	})
+
+	t.Run("Gets nice error if LDAP is enabled, but not configured", func(t *testing.T) {
+		const id int64 = 42
+		origIsLDAPEnabled := isLDAPEnabled
+		origNewLDAP := newLDAP
+		origGetLDAPConfig := getLDAPConfig
+		t.Cleanup(func() {
+			isLDAPEnabled = origIsLDAPEnabled
+			newLDAP = origNewLDAP
+			getLDAPConfig = origGetLDAPConfig
+		})
+
+		isLDAPEnabled = func(*setting.Cfg) bool {
+			return true
+		}
+
+		getLDAPConfig = func(*setting.Cfg) (*ldap.Config, error) {
+			return nil, errors.New("something went wrong")
+		}
+
+		cache := remotecache.NewFakeStore(t)
+
+		auth := prepareMiddleware(t, cache, nil)
+
+		stub := &fakeMultiLDAP{
+			ID: id,
+		}
+
+		newLDAP = func(servers []*ldap.ServerConfig) multildap.IMultiLDAP {
+			return stub
+		}
+
+		gotID, err := auth.Login(logger, false)
+		require.EqualError(t, err, "failed to get the user")
+
+		assert.NotEqual(t, id, gotID)
+		assert.False(t, stub.loginCalled)
+	})
+}

pkg/services/contexthandler/contexthandler.go

@@ -0,0 +1,448 @@
+// Package contexthandler contains the ContextHandler service.
+package contexthandler
+
+import (
+	"context"
+	"errors"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/components/apikeygen"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/network"
+	"github.com/grafana/grafana/pkg/infra/remotecache"
+	"github.com/grafana/grafana/pkg/middleware/cookies"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/registry"
+	"github.com/grafana/grafana/pkg/services/contexthandler/authproxy"
+	"github.com/grafana/grafana/pkg/services/login"
+	"github.com/grafana/grafana/pkg/services/rendering"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/util"
+	"gopkg.in/macaron.v1"
+)
+
+const (
+	InvalidUsernamePassword = "invalid username or password"
+	InvalidAPIKey           = "invalid API key"
+)
+
+const ServiceName = "ContextHandler"
+
+func init() {
+	registry.Register(&registry.Descriptor{
+		Name:         ServiceName,
+		Instance:     &ContextHandler{},
+		InitPriority: registry.High,
+	})
+}
+
+// ContextHandler is a middleware.
+type ContextHandler struct {
+	Cfg              *setting.Cfg             `inject:""`
+	AuthTokenService models.UserTokenService  `inject:""`
+	RemoteCache      *remotecache.RemoteCache `inject:""`
+	RenderService    rendering.Service        `inject:""`
+	SQLStore         *sqlstore.SQLStore       `inject:""`
+
+	// GetTime returns the current time.
+	// Stubbable by tests.
+	GetTime func() time.Time
+}
+
+// Init initializes the service.
+func (h *ContextHandler) Init() error {
+	return nil
+}
+
+// Middleware provides a middleware to initialize the Macaron context.
+func (h *ContextHandler) Middleware(c *macaron.Context) {
+	ctx := &models.ReqContext{
+		Context:        c,
+		SignedInUser:   &models.SignedInUser{},
+		IsSignedIn:     false,
+		AllowAnonymous: false,
+		SkipCache:      false,
+		Logger:         log.New("context"),
+	}
+
+	const headerName = "X-Grafana-Org-Id"
+	orgID := int64(0)
+	orgIDHeader := ctx.Req.Header.Get(headerName)
+	if orgIDHeader != "" {
+		id, err := strconv.ParseInt(orgIDHeader, 10, 64)
+		if err == nil {
+			orgID = id
+		} else {
+			ctx.Logger.Debug("Received invalid header", "header", headerName, "value", orgIDHeader)
+		}
+	}
+
+	// the order in which these are tested are important
+	// look for api key in Authorization header first
+	// then init session and look for userId in session
+	// then look for api key in session (special case for render calls via api)
+	// then test if anonymous access is enabled
+	switch {
+	case h.initContextWithRenderAuth(ctx):
+	case h.initContextWithAPIKey(ctx):
+	case h.initContextWithBasicAuth(ctx, orgID):
+	case h.initContextWithAuthProxy(ctx, orgID):
+	case h.initContextWithToken(ctx, orgID):
+	case h.initContextWithAnonymousUser(ctx):
+	}
+
+	ctx.Logger = log.New("context", "userId", ctx.UserId, "orgId", ctx.OrgId, "uname", ctx.Login)
+	ctx.Data["ctx"] = ctx
+
+	c.Map(ctx)
+
+	// update last seen every 5min
+	if ctx.ShouldUpdateLastSeenAt() {
+		ctx.Logger.Debug("Updating last user_seen_at", "user_id", ctx.UserId)
+		if err := bus.Dispatch(&models.UpdateUserLastSeenAtCommand{UserId: ctx.UserId}); err != nil {
+			ctx.Logger.Error("Failed to update last_seen_at", "error", err)
+		}
+	}
+}
+
+func (h *ContextHandler) initContextWithAnonymousUser(ctx *models.ReqContext) bool {
+	if !h.Cfg.AnonymousEnabled {
+		return false
+	}
+
+	orgQuery := models.GetOrgByNameQuery{Name: h.Cfg.AnonymousOrgName}
+	if err := bus.Dispatch(&orgQuery); err != nil {
+		log.Errorf(3, "Anonymous access organization error: '%s': %s", h.Cfg.AnonymousOrgName, err)
+		return false
+	}
+
+	ctx.IsSignedIn = false
+	ctx.AllowAnonymous = true
+	ctx.SignedInUser = &models.SignedInUser{IsAnonymous: true}
+	ctx.OrgRole = models.RoleType(h.Cfg.AnonymousOrgRole)
+	ctx.OrgId = orgQuery.Result.Id
+	ctx.OrgName = orgQuery.Result.Name
+	return true
+}
+
+func (h *ContextHandler) initContextWithAPIKey(ctx *models.ReqContext) bool {
+	header := ctx.Req.Header.Get("Authorization")
+	parts := strings.SplitN(header, " ", 2)
+	var keyString string
+	if len(parts) == 2 && parts[0] == "Bearer" {
+		keyString = parts[1]
+	} else {
+		username, password, err := util.DecodeBasicAuthHeader(header)
+		if err == nil && username == "api_key" {
+			keyString = password
+		}
+	}
+
+	if keyString == "" {
+		return false
+	}
+
+	// base64 decode key
+	decoded, err := apikeygen.Decode(keyString)
+	if err != nil {
+		ctx.JsonApiErr(401, InvalidAPIKey, err)
+		return true
+	}
+
+	// fetch key
+	keyQuery := models.GetApiKeyByNameQuery{KeyName: decoded.Name, OrgId: decoded.OrgId}
+	if err := bus.Dispatch(&keyQuery); err != nil {
+		ctx.JsonApiErr(401, InvalidAPIKey, err)
+		return true
+	}
+
+	apikey := keyQuery.Result
+
+	// validate api key
+	isValid, err := apikeygen.IsValid(decoded, apikey.Key)
+	if err != nil {
+		ctx.JsonApiErr(500, "Validating API key failed", err)
+		return true
+	}
+	if !isValid {
+		ctx.JsonApiErr(401, InvalidAPIKey, err)
+		return true
+	}
+
+	// check for expiration
+	getTime := h.GetTime
+	if getTime == nil {
+		getTime = time.Now
+	}
+	if apikey.Expires != nil && *apikey.Expires <= getTime().Unix() {
+		ctx.JsonApiErr(401, "Expired API key", err)
+		return true
+	}
+
+	ctx.IsSignedIn = true
+	ctx.SignedInUser = &models.SignedInUser{}
+	ctx.OrgRole = apikey.Role
+	ctx.ApiKeyId = apikey.Id
+	ctx.OrgId = apikey.OrgId
+	return true
+}
+
+func (h *ContextHandler) initContextWithBasicAuth(ctx *models.ReqContext, orgID int64) bool {
+	if !h.Cfg.BasicAuthEnabled {
+		return false
+	}
+
+	header := ctx.Req.Header.Get("Authorization")
+	if header == "" {
+		return false
+	}
+
+	username, password, err := util.DecodeBasicAuthHeader(header)
+	if err != nil {
+		ctx.JsonApiErr(401, "Invalid Basic Auth Header", err)
+		return true
+	}
+
+	authQuery := models.LoginUserQuery{
+		Username: username,
+		Password: password,
+		Cfg:      h.Cfg,
+	}
+	if err := bus.Dispatch(&authQuery); err != nil {
+		ctx.Logger.Debug(
+			"Failed to authorize the user",
+			"username", username,
+			"err", err,
+		)
+
+		if errors.Is(err, models.ErrUserNotFound) {
+			err = login.ErrInvalidCredentials
+		}
+		ctx.JsonApiErr(401, InvalidUsernamePassword, err)
+		return true
+	}
+
+	user := authQuery.User
+
+	query := models.GetSignedInUserQuery{UserId: user.Id, OrgId: orgID}
+	if err := bus.Dispatch(&query); err != nil {
+		ctx.Logger.Error(
+			"Failed at user signed in",
+			"id", user.Id,
+			"org", orgID,
+		)
+		ctx.JsonApiErr(401, InvalidUsernamePassword, err)
+		return true
+	}
+
+	ctx.SignedInUser = query.Result
+	ctx.IsSignedIn = true
+	return true
+}
+
+func (h *ContextHandler) initContextWithToken(ctx *models.ReqContext, orgID int64) bool {
+	if h.Cfg.LoginCookieName == "" {
+		return false
+	}
+
+	rawToken := ctx.GetCookie(h.Cfg.LoginCookieName)
+	if rawToken == "" {
+		return false
+	}
+
+	token, err := h.AuthTokenService.LookupToken(ctx.Req.Context(), rawToken)
+	if err != nil {
+		ctx.Logger.Error("Failed to look up user based on cookie", "error", err)
+		cookies.WriteSessionCookie(ctx, h.Cfg, "", -1)
+		return false
+	}
+
+	query := models.GetSignedInUserQuery{UserId: token.UserId, OrgId: orgID}
+	if err := bus.Dispatch(&query); err != nil {
+		ctx.Logger.Error("Failed to get user with id", "userId", token.UserId, "error", err)
+		return false
+	}
+
+	ctx.SignedInUser = query.Result
+	ctx.IsSignedIn = true
+	ctx.UserToken = token
+
+	// 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(h.rotateEndOfRequestFunc(ctx, h.AuthTokenService, token))
+
+	return true
+}
+
+func (h *ContextHandler) rotateEndOfRequestFunc(ctx *models.ReqContext, authTokenService models.UserTokenService,
+	token *models.UserToken) macaron.BeforeFunc {
+	return func(w macaron.ResponseWriter) {
+		// if response has already been written, skip.
+		if w.Written() {
+			return
+		}
+
+		// if the request is cancelled by the client we should not try
+		// to rotate the token since the client would not accept any result.
+		if errors.Is(ctx.Context.Req.Context().Err(), context.Canceled) {
+			return
+		}
+
+		addr := ctx.RemoteAddr()
+		ip, err := network.GetIPFromAddress(addr)
+		if err != nil {
+			ctx.Logger.Debug("Failed to get client IP address", "addr", addr, "err", err)
+			ip = nil
+		}
+		rotated, err := authTokenService.TryRotateToken(ctx.Req.Context(), token, ip, ctx.Req.UserAgent())
+		if err != nil {
+			ctx.Logger.Error("Failed to rotate token", "error", err)
+			return
+		}
+
+		if rotated {
+			cookies.WriteSessionCookie(ctx, h.Cfg, token.UnhashedToken, h.Cfg.LoginMaxLifetime)
+		}
+	}
+}
+
+func (h *ContextHandler) initContextWithRenderAuth(ctx *models.ReqContext) bool {
+	key := ctx.GetCookie("renderKey")
+	if key == "" {
+		return false
+	}
+
+	renderUser, exists := h.RenderService.GetRenderUser(key)
+	if !exists {
+		ctx.JsonApiErr(401, "Invalid Render Key", nil)
+		return true
+	}
+
+	ctx.IsSignedIn = true
+	ctx.SignedInUser = &models.SignedInUser{
+		OrgId:   renderUser.OrgID,
+		UserId:  renderUser.UserID,
+		OrgRole: models.RoleType(renderUser.OrgRole),
+	}
+	ctx.IsRenderCall = true
+	ctx.LastSeenAt = time.Now()
+	return true
+}
+
+func logUserIn(auth *authproxy.AuthProxy, username string, logger log.Logger, ignoreCache bool) (int64, error) {
+	logger.Debug("Trying to log user in", "username", username, "ignoreCache", ignoreCache)
+	// Try to log in user via various providers
+	id, err := auth.Login(logger, ignoreCache)
+	if err != nil {
+		details := err
+		var e authproxy.Error
+		if errors.As(err, &e) {
+			details = e.DetailsError
+		}
+		logger.Error("Failed to login", "username", username, "message", err.Error(), "error", details,
+			"ignoreCache", ignoreCache)
+		return 0, err
+	}
+	return id, nil
+}
+
+func handleError(ctx *models.ReqContext, err error, statusCode int, cb func(error)) {
+	details := err
+	var e authproxy.Error
+	if errors.As(err, &e) {
+		details = e.DetailsError
+	}
+	ctx.Handle(statusCode, err.Error(), details)
+
+	if cb != nil {
+		cb(details)
+	}
+}
+
+func (h *ContextHandler) initContextWithAuthProxy(ctx *models.ReqContext, orgID int64) bool {
+	username := ctx.Req.Header.Get(h.Cfg.AuthProxyHeaderName)
+	auth := authproxy.New(h.Cfg, &authproxy.Options{
+		RemoteCache: h.RemoteCache,
+		Ctx:         ctx,
+		OrgID:       orgID,
+	})
+
+	logger := log.New("auth.proxy")
+
+	// Bail if auth proxy is not enabled
+	if !auth.IsEnabled() {
+		return false
+	}
+
+	// If there is no header - we can't move forward
+	if !auth.HasHeader() {
+		return false
+	}
+
+	// Check if allowed to continue with this IP
+	if err := auth.IsAllowedIP(); err != nil {
+		handleError(ctx, err, 407, func(details error) {
+			logger.Error("Failed to check whitelisted IP addresses", "message", err.Error(), "error", details)
+		})
+		return true
+	}
+
+	id, err := logUserIn(auth, username, logger, false)
+	if err != nil {
+		handleError(ctx, err, 407, nil)
+		return true
+	}
+
+	logger.Debug("Got user ID, getting full user info", "userID", id)
+
+	user, err := auth.GetSignedInUser(id)
+	if err != nil {
+		// The reason we couldn't find the user corresponding to the ID might be that the ID was found from a stale
+		// cache entry. For example, if a user is deleted via the API, corresponding cache entries aren't invalidated
+		// because cache keys are computed from request header values and not just the user ID. Meaning that
+		// we can't easily derive cache keys to invalidate when deleting a user. To work around this, we try to
+		// log the user in again without the cache.
+		logger.Debug("Failed to get user info given ID, retrying without cache", "userID", id)
+		if err := auth.RemoveUserFromCache(logger); err != nil {
+			if !errors.Is(err, remotecache.ErrCacheItemNotFound) {
+				logger.Error("Got unexpected error when removing user from auth cache", "error", err)
+			}
+		}
+		id, err = logUserIn(auth, username, logger, true)
+		if err != nil {
+			handleError(ctx, err, 407, nil)
+			return true
+		}
+
+		user, err = auth.GetSignedInUser(id)
+		if err != nil {
+			handleError(ctx, err, 407, nil)
+			return true
+		}
+	}
+
+	logger.Debug("Successfully got user info", "userID", user.UserId, "username", user.Login)
+
+	// Add user info to context
+	ctx.SignedInUser = user
+	ctx.IsSignedIn = true
+
+	// Remember user data in cache
+	if err := auth.Remember(id); err != nil {
+		handleError(ctx, err, 500, func(details error) {
+			logger.Error(
+				"Failed to store user in cache",
+				"username", username,
+				"message", err.Error(),
+				"error", details,
+			)
+		})
+		return true
+	}
+
+	return true
+}

pkg/services/contexthandler/contexthandler_test.go

@@ -0,0 +1,126 @@
+package contexthandler
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/components/gtime"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/auth"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/util"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	macaron "gopkg.in/macaron.v1"
+)
+
+func TestDontRotateTokensOnCancelledRequests(t *testing.T) {
+	ctxHdlr := getContextHandler(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	reqContext, _, err := initTokenRotationScenario(ctx, t)
+	require.NoError(t, err)
+
+	tryRotateCallCount := 0
+	uts := &auth.FakeUserAuthTokenService{
+		TryRotateTokenProvider: func(ctx context.Context, token *models.UserToken, clientIP net.IP,
+			userAgent string) (bool, error) {
+			tryRotateCallCount++
+			return false, nil
+		},
+	}
+
+	token := &models.UserToken{AuthToken: "oldtoken"}
+
+	fn := ctxHdlr.rotateEndOfRequestFunc(reqContext, uts, token)
+	cancel()
+	fn(reqContext.Resp)
+
+	assert.Equal(t, 0, tryRotateCallCount, "Token rotation was attempted")
+}
+
+func TestTokenRotationAtEndOfRequest(t *testing.T) {
+	ctxHdlr := getContextHandler(t)
+
+	reqContext, rr, err := initTokenRotationScenario(context.Background(), t)
+	require.NoError(t, err)
+
+	uts := &auth.FakeUserAuthTokenService{
+		TryRotateTokenProvider: func(ctx context.Context, token *models.UserToken, clientIP net.IP,
+			userAgent string) (bool, error) {
+			newToken, err := util.RandomHex(16)
+			require.NoError(t, err)
+			token.AuthToken = newToken
+			return true, nil
+		},
+	}
+
+	token := &models.UserToken{AuthToken: "oldtoken"}
+
+	ctxHdlr.rotateEndOfRequestFunc(reqContext, uts, token)(reqContext.Resp)
+
+	foundLoginCookie := false
+	resp := rr.Result()
+	defer resp.Body.Close()
+	for _, c := range resp.Cookies() {
+		if c.Name == "login_token" {
+			foundLoginCookie = true
+			require.NotEqual(t, token.AuthToken, c.Value, "Auth token is still the same")
+		}
+	}
+
+	assert.True(t, foundLoginCookie, "Could not find cookie")
+}
+
+func initTokenRotationScenario(ctx context.Context, t *testing.T) (*models.ReqContext, *httptest.ResponseRecorder, error) {
+	t.Helper()
+
+	origLoginCookieName := setting.LoginCookieName
+	origLoginMaxLifetime := setting.LoginMaxLifetime
+	t.Cleanup(func() {
+		setting.LoginCookieName = origLoginCookieName
+		setting.LoginMaxLifetime = origLoginMaxLifetime
+	})
+	setting.LoginCookieName = "login_token"
+	var err error
+	setting.LoginMaxLifetime, err = gtime.ParseDuration("7d")
+	if err != nil {
+		return nil, nil, err
+	}
+
+	rr := httptest.NewRecorder()
+	req, err := http.NewRequestWithContext(ctx, "", "", nil)
+	if err != nil {
+		return nil, nil, err
+	}
+	reqContext := &models.ReqContext{
+		Context: &macaron.Context{
+			Req: macaron.Request{
+				Request: req,
+			},
+		},
+		Logger: log.New("testlogger"),
+	}
+
+	mw := mockWriter{rr}
+	reqContext.Resp = mw
+
+	return reqContext, rr, nil
+}
+
+type mockWriter struct {
+	*httptest.ResponseRecorder
+}
+
+func (mw mockWriter) Flush()                    {}
+func (mw mockWriter) Status() int               { return 0 }
+func (mw mockWriter) Size() int                 { return 0 }
+func (mw mockWriter) Written() bool             { return false }
+func (mw mockWriter) Before(macaron.BeforeFunc) {}
+func (mw mockWriter) Push(target string, opts *http.PushOptions) error {
+	return nil
+}

pkg/services/ldap/settings.go

@@ -94,8 +94,12 @@ var config *Config
 
 // GetConfig returns the LDAP config if LDAP is enabled otherwise it returns nil. It returns either cached value of
 // the config or it reads it and caches it first.
-func GetConfig() (*Config, error) {
-	if !IsEnabled() {
+func GetConfig(cfg *setting.Cfg) (*Config, error) {
+	if cfg != nil {
+		if !cfg.LDAPEnabled {
+			return nil, nil
+		}
+	} else if !IsEnabled() {
 		return nil, nil
 	}
 

pkg/services/rendering/rendering.go

@@ -25,12 +25,13 @@ import (
 func init() {
 	remotecache.Register(&RenderUser{})
 	registry.Register(&registry.Descriptor{
-		Name:         "RenderingService",
+		Name:         ServiceName,
 		Instance:     &RenderingService{},
 		InitPriority: registry.High,
 	})
 }
 
+const ServiceName = "RenderingService"
 const renderKeyPrefix = "render-%s"
 
 type RenderUser struct {
@@ -226,8 +227,8 @@ func (rs *RenderingService) getURL(path string) string {
 		return fmt.Sprintf("%s%s&render=1", rs.Cfg.RendererCallbackUrl, path)
 	}
 
-	protocol := setting.Protocol
-	switch setting.Protocol {
+	protocol := rs.Cfg.Protocol
+	switch protocol {
 	case setting.HTTPScheme:
 		protocol = "http"
 	case setting.HTTP2Scheme, setting.HTTPSScheme:

pkg/services/rendering/rendering_test.go

@@ -28,7 +28,7 @@ func TestGetUrl(t *testing.T) {
 		t.Run("And protocol HTTP configured should return expected path", func(t *testing.T) {
 			rs.Cfg.ServeFromSubPath = false
 			rs.Cfg.AppSubURL = ""
-			setting.Protocol = setting.HTTPScheme
+			rs.Cfg.Protocol = setting.HTTPScheme
 			url := rs.getURL(path)
 			require.Equal(t, "http://localhost:3000/"+path+"&render=1", url)
 
@@ -43,7 +43,7 @@ func TestGetUrl(t *testing.T) {
 		t.Run("And protocol HTTPS configured should return expected path", func(t *testing.T) {
 			rs.Cfg.ServeFromSubPath = false
 			rs.Cfg.AppSubURL = ""
-			setting.Protocol = setting.HTTPSScheme
+			rs.Cfg.Protocol = setting.HTTPSScheme
 			url := rs.getURL(path)
 			require.Equal(t, "https://localhost:3000/"+path+"&render=1", url)
 		})
@@ -51,7 +51,7 @@ func TestGetUrl(t *testing.T) {
 		t.Run("And protocol HTTP2 configured should return expected path", func(t *testing.T) {
 			rs.Cfg.ServeFromSubPath = false
 			rs.Cfg.AppSubURL = ""
-			setting.Protocol = setting.HTTP2Scheme
+			rs.Cfg.Protocol = setting.HTTP2Scheme
 			url := rs.getURL(path)
 			require.Equal(t, "https://localhost:3000/"+path+"&render=1", url)
 		})

pkg/services/sqlstore/preferences.go

@@ -6,8 +6,6 @@ import (
 
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/models"
-
-	"github.com/grafana/grafana/pkg/setting"
 )
 
 func (ss *SQLStore) addPreferencesQueryAndCommandHandlers() {
@@ -42,7 +40,7 @@ func (ss *SQLStore) GetPreferencesWithDefaults(query *models.GetPreferencesWithD
 	}
 
 	res := &models.Preferences{
-		Theme:           setting.DefaultTheme,
+		Theme:           ss.Cfg.DefaultTheme,
 		Timezone:        ss.Cfg.DateFormats.DefaultTimezone,
 		HomeDashboardId: 0,
 	}

pkg/services/sqlstore/preferences_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 
 	"github.com/grafana/grafana/pkg/models"
-	"github.com/grafana/grafana/pkg/setting"
 	"github.com/stretchr/testify/require"
 )
 
@@ -14,7 +13,7 @@ func TestPreferencesDataAccess(t *testing.T) {
 	ss := InitTestDB(t)
 
 	t.Run("GetPreferencesWithDefaults with no saved preferences should return defaults", func(t *testing.T) {
-		setting.DefaultTheme = "light"
+		ss.Cfg.DefaultTheme = "light"
 		ss.Cfg.DateFormats.DefaultTimezone = "UTC"
 
 		query := &models.GetPreferencesWithDefaultsQuery{User: &models.SignedInUser{}}

pkg/services/sqlstore/sqlstore.go

@@ -39,16 +39,21 @@ var (
 // ContextSessionKey is used as key to save values in `context.Context`
 type ContextSessionKey struct{}
 
+const ServiceName = "SqlStore"
+const InitPriority = registry.High
+
 func init() {
+	ss := &SQLStore{}
+
 	// This change will make xorm use an empty default schema for postgres and
 	// by that mimic the functionality of how it was functioning before
 	// xorm's changes above.
 	xorm.DefaultPostgresSchema = ""
 
 	registry.Register(&registry.Descriptor{
-		Name:         "SQLStore",
-		Instance:     &SQLStore{},
-		InitPriority: registry.High,
+		Name:         ServiceName,
+		Instance:     ss,
+		InitPriority: InitPriority,
 	})
 }
 
@@ -113,13 +118,20 @@ func (ss *SQLStore) Init() error {
 
 func (ss *SQLStore) ensureMainOrgAndAdminUser() error {
 	err := ss.InTransaction(context.Background(), func(ctx context.Context) error {
-		systemUserCountQuery := models.GetSystemUserCountStatsQuery{}
-		err := bus.DispatchCtx(ctx, &systemUserCountQuery)
+		var stats models.SystemUserCountStats
+		err := ss.WithDbSession(ctx, func(sess *DBSession) error {
+			var rawSql = `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
+			if _, err := sess.SQL(rawSql).Get(&stats); err != nil {
+				return fmt.Errorf("could not determine if admin user exists: %w", err)
+			}
+
+			return nil
+		})
 		if err != nil {
-			return fmt.Errorf("could not determine if admin user exists: %w", err)
+			return err
 		}
 
-		if systemUserCountQuery.Result.Count > 0 {
+		if stats.Count > 0 {
 			return nil
 		}
 
@@ -351,7 +363,7 @@ func InitTestDB(t ITestDB) *SQLStore {
 		testSQLStore = &SQLStore{}
 		testSQLStore.Bus = bus.New()
 		testSQLStore.CacheService = localcache.New(5*time.Minute, 10*time.Minute)
-		testSQLStore.skipEnsureDefaultOrgAndUser = true
+		testSQLStore.skipEnsureDefaultOrgAndUser = false
 
 		dbType := migrator.SQLite