[TT-5545] issue when retrieving a list of keys attached to an API - alternative solution

gateway/api.go

@@ -27,6 +27,7 @@
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -37,6 +38,7 @@
 	"net/url"
 	"os"
 	"path/filepath"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -702,31 +704,91 @@
 	APIKeys []string `json:"keys"`
 }
 
-func (gw *Gateway) handleGetAllKeys(filter string) (interface{}, int) {
-	sessions := gw.GlobalSessionManager.Sessions(filter)
-	if filter != "" {
-		filterB64 := base64.StdEncoding.WithPadding(base64.NoPadding).EncodeToString([]byte(fmt.Sprintf(`{"org":"%s"`, filter)))
-		// Remove last 2 digits to look exact match
-		filterB64 = filterB64[0 : len(filterB64)-2]
-		orgIDB64Sessions := gw.GlobalSessionManager.Sessions(filterB64)
-		sessions = append(sessions, orgIDB64Sessions...)
+func (gw *Gateway) handleGetAllKeys(c context.Context, filter string, apiID string, hashed bool) (interface{}, int) {
+	keys := gw.getAllSessionKeys(filter)
+	validSessions := make([]string, 0)
+
+	if apiID == "" {
+		for _, k := range keys {
+			if !strings.HasPrefix(k, QuotaKeyPrefix) && !strings.HasPrefix(k, RateLimitKeyPrefix) {
+				validSessions = append(validSessions, k)
+			}
+		}
+		return apiAllKeys{validSessions}, http.StatusOK
 	}
 
-	fixedSessions := make([]string, 0)
-	for _, s := range sessions {
-		if !strings.HasPrefix(s, QuotaKeyPrefix) && !strings.HasPrefix(s, RateLimitKeyPrefix) {
-			fixedSessions = append(fixedSessions, s)
+	const batchSize = 1000
+	for i := 0; i < len(keys); i += batchSize {
+		select {
+		case <-c.Done():
+			log.WithError(c.Err()).Warn("Context finished while processing keys")
+			return apiError("Request processing terminated"), http.StatusGatewayTimeout
+		default:
+		}
+
+		end := i + batchSize
+		if end > len(keys) {
+			end = len(keys)
+		}
+
+		batch := make([]string, 0, batchSize)
+		for _, k := range keys[i:end] {
+			if !strings.HasPrefix(k, QuotaKeyPrefix) && !strings.HasPrefix(k, RateLimitKeyPrefix) {
+				batch = append(batch, k)
+			}
+		}
+
+		if len(batch) == 0 {
+			continue
+		}
+
+		sessions, err := gw.GlobalSessionManager.SessionDetailBulk(filter, batch, hashed)
+		if err != nil {
+			log.WithError(err).Error("Failed to fetch sessions from storage")
+			return apiError("Internal Server Error"), http.StatusInternalServerError
+		}
+
+		for _, key := range batch {
+			session, found := sessions[key]
+			if !found {
+				continue
+			}
+
+			if _, ok := session.AccessRights[apiID]; ok {
+				validSessions = append(validSessions, key)
+				continue
+			}
+
+			if len(session.AccessRights) == 0 && gw.GetConfig().AllowMasterKeys {
+				validSessions = append(validSessions, key)
+			}
 		}
 	}
 
-	sessionsObj := apiAllKeys{fixedSessions}
+	sort.Strings(validSessions)
 
 	log.WithFields(logrus.Fields{
 		"prefix": "api",
 		"status": "ok",
+		"count":  len(validSessions),
 	}).Info("Retrieved key list.")
 
-	return sessionsObj, http.StatusOK
+	return apiAllKeys{validSessions}, http.StatusOK
+}
+
+func (gw *Gateway) getAllSessionKeys(filter string) []string {
+	if strings.ContainsAny(filter, "*?\"{}[]") {
+		log.WithField("filter", filter).Warn("Rejected potentially unsafe filter characters")
+		return []string{}
+	}
+	keys := gw.GlobalSessionManager.Sessions(filter)
+	if filter != "" {
+		filterB64 := base64.StdEncoding.WithPadding(base64.NoPadding).EncodeToString([]byte(fmt.Sprintf(`{"org":"%s"`, filter)))
+		filterB64 = filterB64[0 : len(filterB64)-2]
+		orgIDB64Sessions := gw.GlobalSessionManager.Sessions(filterB64)
+		keys = append(keys, orgIDB64Sessions...)
+	}
+	return keys
 }
 
 func (gw *Gateway) handleAddKey(keyName, sessionString, orgId string) {
@@ -1688,11 +1750,10 @@
 					return
 				}
 
-				// we don't use filter for hashed keys
-				obj, code = gw.handleGetAllKeys("")
+				obj, code = gw.handleGetAllKeys(r.Context(), "", apiID, true)
 			} else {
 				filter := r.URL.Query().Get("filter")
-				obj, code = gw.handleGetAllKeys(filter)
+				obj, code = gw.handleGetAllKeys(r.Context(), filter, apiID, false)
 			}
 		}
 

gateway/api_test.go

@@ -4203,3 +4203,84 @@ func TestPurgeOAuthClientTokensEndpoint(t *testing.T) {
 		assertTokensLen(t, storageManager, storageKey2, 0)
 	})
 }
