feat: improve sessionmanager logic (#301)

* feat: improve sessionmanager logic

* feat: improve sessionmanager logic

* feat: fix linter issues

* feat: fix linter issues

* feat: fix linter issues

changelog/CHANGELOG.md

@@ -0,0 +1,2 @@
+### Features
+- feat: improve SessionManager

server/application/terminal.go

@@ -38,12 +38,12 @@ type terminalHandler struct {
 	allowedShells     []string
 	namespace         string
 	enabledNamespaces []string
-	sessionManager    util_session.SessionManager
+	sessionManager    *util_session.SessionManager
 }
 
 // NewHandler returns a new terminal handler.
 func NewHandler(appLister applisters.ApplicationLister, namespace string, enabledNamespaces []string, db db.ArgoDB, enf *rbac.Enforcer, cache *servercache.Cache,
-	appResourceTree AppResourceTreeFn, allowedShells []string, sessionManager util_session.SessionManager) *terminalHandler {
+	appResourceTree AppResourceTreeFn, allowedShells []string, sessionManager *util_session.SessionManager) *terminalHandler {
 	return &terminalHandler{
 		appLister:         appLister,
 		db:                db,

server/application/websocket.go

@@ -37,7 +37,7 @@ type terminalSession struct {
 	tty            bool
 	readLock       sync.Mutex
 	writeLock      sync.Mutex
-	sessionManager util_session.SessionManager
+	sessionManager *util_session.SessionManager
 	token          *string
 }
 
@@ -48,7 +48,7 @@ func getToken(r *http.Request) (string, error) {
 }
 
 // newTerminalSession create terminalSession
-func newTerminalSession(w http.ResponseWriter, r *http.Request, responseHeader http.Header, sessionManager util_session.SessionManager) (*terminalSession, error) {
+func newTerminalSession(w http.ResponseWriter, r *http.Request, responseHeader http.Header, sessionManager *util_session.SessionManager) (*terminalSession, error) {
 	token, err := getToken(r)
 	if err != nil {
 		return nil, err

server/server.go

@@ -992,7 +992,7 @@ func (a *ArgoCDServer) newHTTPServer(ctx context.Context, port int, grpcWebHandl
 	}
 	mux.Handle("/api/", handler)
 
-	terminal := application.NewHandler(a.appLister, a.Namespace, a.ApplicationNamespaces, a.db, a.enf, a.Cache, appResourceTreeFn, a.settings.ExecShells, *a.sessionMgr).
+	terminal := application.NewHandler(a.appLister, a.Namespace, a.ApplicationNamespaces, a.db, a.enf, a.Cache, appResourceTreeFn, a.settings.ExecShells, a.sessionMgr).
 		WithFeatureFlagMiddleware(a.settingsMgr.GetSettings)
 	th := util_session.WithAuthMiddleware(a.DisableAuth, a.sessionMgr, terminal)
 	mux.Handle("/terminal", th)

util/session/sessionmanager.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/coreos/go-oidc/v3/oidc"
@@ -41,6 +42,7 @@ type SessionManager struct {
 	storage                       UserStateStorage
 	sleep                         func(d time.Duration)
 	verificationDelayNoiseEnabled bool
+	failedLock                    sync.RWMutex
 }
 
 // LoginAttempts is a timestamped counter for failed login attempts
@@ -69,7 +71,7 @@ const (
 	// Maximum length of username, too keep the cache's memory signature low
 	maxUsernameLength = 32
 	// The default maximum session cache size
-	defaultMaxCacheSize = 1000
+	defaultMaxCacheSize = 10000
 	// The default number of maximum login failures before delay kicks in
 	defaultMaxLoginFailures = 5
 	// The default time in seconds for the failure window
@@ -284,7 +286,7 @@ func (mgr *SessionManager) Parse(tokenString string) (jwt.Claims, string, error)
 	return token.Claims, newToken, nil
 }
 
-// GetLoginFailures retrieves the login failure information from the cache
+// GetLoginFailures retrieves the login failure information from the cache. Any modifications to the LoginAttemps map must be done in a thread-safe manner.
 func (mgr *SessionManager) GetLoginFailures() map[string]LoginAttempts {
 	// Get failures from the cache
 	var failures map[string]LoginAttempts
@@ -299,25 +301,43 @@ func (mgr *SessionManager) GetLoginFailures() map[string]LoginAttempts {
 	return failures
 }
 
-func expireOldFailedAttempts(maxAge time.Duration, failures *map[string]LoginAttempts) int {
+func expireOldFailedAttempts(maxAge time.Duration, failures map[string]LoginAttempts) int {
 	expiredCount := 0
-	for key, attempt := range *failures {
+	for key, attempt := range failures {
 		if time.Since(attempt.LastFailed) > maxAge*time.Second {
 			expiredCount += 1
-			delete(*failures, key)
+			delete(failures, key)
 		}
 	}
 	return expiredCount
 }
 
+// Protect admin user from login attempt reset caused by attempts to overflow cache in a brute force attack. Instead remove random non-admin to make room in cache.
+func pickRandomNonAdminLoginFailure(failures map[string]LoginAttempts, username string) *string {
+	idx := rand.Intn(len(failures) - 1)
+	i := 0
+	for key := range failures {
+		if i == idx {
+			if key == common.ArgoCDAdminUsername || key == username {
+				return pickRandomNonAdminLoginFailure(failures, username)
+			}
+			return &key
+		}
+		i++
+	}
+	return nil
+}
+
 // Updates the failure count for a given username. If failed is true, increases the counter. Otherwise, sets counter back to 0.
 func (mgr *SessionManager) updateFailureCount(username string, failed bool) {
+	mgr.failedLock.Lock()
+	defer mgr.failedLock.Unlock()
 
 	failures := mgr.GetLoginFailures()
 
 	// Expire old entries in the cache if we have a failure window defined.
 	if window := getLoginFailureWindow(); window > 0 {
-		count := expireOldFailedAttempts(window, &failures)
+		count := expireOldFailedAttempts(window, failures)
 		if count > 0 {
 			log.Infof("Expired %d entries from session cache due to max age reached", count)
 		}
@@ -332,18 +352,11 @@ func (mgr *SessionManager) updateFailureCount(username string, failed bool) {
 	// chance is low (1:cache_size)
 	if failed && len(failures) >= getMaximumCacheSize() {
 		log.Warnf("Session cache size exceeds %d entries, removing random entry", getMaximumCacheSize())
-		idx := rand.Intn(len(failures) - 1)
-		var rmUser string
-		i := 0
-		for key := range failures {
-			if i == idx {
-				rmUser = key
-				delete(failures, key)
-				break
-			}
-			i++
+		rmUser := pickRandomNonAdminLoginFailure(failures, username)
+		if rmUser != nil {
+			delete(failures, *rmUser)
+			log.Infof("Deleted entry for user %s from cache", *rmUser)
 		}
-		log.Infof("Deleted entry for user %s from cache", rmUser)
 	}
 
 	attempt, ok := failures[username]
@@ -374,6 +387,8 @@ func (mgr *SessionManager) updateFailureCount(username string, failed bool) {
 
 // Get the current login failure attempts for given username
 func (mgr *SessionManager) getFailureCount(username string) LoginAttempts {
+	mgr.failedLock.RLock()
+	defer mgr.failedLock.RUnlock()
 	failures := mgr.GetLoginFailures()
 	attempt, ok := failures[username]
 	if !ok {

util/session/sessionmanager_test.go

@@ -1173,3 +1173,42 @@ requestedScopes: ["oidc"]`, oidcTestServer.URL),
 		assert.ErrorIs(t, err, common.TokenVerificationErr)
 	})
 }
+
+func Test_PickFailureAttemptWhenOverflowed(t *testing.T) {
+	t.Run("Not pick admin user from the queue", func(t *testing.T) {
+		failures := map[string]LoginAttempts{
+			"admin": {
+				FailCount: 1,
+			},
+			"test2": {
+				FailCount: 1,
+			},
+		}
+
+		// inside pickRandomNonAdminLoginFailure, it uses random, so we need to test it multiple times
+		for i := 0; i < 1000; i++ {
+			user := pickRandomNonAdminLoginFailure(failures, "test")
+			assert.Equal(t, "test2", *user)
+		}
+	})
+
+	t.Run("Not pick admin user and current user from the queue", func(t *testing.T) {
+		failures := map[string]LoginAttempts{
+			"test": {
+				FailCount: 1,
+			},
+			"admin": {
+				FailCount: 1,
+			},
+			"test2": {
+				FailCount: 1,
+			},
+		}
+
+		// inside pickRandomNonAdminLoginFailure, it uses random, so we need to test it multiple times
+		for i := 0; i < 1000; i++ {
+			user := pickRandomNonAdminLoginFailure(failures, "test")
+			assert.Equal(t, "test2", *user)
+		}
+	})
+}