[TT-15954] Make org session fetch non-blocking

gateway/auth_manager.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"strings"
@@ -21,6 +22,7 @@ type SessionHandler interface {
 	UpdateSession(keyName string, session *user.SessionState, resetTTLTo int64, hashed bool) error
 	RemoveSession(orgID string, keyName string, hashed bool) bool
 	SessionDetail(orgID string, keyName string, hashed bool) (user.SessionState, bool)
+	SessionDetailContext(ctx context.Context, orgID string, keyName string, hashed bool) (user.SessionState, bool)
 	KeyExpired(newSession *user.SessionState) bool
 	Sessions(filter string) []string
 	ResetQuota(string, *user.SessionState, bool)
@@ -215,6 +217,73 @@ func (b *DefaultSessionManager) SessionDetail(orgID string, keyName string, hash
 	return session.Clone(), true
 }
 
+// SessionDetailContext returns the session detail using the storage engine with context support for cancellation
+func (b *DefaultSessionManager) SessionDetailContext(ctx context.Context, orgID string, keyName string, hashed bool) (user.SessionState, bool) {
+	select {
+	case <-ctx.Done():
+		log.WithFields(logrus.Fields{
+			"prefix":      "auth-mgr",
+			"inbound-key": b.Gw.obfuscateKey(keyName),
+		}).Debug("Context cancelled")
+		return user.SessionState{}, false
+	default:
+	}
+
+	var jsonKeyVal string
+	var err error
+	keyId := keyName
+
+	if hashed {
+		jsonKeyVal, err = b.store.GetRawKeyContext(ctx, b.store.GetKeyPrefix()+keyName)
+	} else {
+		if storage.TokenOrg(keyName) != orgID {
+			// try to get legacy and new format key at once
+			toSearchList := []string{}
+			if !b.Gw.GetConfig().DisableKeyActionsByUsername {
+				toSearchList = append(toSearchList, b.Gw.generateToken(orgID, keyName))
+			}
+
+			toSearchList = append(toSearchList, keyName)
+			for _, fallback := range b.Gw.GetConfig().HashKeyFunctionFallback {
+				if !b.Gw.GetConfig().DisableKeyActionsByUsername {
+					toSearchList = append(toSearchList, b.Gw.generateToken(orgID, keyName, fallback))
+				}
+			}
+
+			var jsonKeyValList []string
+
+			jsonKeyValList, err = b.store.GetMultiKeyContext(ctx, toSearchList)
+			// pick the 1st non empty from the returned list
+			for idx, val := range jsonKeyValList {
+				if val != "" {
+					jsonKeyVal = val
+					keyId = toSearchList[idx]
+					break
+				}
+			}
+		} else {
+			// key is not an imported one
+			jsonKeyVal, err = b.store.GetKeyContext(ctx, keyName)
+		}
+	}
+
+	if err != nil {
+		log.WithFields(logrus.Fields{
+			"prefix":      "auth-mgr",
+			"inbound-key": b.Gw.obfuscateKey(keyName),
+			"err":         err,
+		}).Debug("Could not get session detail, key not found")
+		return user.SessionState{}, false
+	}
+	session := &user.SessionState{}
+	if err := json.Unmarshal([]byte(jsonKeyVal), &session); err != nil {
+		log.Error("Couldn't unmarshal session object (may be cache miss): ", err)
+		return user.SessionState{}, false
+	}
+	session.KeyID = keyId
+	return session.Clone(), true
+}
+
 func (b *DefaultSessionManager) Stop() {}
 
 // Sessions returns all sessions in the key store that match a filter key (a prefix)

gateway/auth_manager_test.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"crypto/x509"
 	"encoding/json"
 	"fmt"
@@ -578,3 +579,15 @@ func (c *countingStorageHandler) AppendToSet(s string, s2 string) {}
 func (c *countingStorageHandler) Exists(s string) (bool, error) {
 	return false, nil
 }
+
+func (c *countingStorageHandler) GetKeyContext(ctx context.Context, s string) (string, error) {
+	return "", nil
+}
+
+func (c *countingStorageHandler) GetRawKeyContext(ctx context.Context, key string) (string, error) {
+	return "", nil
+}
+
+func (c *countingStorageHandler) GetMultiKeyContext(ctx context.Context, keys []string) ([]string, error) {
+	return nil, nil
+}

gateway/ldap_auth_handler.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"errors"
 	"strings"
 
@@ -242,3 +243,18 @@
 	log.Error("Not implemented")
 	return false, nil
 }
+
+// GetKeyContext retrieves a key with context support.
+func (l *LDAPStorageHandler) GetKeyContext(ctx context.Context, keyName string) (string, error) {
+	return l.GetKey(keyName)
+}
+
+// GetRawKeyContext retrieves a raw key with context support.
+func (l *LDAPStorageHandler) GetRawKeyContext(ctx context.Context, keyName string) (string, error) {
+	return l.GetRawKey(keyName)
+}
+
+// GetMultiKeyContext retrieves multiple keys with context support.
+func (l *LDAPStorageHandler) GetMultiKeyContext(ctx context.Context, keyNames []string) ([]string, error) {
+	return l.GetMultiKey(keyNames)
+}

gateway/middleware.go

@@ -39,7 +39,8 @@
 
 const (
 	DEFAULT_ORG_SESSION_EXPIRATION = int64(604800)
+	orgSessionFetchTimeout         = 2 * time.Second
 )
 
 var (
 	GlobalRate            = ratecounter.NewRateCounter(1 * time.Second)
@@ -343,55 +344,113 @@
 	return t.Spec
 }
 
