AuthProxy: Invalidate previous cached item for user when changes are made to any header (#81445)

* fix: sign in using auth_proxy with role a -> b -> a would end up with role b

* Update pkg/services/authn/clients/proxy.go

Co-authored-by: Karl Persson <kalle.persson92@gmail.com>

* Update pkg/services/authn/clients/proxy.go

Co-authored-by: Karl Persson <kalle.persson92@gmail.com>

pkg/services/authn/clients/proxy.go

@@ -54,6 +54,7 @@ func ProvideProxy(cfg *setting.Cfg, cache proxyCache, userSrv user.Service, clie
 type proxyCache interface {
 	Get(ctx context.Context, key string) ([]byte, error)
 	Set(ctx context.Context, key string, value []byte, expire time.Duration) error
+	Delete(ctx context.Context, key string) error
 }
 
 type Proxy struct {
@@ -146,13 +147,32 @@ func (c *Proxy) Hook(ctx context.Context, identity *authn.Identity, r *authn.Req
 		c.log.Warn("Failed to cache proxy user", "error", err, "userId", identifier, "err", err)
 		return nil
 	}
+
+	// User's role would not be updated if the cache hit. If requests arrive in the following order:
+	// 1. Name = x; Role = Admin			# cache missed, new user created and cached with key Name=x;Role=Admin
+	// 2. Name = x; Role = Editor			# cache missed, the user got updated and cached with key Name=x;Role=Editor
+	// 3. Name = x; Role = Admin			# cache hit with key Name=x;Role=Admin, no update, the user stays with Role=Editor
+	// To avoid such a problem we also cache the key used using `prefix:[username]`.
+	// Then whenever we get a cache miss due to changes in any header we use it to invalidate the previous item.
+	username := getProxyHeader(r, c.cfg.AuthProxyHeaderName, c.cfg.AuthProxyHeadersEncoded)
+	userKey := fmt.Sprintf("%s:%s", proxyCachePrefix, username)
+
+	// invalidate previously cached user id
+	if prevCacheKey, err := c.cache.Get(ctx, userKey); err == nil && len(prevCacheKey) > 0 {
+		if err := c.cache.Delete(ctx, string(prevCacheKey)); err != nil {
+			return err
+		}
+	}
+
 	c.log.FromContext(ctx).Debug("Cache proxy user", "userId", id)
 	bytes := []byte(strconv.FormatInt(id, 10))
-	if err := c.cache.Set(ctx, identity.ClientParams.CacheAuthProxyKey, bytes, time.Duration(c.cfg.AuthProxySyncTTL)*time.Minute); err != nil {
+	duration := time.Duration(c.cfg.AuthProxySyncTTL) * time.Minute
+	if err := c.cache.Set(ctx, identity.ClientParams.CacheAuthProxyKey, bytes, duration); err != nil {
 		c.log.Warn("Failed to cache proxy user", "error", err, "userId", id)
 	}
 
-	return nil
+	// store current cacheKey for the user
+	return c.cache.Set(ctx, userKey, []byte(identity.ClientParams.CacheAuthProxyKey), duration)
 }
 
 func (c *Proxy) isAllowedIP(r *authn.Request) bool {

pkg/services/authn/clients/proxy_test.go

@@ -3,6 +3,7 @@ package clients
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net/http"
 	"testing"
 	"time"
@@ -112,7 +113,7 @@ func TestProxy_Authenticate(t *testing.T) {
 				calledAdditional = additional
 				return nil, nil
 			}}
-			c, err := ProvideProxy(cfg, fakeCache{expectedErr: errors.New("")}, usertest.NewUserServiceFake(), proxyClient)
+			c, err := ProvideProxy(cfg, &fakeCache{expectedErr: errors.New("")}, usertest.NewUserServiceFake(), proxyClient)
 			require.NoError(t, err)
 
 			_, err = c.Authenticate(context.Background(), tt.req)
@@ -177,14 +178,65 @@ func TestProxy_Test(t *testing.T) {
 var _ proxyCache = new(fakeCache)
 
 type fakeCache struct {
-	expectedErr  error
-	expectedItem []byte
+	data        map[string][]byte
+	expectedErr error
 }
 
-func (f fakeCache) Get(ctx context.Context, key string) ([]byte, error) {
-	return f.expectedItem, f.expectedErr
+func (f *fakeCache) Get(ctx context.Context, key string) ([]byte, error) {
+	return f.data[key], f.expectedErr
 }
 
-func (f fakeCache) Set(ctx context.Context, key string, value []byte, expire time.Duration) error {
+func (f *fakeCache) Set(ctx context.Context, key string, value []byte, expire time.Duration) error {
+	f.data[key] = value
 	return f.expectedErr
 }
+
+func (f fakeCache) Delete(ctx context.Context, key string) error {
+	delete(f.data, key)
+	return f.expectedErr
+}
+
+func TestProxy_Hook(t *testing.T) {
+	cfg := setting.NewCfg()
+	cfg.AuthProxyHeaderName = "X-Username"
+	cfg.AuthProxyHeaders = map[string]string{
+		proxyFieldRole: "X-Role",
+	}
+	cache := &fakeCache{data: make(map[string][]byte)}
+	userId := 1
+	userID := fmt.Sprintf("%s:%d", authn.NamespaceUser, userId)
+
+	// withRole creates a test case for a user with a specific role.
+	withRole := func(role string) func(t *testing.T) {
+		cacheKey := fmt.Sprintf("users:johndoe-%s", role)
+		return func(t *testing.T) {
+			c, err := ProvideProxy(cfg, cache, usertest.NewUserServiceFake(), authntest.MockProxyClient{})
+			require.NoError(t, err)
+			userIdentity := &authn.Identity{
+				ID: userID,
+				ClientParams: authn.ClientParams{
+					CacheAuthProxyKey: cacheKey,
+				},
+			}
+			userReq := &authn.Request{
+				HTTPRequest: &http.Request{
+					Header: map[string][]string{
+						"X-Username": {"johndoe"},
+						"X-Role":     {role},
+					},
+				},
+			}
+			err = c.Hook(context.Background(), userIdentity, userReq)
+			assert.NoError(t, err)
+			expectedCache := map[string][]byte{
+				cacheKey: []byte(fmt.Sprintf("%d", userId)),
+				fmt.Sprintf("%s:%s", proxyCachePrefix, "johndoe"): []byte(fmt.Sprintf("users:johndoe-%s", role)),
+			}
+			assert.Equal(t, expectedCache, cache.data)
+		}
+	}
+
+	t.Run("step 1: new user with role Admin", withRole("Admin"))
+	t.Run("step 2: cached user with new Role Viewer", withRole("Viewer"))
+	t.Run("step 3: cached user get changed back to Admin", withRole("Admin"))
+}