Avoid unnecessary/buggy readHead access during command processing

libs/server/Resp/Objects/ListCommands.cs

@@ -226,10 +226,7 @@ private unsafe bool ListPopMultiple<TGarnetApi>(ref TGarnetApi storageApi)
             }
 
             // Get the direction
-            var dir = parseState.GetArgSliceByRef(currTokenId++);
-            var popDirection = GetOperationDirection(dir);
-
-            if (popDirection == OperationDirection.Unknown)
+            if (!parseState.TryGetOperationDirection(currTokenId++, out var popDirection))
             {
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
             }
@@ -355,8 +352,12 @@ private unsafe bool ListBlockingMove()
 
             var srcKey = parseState.GetArgSliceByRef(0);
             var dstKey = parseState.GetArgSliceByRef(1);
-            var srcDir = parseState.GetArgSliceByRef(2);
-            var dstDir = parseState.GetArgSliceByRef(3);
+
+            if (!parseState.TryGetOperationDirection(2, out var srcDir) ||
+                !parseState.TryGetOperationDirection(3, out var dstDir))
+            {
+                return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
+            }
 
             if (!parseState.TryGetTimeout(4, out var timeout, out var error))
             {
@@ -379,27 +380,26 @@ private bool ListBlockingPopPush()
 
             var srcKey = parseState.GetArgSliceByRef(0);
             var dstKey = parseState.GetArgSliceByRef(1);
-            var rightOption = ArgSlice.FromPinnedSpan(CmdStrings.RIGHT);
-            var leftOption = ArgSlice.FromPinnedSpan(CmdStrings.LEFT);
 
             if (!parseState.TryGetTimeout(2, out var timeout, out var error))
             {
                 return AbortWithErrorMessage(error);
             }
 
-            return ListBlockingMove(srcKey, dstKey, rightOption, leftOption, timeout);
+            return ListBlockingMove(srcKey, dstKey, OperationDirection.Right,
+                                    OperationDirection.Left, timeout);
         }
 
-        private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey, ArgSlice srcDir, ArgSlice dstDir, double timeout)
+        private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey,
+                                      OperationDirection sourceDirection,
+                                      OperationDirection destinationDirection,
+                                      double timeout)
         {
             var cmdArgs = new ArgSlice[] { default, default, default };
 
             // Read destination key
             cmdArgs[0] = dstKey;
 
-            var sourceDirection = GetOperationDirection(srcDir);
-            var destinationDirection = GetOperationDirection(dstDir);
-
             if (sourceDirection == OperationDirection.Unknown || destinationDirection == OperationDirection.Unknown)
             {
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
@@ -770,13 +770,8 @@ private bool ListMove<TGarnetApi>(ref TGarnetApi storageApi)
             var srcKey = parseState.GetArgSliceByRef(0);
             var dstKey = parseState.GetArgSliceByRef(1);
 
-            var srcDirSlice = parseState.GetArgSliceByRef(2);
-            var dstDirSlice = parseState.GetArgSliceByRef(3);
-
-            var sourceDirection = GetOperationDirection(srcDirSlice);
-            var destinationDirection = GetOperationDirection(dstDirSlice);
-
-            if (sourceDirection == OperationDirection.Unknown || destinationDirection == OperationDirection.Unknown)
+            if (!parseState.TryGetOperationDirection(2, out var sourceDirection) ||
+                !parseState.TryGetOperationDirection(3, out var destinationDirection))
             {
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
             }
@@ -964,10 +959,7 @@ private unsafe bool ListBlockingPopMultiple()
             var cmdArgs = new ArgSlice[2];
 
             // Get the direction
-            var dir = parseState.GetArgSliceByRef(currTokenId++);
-            var popDirection = GetOperationDirection(dir);
-
-            if (popDirection == OperationDirection.Unknown)
+            if (!parseState.TryGetOperationDirection(currTokenId++, out var popDirection))
             {
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
             }

libs/server/Resp/Objects/ObjectStoreUtils.cs

@@ -99,27 +99,5 @@ private bool AbortWithErrorMessage(string format, params object[] args)
         {
             return AbortWithErrorMessage(Encoding.ASCII.GetBytes(string.Format(format, args)));
         }
-
-        /// <summary>
-        /// Tries to parse the input as "LEFT" or "RIGHT" and returns the corresponding OperationDirection.
-        /// If parsing fails, returns OperationDirection.Unknown.
-        /// </summary>
-        /// <param name="input">The input to parse.</param>
-        /// <returns>The parsed OperationDirection, or OperationDirection.Unknown if parsing fails.</returns>
-        public OperationDirection GetOperationDirection(ArgSlice input)
-        {
-            // Optimize for the common case
-            if (input.ReadOnlySpan.SequenceEqual("LEFT"u8))
-                return OperationDirection.Left;
-            if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8))
-                return OperationDirection.Right;
-            // Rare case: try making upper case and retry
-            MakeUpperCase(input.ptr);
-            if (input.ReadOnlySpan.SequenceEqual("LEFT"u8))
-                return OperationDirection.Left;
-            if (input.ReadOnlySpan.SequenceEqual("RIGHT"u8))
-                return OperationDirection.Right;
-            return OperationDirection.Unknown;
-        }
     }
 }