+
+// TestKeyHandler_BatchFiltering_Integration is only able to cover the fallback redis standalone case
+// as redis cluster is not running in the pipeline
+func TestKeyHandler_BatchFiltering_Integration(t *testing.T) {
+	ts := StartTest(func(globalConf *config.Config) {
+		globalConf.HashKeys = true
+		globalConf.EnableHashedKeysListing = true
+		globalConf.AllowMasterKeys = true
+	})
+	defer ts.Close()
+
+	keyA := "key-access-api-one"
+	hashA := storage.HashKey(keyA, true)
+	sessionA := CreateStandardSession()
+	sessionA.AccessRights = map[string]user.AccessDefinition{
+		"api-one": {APIID: "api-one", Versions: []string{"Default"}},
+	}
+	assert.NoError(t, ts.Gw.GlobalSessionManager.UpdateSession(hashA, sessionA, 100, true))
+
+	keyB := "key-access-api-two"
+	hashB := storage.HashKey(keyB, true)
+	sessionB := CreateStandardSession()
+	sessionB.AccessRights = map[string]user.AccessDefinition{
+		"api-two": {APIID: "api-two", Versions: []string{"Default"}},
+	}
+	assert.NoError(t, ts.Gw.GlobalSessionManager.UpdateSession(hashB, sessionB, 100, true))
+
+	keyC := "key-master"
+	hashC := storage.HashKey(keyC, true)
+	sessionC := CreateStandardSession()
+	sessionC.AccessRights = map[string]user.AccessDefinition{}
+	assert.NoError(t, ts.Gw.GlobalSessionManager.UpdateSession(hashC, sessionC, 100, true))
+
+	t.Cleanup(func() {
+		ts.Gw.GlobalSessionManager.RemoveSession("", keyA, true)
+		ts.Gw.GlobalSessionManager.RemoveSession("", keyB, true)
+		ts.Gw.GlobalSessionManager.RemoveSession("", keyC, true)
+	})
+
+	t.Run("Filter API One - Should find Key A and Master Key", func(t *testing.T) {
+		uri := "/tyk/keys/?api_id=api-one"
+		req := ts.withAuth(TestReq(t, "GET", uri, nil))
+		rec := httptest.NewRecorder()
+		ts.mainRouter().ServeHTTP(rec, req)
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		body := rec.Body.String()
+
+		assert.Contains(t, body, hashA, "Should contain key for api-one")
+		assert.Contains(t, body, hashC, "Should contain master key")
+		assert.NotContains(t, body, hashB, "Should NOT contain key for api-two")
+	})
+
+	t.Run("Filter API Two - Should find Key B and Master Key", func(t *testing.T) {
+		uri := "/tyk/keys/?api_id=api-two"
+		req := ts.withAuth(TestReq(t, "GET", uri, nil))
+		rec := httptest.NewRecorder()
+		ts.mainRouter().ServeHTTP(rec, req)
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		body := rec.Body.String()
+
+		assert.NotContains(t, body, hashA, "Should NOT contain key for api-one")
+		assert.Contains(t, body, hashB, "Should contain key for api-two")
+		assert.Contains(t, body, hashC, "Should contain master key")
+	})
+
+	t.Run("Filter Unknown API - Should only find Master Key", func(t *testing.T) {
+		uri := "/tyk/keys/?api_id=api-unknown"
+		req := ts.withAuth(TestReq(t, "GET", uri, nil))
+		rec := httptest.NewRecorder()
+		ts.mainRouter().ServeHTTP(rec, req)
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		body := rec.Body.String()
+
+		assert.NotContains(t, body, hashA)
+		assert.NotContains(t, body, hashB)
+		assert.Contains(t, body, hashC, "Master key is valid for all APIs")
+	})
+}

