From b144d934ccc025784261393e36a2b4d97a93eb6f Mon Sep 17 00:00:00 2001
From: Vasileios Zois <vazois@microsoft.com>
Date: Tue, 2 Jul 2024 11:30:59 -0700
Subject: [PATCH] fix edge cases for centralized slot verification

---
 libs/server/Resp/Parser/RespCommand.cs          | 12 +++++++++++-
 libs/server/Resp/RespServerSession.cs           |  6 ++++--
 libs/server/Resp/RespServerSessionSlotVerify.cs |  7 +------
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/libs/server/Resp/Parser/RespCommand.cs b/libs/server/Resp/Parser/RespCommand.cs
index b07c1dd4b1a..febefbe30bf 100644
--- a/libs/server/Resp/Parser/RespCommand.cs
+++ b/libs/server/Resp/Parser/RespCommand.cs
@@ -334,7 +334,17 @@ public static bool IsReadOnly(this RespCommand cmd)
             => cmd <= LastReadCommand();
 
         public static bool IsDataCommand(this RespCommand cmd)
-            => cmd >= FirstReadCommand() && cmd <= LastWriteCommand();
+        {
+            return cmd switch
+            {
+                // TODO: validate if these cases need to be excluded
+                RespCommand.MIGRATE => false,
+                RespCommand.DBSIZE => false,
+                RespCommand.MEMORY_USAGE => false,
+                RespCommand.FLUSHDB => false,
+                _ => cmd >= FirstReadCommand() && cmd <= LastWriteCommand()
+            };
+        }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static bool IsWriteOnly(this RespCommand cmd)
diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs
index 0084b7225ff..94ac2886649 100644
--- a/libs/server/Resp/RespServerSession.cs
+++ b/libs/server/Resp/RespServerSession.cs
@@ -351,7 +351,7 @@ private void ProcessMessages()
                 }
 
                 // Check ACL permissions for the command
-                if (cmd != RespCommand.INVALID && CanServeSlot(cmd) && CheckACLPermissions(cmd))
+                if (cmd != RespCommand.INVALID && CheckACLPermissions(cmd))
                 {
                     if (txnManager.state != TxnState.None)
                     {
@@ -370,7 +370,9 @@ private void ProcessMessages()
                     }
                     else
                     {
-                        _ = ProcessBasicCommands(cmd, ref basicGarnetApi);
+                        // Check if can serve slot for provided keys
+                        if (CanServeSlot(cmd))
+                            _ = ProcessBasicCommands(cmd, ref basicGarnetApi);
                     }
                 }
 
diff --git a/libs/server/Resp/RespServerSessionSlotVerify.cs b/libs/server/Resp/RespServerSessionSlotVerify.cs
index 0e138d57a36..6b1519ac2a6 100644
--- a/libs/server/Resp/RespServerSessionSlotVerify.cs
+++ b/libs/server/Resp/RespServerSessionSlotVerify.cs
@@ -36,13 +36,8 @@ bool CanServeSlot(RespCommand cmd)
             if (clusterSession == null)
                 return true;
 
-            // If this is a UNKNOWN command let process message generate the appropriate response
-            if (cmd == RespCommand.NONE)
-                return true;
-
             // Verify slot for command if it falls into data command category
-            // TODO: Skip validation of MIGRATE until refactoring of parsing is complete.
-            if (!cmd.IsDataCommand() || cmd == RespCommand.MIGRATE)
+            if (!cmd.IsDataCommand())
                 return true;
 
             cmd = cmd.NormalizeForACLs();