-func (t *BaseMiddleware) OrgSession(orgID string) (user.SessionState, bool) {
+// fetchOrgSessionWithTimeout fetches org session with a timeout to prevent hanging
+func (t *BaseMiddleware) fetchOrgSessionWithTimeout(orgID string) (user.SessionState, bool) {
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), orgSessionFetchTimeout)
+	defer cancel()
+
+	resultChan := make(chan struct {
+		session user.SessionState
+		found   bool
+	}, 1)
+
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				t.Logger().Errorf("Panic recovered during org session fetch for org %s: %v", orgID, r)
+				select {
+				case resultChan <- struct {
+					session user.SessionState
+					found   bool
+				}{session: user.SessionState{}, found: false}:
+				case <-timeoutCtx.Done():
+				}
+			}
+		}()
 
-	if rpc.IsEmergencyMode() {
+		session, found := t.Spec.OrgSessionManager.SessionDetailContext(timeoutCtx, orgID, orgID, false)
+		if found && t.Spec.GlobalConfig.EnforceOrgDataAge {
+			t.Logger().Debug("Setting data expiry: ", orgID)
+			t.Gw.ExpiryCache.Set(session.OrgID, session.DataExpires, cache.DefaultExpiration)
+		}
+		session.SetKeyHash(storage.HashKey(orgID, t.Gw.GetConfig().HashKeys))
+
+		select {
+		case resultChan <- struct {
+			session user.SessionState
+			found   bool
+		}{session: session.Clone(), found: found}:
+		case <-timeoutCtx.Done():
+			return
+		}
+	}()
+
+	select {
+	case res := <-resultChan:
+		return res.session, res.found
+	case <-timeoutCtx.Done():
+		t.Logger().Warning("Org session fetch timed out after 2s")
 		return user.SessionState{}, false
 	}
+}
 
-	// Try and get the session from the session store
-	session, found := t.Spec.OrgSessionManager.SessionDetail(orgID, orgID, false)
-	if found && t.Spec.GlobalConfig.EnforceOrgDataAge {
-		// If exists, assume it has been authorized and pass on
-		// We cache org expiry data
-		t.Logger().Debug("Setting data expiry: ", orgID)
-
-		t.Gw.ExpiryCache.Set(session.OrgID, session.DataExpires, cache.DefaultExpiration)
+func (t *BaseMiddleware) OrgSession(orgID string) (user.SessionState, bool) {
+	if rpc.IsEmergencyMode() {
+		return user.SessionState{}, false
 	}
 
-	session.SetKeyHash(storage.HashKey(orgID, t.Gw.GetConfig().HashKeys))
+	// Fetch with timeout to prevent blocking
+	t.Logger().Debug("Fetching org session with timeout")
+	session, found := t.fetchOrgSessionWithTimeout(orgID)
 
-	return session.Clone(), found
+	return session, found
 }
 
 func (t *BaseMiddleware) SetOrgExpiry(orgid string, expiry int64) {
 	t.Gw.ExpiryCache.Set(orgid, expiry, cache.DefaultExpiration)
 }
 
 func (t *BaseMiddleware) OrgSessionExpiry(orgid string) int64 {
 	t.Logger().Debug("Checking: ", orgid)
 
-	// Cache failed attempt
-	id, err, _ := orgSessionExpiryCache.Do(orgid, func() (interface{}, error) {
-		cachedVal, found := t.Gw.ExpiryCache.Get(orgid)
-		if found {
-			return cachedVal, nil
+	// Check cache first
+	cachedVal, found := t.Gw.ExpiryCache.Get(orgid)
+	if found {
+		val, ok := cachedVal.(int64)
+		if !ok {
+			t.Logger().Error("Type assertion failed")
+			return DEFAULT_ORG_SESSION_EXPIRATION
 		}
+		return val
+	}
 
-		s, found := t.OrgSession(orgid)
-		if found && t.Spec.GlobalConfig.EnforceOrgDataAge {
-			return s.DataExpires, nil
-		}
-		return 0, errors.New("missing session")
+	// Cache miss
+	go orgSessionExpiryCache.Do(orgid, func() (interface{}, error) {
+		return t.refreshOrgSessionExpiry(orgid)
 	})
 
-	if err != nil {
-		t.Logger().Debug("no cached entry found, returning 7 days")
-		t.SetOrgExpiry(orgid, DEFAULT_ORG_SESSION_EXPIRATION)
-		return DEFAULT_ORG_SESSION_EXPIRATION
+	t.Logger().Debug("no cached entry found, returning 7 days (async refresh started)")
+	return DEFAULT_ORG_SESSION_EXPIRATION
+}
+
+// refreshOrgSessionExpiry fetches org session expiry in the background
+func (t *BaseMiddleware) refreshOrgSessionExpiry(orgid string) (interface{}, error) {
+	defer func() {
+		if r := recover(); r != nil {
+			t.Logger().Errorf("Panic recovered during org session expiry refresh for org %s: %v", orgid, r)
+		}
+	}()
+
+	s, found := t.OrgSession(orgid) // RPC call happens in background
+	if found && t.Spec.GlobalConfig.EnforceOrgDataAge {
+		t.Logger().Debug("Background refresh: setting data expiry for org: ", orgid)
+		t.SetOrgExpiry(orgid, s.DataExpires)
+		return s.DataExpires, nil
 	}
 
-	return id.(int64)
+	t.Logger().Debug("Background refresh: org session not found, setting default expiry")
+	t.SetOrgExpiry(orgid, DEFAULT_ORG_SESSION_EXPIRATION)
+	return DEFAULT_ORG_SESSION_EXPIRATION, nil
 }
 
 func (t *BaseMiddleware) UpdateRequestSession(r *http.Request) bool {