gateway/auth_manager.go

@@ -21,6 +21,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)
+	SessionDetailBulk(orgID string, keyNames []string, hashed bool) (map[string]user.SessionState, error)
 	KeyExpired(newSession *user.SessionState) bool
 	Sessions(filter string) []string
 	ResetQuota(string, *user.SessionState, bool)
@@ -215,6 +216,95 @@ func (b *DefaultSessionManager) SessionDetail(orgID string, keyName string, hash
 	return session.Clone(), true
 }
 
+func (b *DefaultSessionManager) SessionDetailBulk(orgID string, keyNames []string, hashed bool) (map[string]user.SessionState, error) {
+	result := make(map[string]user.SessionState)
+	if len(keyNames) == 0 {
+		return result, nil
+	}
+
+	if hashed {
+		prefix := b.store.GetKeyPrefix()
+		prefixedKeys := make([]string, len(keyNames))
+		for i, keyName := range keyNames {
+			prefixedKeys[i] = prefix + keyName
+		}
+
+		jsonValues, err := b.store.GetRawMultiKey(prefixedKeys)
+		if err != nil {
+			log.WithError(err).Debug("Failed to bulk fetch hashed sessions")
+			return nil, err
+		}
+
+		for i, jsonVal := range jsonValues {
+			if jsonVal == "" {
+				continue
+			}
+			session := &user.SessionState{}
+			if err := json.Unmarshal([]byte(jsonVal), session); err != nil {
+				log.WithField("key", keyNames[i]).Error("Failed to unmarshal session in bulk fetch")
+				continue
+			}
+			session.KeyID = keyNames[i]
+			result[keyNames[i]] = *session
+		}
+		return result, nil
+	}
+
+	allKeysToSearch := make([]string, 0, len(keyNames)*2)
+	searchToOriginal := make(map[string]string)
+
+	for _, keyName := range keyNames {
+		if storage.TokenOrg(keyName) != orgID {
+			if !b.Gw.GetConfig().DisableKeyActionsByUsername {
+				legacyKey := b.Gw.generateToken(orgID, keyName)
+				allKeysToSearch = append(allKeysToSearch, legacyKey)
+				searchToOriginal[legacyKey] = keyName
+			}
+			allKeysToSearch = append(allKeysToSearch, keyName)
+			searchToOriginal[keyName] = keyName
+
+			for _, fallback := range b.Gw.GetConfig().HashKeyFunctionFallback {
+				if !b.Gw.GetConfig().DisableKeyActionsByUsername {
+					fallbackKey := b.Gw.generateToken(orgID, keyName, fallback)
+					allKeysToSearch = append(allKeysToSearch, fallbackKey)
+					searchToOriginal[fallbackKey] = keyName
+				}
+			}
+		} else {
+			allKeysToSearch = append(allKeysToSearch, keyName)
+			searchToOriginal[keyName] = keyName
+		}
+	}
+
+	jsonValues, err := b.store.GetMultiKey(allKeysToSearch)
+	if err != nil {
+		log.WithError(err).Debug("Failed to bulk fetch legacy sessions")
+		return nil, err
+	}
+
+	for i, jsonVal := range jsonValues {
+		if jsonVal == "" {
+			continue
+		}
+		foundKey := allKeysToSearch[i]
+		originalKey := searchToOriginal[foundKey]
+
+		if _, exists := result[originalKey]; exists {
+			continue
+		}
+
+		session := &user.SessionState{}
+		if err := json.Unmarshal([]byte(jsonVal), session); err != nil {
+			log.WithField("key", foundKey).WithError(err).Error("Failed to unmarshal session in bulk fetch")
+			continue
+		}
+		session.KeyID = foundKey
+		result[originalKey] = *session
+	}
+
+	return result, nil
+}
+
 func (b *DefaultSessionManager) Stop() {}
 
 // Sessions returns all sessions in the key store that match a filter key (a prefix)