Handle quota reset for RPC key upadates (TykTechnologies#2408)

Fix TykTechnologies#2327
anuvu · Jul 16, 2019 · 692446c · 692446c
1 parent 82f94a5
commit 692446c
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 14 deletions.
diff --git a/gateway/api.go b/gateway/api.go
@@ -555,7 +555,7 @@ func handleAddKey(keyName, hashedName, sessionString, apiID string) {
 	}).Info("Updated hashed key in slave storage.")
 }
 
-func handleDeleteKey(keyName, apiID string) (interface{}, int) {
+func handleDeleteKey(keyName, apiID string, resetQuota bool) (interface{}, int) {
 	if apiID == "-1" {
 		// Go through ALL managed API's and delete the key
 		apisMu.RLock()
@@ -601,7 +601,10 @@ func handleDeleteKey(keyName, apiID string) (interface{}, int) {
 		}).Error("Failed to remove the key")
 		return apiError("Failed to remove the key"), http.StatusBadRequest
 	}
-	sessionManager.ResetQuota(keyName, &user.SessionState{}, false)
+
+	if resetQuota {
+		sessionManager.ResetQuota(keyName, &user.SessionState{}, false)
+	}
 
 	statusObj := apiModifyKeySuccess{
 		Key:    keyName,
@@ -624,7 +627,7 @@ func handleDeleteKey(keyName, apiID string) (interface{}, int) {
 	return statusObj, http.StatusOK
 }
 
-func handleDeleteHashedKey(keyName, apiID string) (interface{}, int) {
+func handleDeleteHashedKey(keyName, apiID string, resetQuota bool) (interface{}, int) {
 	if apiID == "-1" {
 		// Go through ALL managed API's and delete the key
 		removed := false
@@ -668,6 +671,10 @@ func handleDeleteHashedKey(keyName, apiID string) (interface{}, int) {
 		return apiError("Failed to remove the key"), http.StatusBadRequest
 	}
 
+	if resetQuota {
+		sessionManager.ResetQuota(keyName, &user.SessionState{}, true)
+	}
+
 	statusObj := apiModifyKeySuccess{
 		Key:    keyName,
 		Status: "ok",
@@ -875,16 +882,16 @@ func keyHandler(w http.ResponseWriter, r *http.Request) {
 	case http.MethodDelete:
 		// Remove a key
 		if !isHashed {
-			obj, code = handleDeleteKey(keyName, apiID)
+			obj, code = handleDeleteKey(keyName, apiID, true)
 		} else {
-			obj, code = handleDeleteHashedKey(keyName, apiID)
+			obj, code = handleDeleteHashedKey(keyName, apiID, true)
 		}
 		if code != http.StatusOK && hashKeyFunction != "" {
 			// try to use legacy key format
 			if !isHashed {
-				obj, code = handleDeleteKey(origKeyName, apiID)
+				obj, code = handleDeleteKey(origKeyName, apiID, true)
 			} else {
-				obj, code = handleDeleteHashedKey(origKeyName, apiID)
+				obj, code = handleDeleteHashedKey(origKeyName, apiID, true)
 			}
 		}
 	}

diff --git a/gateway/rpc_storage_handler.go b/gateway/rpc_storage_handler.go
@@ -819,18 +819,26 @@ func getSessionAndCreate(keyName string, r *RPCStorageHandler) {
 }
 
 func (r *RPCStorageHandler) ProcessKeySpaceChanges(keys []string) {
+	keysToReset := map[string]bool{}
+
+	for _, key := range keys {
+		splitKeys := strings.Split(key, ":")
+		if len(splitKeys) > 1 && splitKeys[1] == "resetQuota" {
+			keysToReset[splitKeys[0]] = true
+		}
+	}
+
 	for _, key := range keys {
 		splitKeys := strings.Split(key, ":")
-		if len(splitKeys) > 1 {
+		_, resetQuota := keysToReset[splitKeys[0]]
+		if len(splitKeys) > 1 && splitKeys[1] == "hashed" {
 			key = splitKeys[0]
-			if splitKeys[1] == "hashed" {
-				log.Info("--> removing cached (hashed) key: ", splitKeys[0])
-				handleDeleteHashedKey(splitKeys[0], "")
-				getSessionAndCreate(splitKeys[0], r)
-			}
+			log.Info("--> removing cached (hashed) key: ", splitKeys[0])
+			handleDeleteHashedKey(splitKeys[0], "", resetQuota)
+			getSessionAndCreate(splitKeys[0], r)
 		} else {
 			log.Info("--> removing cached key: ", key)
-			handleDeleteKey(key, "-1")
+			handleDeleteKey(key, "-1", resetQuota)
 			getSessionAndCreate(splitKeys[0], r)
 		}
 		SessionCache.Delete(key)