Fix parsing for object store commands with wrong number of arguments

libs/server/Resp/Objects/HashCommands.cs

@@ -119,9 +119,7 @@ private unsafe bool HashGet<TGarnetApi>(int count, byte* ptr, HashOperation op,
                 (op == HashOperation.HGET && count != 3) ||
                 (op == HashOperation.HMGET && count < 3))
             {
-                // Send error to output
-                WriteErrorTokenNumberInCommand(op.ToString());
-                ReadLeftToken(count - 1, ref ptr);
+                return AbortWithWrongNumberOfArguments(op.ToString(), count);
             }
             else
             {
@@ -221,11 +219,10 @@ private unsafe bool HashLength<TGarnetApi>(int count, byte* ptr, ref TGarnetApi
         {
             ptr += 10;
 
-            if (count - 2 < 0)
+            if (count != 2)
             {
                 hashItemsDoneCount = hashOpsCount = 0;
-                // Send error to output
-                WriteErrorTokenNumberInCommand("HLEN");
+                return AbortWithWrongNumberOfArguments("HLEN", count);
             }
             else
             {
@@ -293,11 +290,10 @@ private unsafe bool HashDelete<TGarnetApi>(int count, byte* ptr, ref TGarnetApi
         {
             ptr += 10;
 
-            if (count - 2 <= 0)
+            if (count < 2)
             {
                 hashItemsDoneCount = hashOpsCount = 0;
-                // Send error to output
-                WriteErrorTokenNumberInCommand("HDEL");
+                return AbortWithWrongNumberOfArguments("HDEL", count);
             }
             else
             {
@@ -375,11 +371,10 @@ private unsafe bool HashExists<TGarnetApi>(int count, byte* ptr, ref TGarnetApi
         {
             ptr += 13;
 
-            if (count < 3)
+            if (count != 3)
             {
                 hashItemsDoneCount = hashOpsCount = 0;
-                // Send error to output
-                WriteErrorTokenNumberInCommand("HEXISTS");
+                return AbortWithWrongNumberOfArguments("HEXISTS", count);
             }
             else
             {
@@ -534,11 +529,9 @@ private unsafe bool HashIncrement<TGarnetApi>(int count, byte* ptr, HashOperatio
             ptr += op == HashOperation.HINCRBY ? 13 : 19;
 
             // Check if parameters number is right
-            if (count < 4)
+            if (count != 4)
             {
-                // Send error to output
-                WriteErrorTokenNumberInCommand(op == HashOperation.HINCRBY ? "HINCRBY" : "HINCRBYFLOAT");
-                ReadLeftToken(count - 1, ref ptr);
+                return AbortWithWrongNumberOfArguments(op == HashOperation.HINCRBY ? "HINCRBY" : "HINCRBYFLOAT", count);
             }
             else
             {

libs/server/Resp/Objects/ListCommands.cs

@@ -268,11 +268,7 @@ private unsafe bool ListTrim<TGarnetApi>(int count, byte* ptr, ref TGarnetApi st
 
             if (count != 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("LTRIM");
+                return AbortWithWrongNumberOfArguments("LTRIM", count);
             }
             else
             {
@@ -342,11 +338,7 @@ private unsafe bool ListRange<TGarnetApi>(int count, byte* ptr, ref TGarnetApi s
 
             if (count != 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("LRANGE");
+                return AbortWithWrongNumberOfArguments("LRANGE", count);
             }
             else
             {
@@ -424,11 +416,7 @@ private unsafe bool ListIndex<TGarnetApi>(int count, byte* ptr, ref TGarnetApi s
 
             if (count != 3)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("LINDEX");
+                return AbortWithWrongNumberOfArguments("LINDEX", count);
             }
             else
             {
@@ -511,11 +499,7 @@ private unsafe bool ListInsert<TGarnetApi>(int count, byte* ptr, ref TGarnetApi
 
             if (count != 5)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("LINSERT");
+                return AbortWithWrongNumberOfArguments("LINSERT", count);
             }
             else
             {
@@ -599,11 +583,7 @@ private unsafe bool ListRemove<TGarnetApi>(int count, byte* ptr, ref TGarnetApi
             // if params are missing return error
             if (count != 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("LREM");
+                return AbortWithWrongNumberOfArguments("LREM", count);
             }
             else
             {
@@ -684,11 +664,7 @@ private unsafe bool ListMove<TGarnetApi>(int count, byte* ptr, ref TGarnetApi st
 
             if (count != 5)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("LMOVE");
+                return AbortWithWrongNumberOfArguments("LMOVE", count);
             }
             else
             {
@@ -742,11 +718,7 @@ private unsafe bool ListRightPopLeftPush<TGarnetApi>(int count, byte* ptr, ref T
 
             if (count != 3)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("RPOPLPUSH");
+                return AbortWithWrongNumberOfArguments("RPOPLPUSH", count);
             }
             else
             {

libs/server/Resp/Objects/ObjectStoreUtils.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+using System;
 using System.Text;
 using Garnet.common;
 
@@ -11,18 +12,6 @@ namespace Garnet.server
     /// </summary>
     internal sealed unsafe partial class RespServerSession : ServerSessionBase
     {
-        /// <summary>
-        /// Send the error of missing arguments for any command
-        /// </summary>
-        /// <param name="cmdName"></param>
-        //todo move this method to write utils ??
-        private void WriteErrorTokenNumberInCommand(string cmdName)
-        {
-            var errorMessage = Encoding.ASCII.GetBytes($"-ERR wrong number of arguments for {cmdName} command.\r\n");
-            while (!RespWriteUtils.WriteResponse(errorMessage, ref dcurr, dend))
-                SendAndReset();
-        }
-
         /// <summary>
         /// Reads the n tokens from the current buffer, and returns
         /// total tokens read
@@ -41,5 +30,28 @@ private int ReadLeftToken(int count, ref byte* ptr)
 
             return totalTokens;
         }
+
+        /// <summary>
+        /// Aborts the execution of the current object store command and outputs
+        /// an error message to indicate a wrong number of arguments for the given command.
+        /// </summary>
+        /// <param name="cmdName">Name of the command that caused the error message.</param>
+        /// <param name="count">Number of remaining tokens belonging to this command on the receive buffer.</param>
+        /// <returns>true if the command was completely consumed, false if the input on the receive buffer was incomplete.</returns>
+        private bool AbortWithWrongNumberOfArguments(string cmdName, int count)
+        {
+            // Abort command and discard any remaining tokens on the input buffer
+            var bufSpan = new ReadOnlySpan<byte>(recvBufferPtr, bytesRead);
+
+            if (!DrainCommands(bufSpan, count))
+                return false;
+
+            // Print error message to result stream
+            var errorMessage = Encoding.ASCII.GetBytes($"-ERR wrong number of arguments for {cmdName} command.\r\n");
+            while (!RespWriteUtils.WriteResponse(errorMessage, ref dcurr, dend))
+                SendAndReset();
+
+            return true;
+        }
     }
 }

libs/server/Resp/Objects/SetCommands.cs

@@ -177,11 +177,11 @@ private unsafe bool SetLength<TGarnetApi>(int count, byte* ptr, ref TGarnetApi s
         {
             ptr += 11;
 
-            if (count - 2 < 0)
+            if (count != 2)
             {
                 setItemsDoneCount = setOpsCount = 0;
-                // Send error to output
-                WriteErrorTokenNumberInCommand("SCARD");
+                return AbortWithWrongNumberOfArguments("SCARD", count);
+
             }
             else
             {

libs/server/Resp/Objects/SortedSetCommands.cs

@@ -103,11 +103,10 @@ private unsafe bool SortedSetRemove<TGarnetApi>(int count, byte* ptr, ref TGarne
         {
             ptr += 10;
 
-            if (count - 2 <= 0)
+            if (count < 3)
             {
                 zaddDoneCount = zaddAddCount = 0;
-                // send error to output and forward pointers to the end
-                WriteErrorTokenNumberInCommand("ZREM");
+                return AbortWithWrongNumberOfArguments("ZREM", count);
             }
             else
             {
@@ -191,11 +190,10 @@ private unsafe bool SortedSetLength<TGarnetApi>(int count, byte* ptr, ref TGarne
         {
             ptr += 11;
 
-            if (count - 2 < 0)
+            if (count != 2)
             {
                 zaddDoneCount = zaddAddCount = 0;
-                // send error to output
-                WriteErrorTokenNumberInCommand("ZCARD");
+                return AbortWithWrongNumberOfArguments("ZCARD", count);
             }
             else
             {
@@ -386,11 +384,9 @@ private unsafe bool SortedSetScore<TGarnetApi>(int count, byte* ptr, ref TGarnet
             ptr += 12;
 
             //validation if minimum args
-            if (count - 3 != 0)
+            if (count != 3)
             {
-                // send error to output
-                WriteErrorTokenNumberInCommand("ZSCORE");
-                ReadLeftToken(count - 1, ref ptr);
+                return AbortWithWrongNumberOfArguments("ZSCORE", count);
             }
             else
             {
@@ -467,11 +463,9 @@ private unsafe bool SortedSetPop<TGarnetApi>(int count, byte* ptr, SortedSetOper
              where TGarnetApi : IGarnetApi
         {
             ptr += 13;
-            if (count - 1 > 2 || count - 1 == 0)
+            if (count < 2 || count > 3)
             {
-                // send error to output
-                WriteErrorTokenNumberInCommand(op == SortedSetOperation.ZPOPMAX ? "ZPOPMAX" : "ZPOPMIN");
-                ReadLeftToken(count - 1, ref ptr);
+                return AbortWithWrongNumberOfArguments(op == SortedSetOperation.ZPOPMAX ? "ZPOPMAX" : "ZPOPMIN", count);
             }
             else
             {
@@ -559,12 +553,7 @@ private unsafe bool SortedSetCount<TGarnetApi>(int count, byte* ptr, ref TGarnet
 
             if (count != 4)
             {
-                zaddDoneCount = zaddAddCount = 0;
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("ZCOUNT");
+                return AbortWithWrongNumberOfArguments("ZCOUNT", count);
             }
             else
             {
@@ -655,11 +644,7 @@ private unsafe bool SortedSetLengthByValue<TGarnetApi>(int count, byte* ptr, Sor
             if (count != 4)
             {
                 zaddDoneCount = zaddAddCount = 0;
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand(op == SortedSetOperation.ZLEXCOUNT ? "ZLEXCOUNT" : "ZREMRANGEBYLEX");
+                return AbortWithWrongNumberOfArguments(op == SortedSetOperation.ZLEXCOUNT ? "ZLEXCOUNT" : "ZREMRANGEBYLEX", count);
             }
             else
             {
@@ -749,11 +734,7 @@ private unsafe bool SortedSetIncrement<TGarnetApi>(int count, byte* ptr, ref TGa
             //validation of required args
             if (count != 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("ZINCRBY");
+                return AbortWithWrongNumberOfArguments("ZINCRBY", count);
             }
             else
             {
@@ -850,11 +831,7 @@ private unsafe bool SortedSetRank<TGarnetApi>(int count, byte* ptr, SortedSetOpe
             //validation of required args
             if (count < 3 || count > 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand(op == SortedSetOperation.ZRANK ? "ZRANK" : "ZREVRANK");
+                return AbortWithWrongNumberOfArguments(op == SortedSetOperation.ZRANK ? "ZRANK" : "ZREVRANK", count);
             }
             else
             {
@@ -935,11 +912,7 @@ private unsafe bool SortedSetRemoveRange<TGarnetApi>(int count, byte* ptr, Sorte
 
             if (count != 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand(op == SortedSetOperation.ZREMRANGEBYRANK ? "ZREMRANGEBYRANK" : "ZREMRANGEBYSCORE");
+                return AbortWithWrongNumberOfArguments(op == SortedSetOperation.ZREMRANGEBYRANK ? "ZREMRANGEBYRANK" : "ZREMRANGEBYSCORE", count);
             }
             else
             {
@@ -1021,13 +994,9 @@ private unsafe bool SortedSetRandomMember<TGarnetApi>(int count, byte* ptr, ref
         {
             ptr += 18;
 
-            if (count - 2 < 0 || count > 4)
+            if (count < 2 || count > 4)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("ZRANDMEMBER");
+                return AbortWithWrongNumberOfArguments("ZRANDMEMBER", count);
             }
             else
             {
@@ -1126,13 +1095,9 @@ private unsafe bool SortedSetDifference<TGarnetApi>(int count, byte* ptr, ref TG
              where TGarnetApi : IGarnetApi
         {
             ptr += 11;
-            if (count - 3 < 0)
+            if (count < 3)
             {
-                var tokens = ReadLeftToken(count - 1, ref ptr);
-                if (tokens < count - 1)
-                    return false;
-                // send error to output
-                WriteErrorTokenNumberInCommand("ZDIFF");
+                return AbortWithWrongNumberOfArguments("ZDIFF", count);
             }
             else
             {