libs/server/Resp/Parser/RespCommand.cs

@@ -2768,7 +2768,7 @@ private RespCommand ArrayParseCommand(bool writeErrorOnFailure, ref int count, r
             var ptr = recvBufferPtr + readHead;
 
             // See if input command is all upper-case. If not, convert and try fast parse pass again.
-            if (MakeUpperCase(ptr))
+            if (MakeUpperCase(ptr, bytesRead - readHead))
             {
                 cmd = FastParseCommand(out count);
                 if (cmd != RespCommand.NONE)

libs/server/Resp/RespServerSession.cs

@@ -657,7 +657,7 @@ private void ProcessMessages()
         }
 
         // Make first command in string as uppercase
-        private bool MakeUpperCase(byte* ptr)
+        private bool MakeUpperCase(byte* ptr, int len)
         {
             // Assume most commands are already upper case.
             // Assume most commands are 2-8 bytes long.
@@ -672,7 +672,6 @@ private bool MakeUpperCase(byte* ptr)
             // Note that _all_ of these bytes are <= 95 in the common case
             // and there's no need to scan the whole string in those cases.
 
-            var len = bytesRead - readHead;
             if (len >= 12)
             {
                 var cmdLen = (uint)(*(ptr + 5) - '2');
@@ -1188,43 +1187,6 @@ ReadOnlySpan<byte> GetUpperCaseCommand(out bool success)
             return result;
         }
 
-        public ArgSlice GetCommandAsArgSlice(out bool success)
-        {
-            if (bytesRead - readHead < 6)
-            {
-                success = false;
-                return default;
-            }
-
-            Debug.Assert(*(recvBufferPtr + readHead) == '$');
-            int psize = *(recvBufferPtr + readHead + 1) - '0';
-            readHead += 2;
-            while (*(recvBufferPtr + readHead) != '\r')
-            {
-                psize = psize * 10 + *(recvBufferPtr + readHead) - '0';
-                if (bytesRead - readHead < 1)
-                {
-                    success = false;
-                    return default;
-                }
-                readHead++;
-            }
-            if (bytesRead - readHead < 2 + psize + 2)
-            {
-                success = false;
-                return default;
-            }
-            Debug.Assert(*(recvBufferPtr + readHead + 1) == '\n');
-
-            var result = new ArgSlice(recvBufferPtr + readHead + 2, psize);
-            Debug.Assert(*(recvBufferPtr + readHead + 2 + psize) == '\r');
-            Debug.Assert(*(recvBufferPtr + readHead + 2 + psize + 1) == '\n');
-
-            readHead += 2 + psize + 2;
-            success = true;
-            return result;
-        }
-
         /// <summary>
         /// Attempt to kill this session.
         ///

libs/server/SessionParseStateExtensions.cs

@@ -614,6 +614,21 @@ internal static bool TryGetManagerType(this SessionParseState parseState, int id
             return true;
         }
 
+        internal static bool TryGetOperationDirection(this SessionParseState parseState, int idx, out OperationDirection value)
+        {
+            value = OperationDirection.Unknown;
+            var sbArg = parseState.GetArgSliceByRef(idx).ReadOnlySpan;
+
+            if (sbArg.EqualsUpperCaseSpanIgnoringCase("LEFT"u8))
+                value = OperationDirection.Left;
+            else if (sbArg.EqualsUpperCaseSpanIgnoringCase("RIGHT"u8))
+                value = OperationDirection.Right;
+            else
+                return false;
+
+            return true;
+        }
+
         /// <summary>
         /// Parse sorted set add option from parse state at specified index
         /// </summary>

libs/server/Transaction/TxnKeyManager.cs

@@ -489,15 +489,18 @@ private int ListKeys(int inputCount, StoreType storeType, LockType type)
         /// </summary>
         private int ListReadKeysWithCount(LockType type, bool isObject = true)
         {
-            var numKeysArg = respSession.GetCommandAsArgSlice(out var success);
-            if (!success) return -2;
+            if (respSession.parseState.Count == 0)
+                return -2;
 
-            if (!NumUtils.TryParse(numKeysArg.ReadOnlySpan, out int numKeys)) return -2;
+            if (!respSession.parseState.TryGetInt(0, out var numKeys))
+                return -2;
 
-            for (var i = 0; i < numKeys; i++)
+            if (numKeys + 1 > respSession.parseState.Count)
+                return -2;
+
+            for (var i = 1; i < numKeys + 1; i++)
             {
-                var key = respSession.GetCommandAsArgSlice(out success);
-                if (!success) return -2;
+                var key = respSession.parseState.GetArgSliceByRef(i);
                 SaveKeyEntryToLock(key, isObject, type);
                 SaveKeyArgSlice(key);
             }