Fix pub/sub: restrict commands in RESP2 subscription mode and fix reentrant lock crash

Fixes #1615. Two related bugs:

1. RESP2 subscription mode allowed arbitrary commands (GET, SET, PUBLISH, etc.)
   instead of restricting to only (P|S)SUBSCRIBE/(P|S)UNSUBSCRIBE/PING/QUIT per
   the Redis protocol. Added IsAllowedInSubscriptionMode() check in ProcessMessages
   that rejects disallowed commands with a Redis-compatible error message.

2. PUBLISH from a subscriber session caused SynchronizationLockException because
   the Publish() callback re-entered the spinlock already held by TryConsumeMessages.
   Fixed by tracking the command-processing thread ID and detecting reentrant calls
   in Publish()/PatternPublish() to skip lock acquire/release when on the same thread.

RESP3 sessions are not restricted since push message types are distinguishable
from regular responses.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: badrishc

libs/server/Resp/CmdStrings.cs

@@ -328,6 +328,7 @@ static partial class CmdStrings
         public const string GenericUnknownClientType = "ERR Unknown client type '{0}'";
         public const string GenericErrDuplicateFilter = "ERR Filter '{0}' defined multiple times";
         public const string GenericPubSubCommandDisabled = "ERR {0} is disabled, enable it with --pubsub option.";
+        public const string GenericPubSubCommandNotAllowed = "ERR Can't execute '{0}': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT are allowed in this context";
         public const string GenericErrLonLat = "ERR invalid longitude,latitude pair {0:F6},{1:F6}";
         public const string GenericErrStoreCommand = "ERR STORE option in {0} is not compatible with WITHDIST, WITHHASH and WITHCOORD options";
         public const string GenericErrCommandDisallowedWithOption =

libs/server/Resp/Parser/RespCommand.cs

@@ -627,6 +627,23 @@ public static bool IsClusterSubCommand(this RespCommand cmd)
             bool inRange = test <= (RespCommand.CLUSTER_SYNC - RespCommand.CLUSTER_ADDSLOTS);
             return inRange;
         }
+
+        /// <summary>
+        /// Returns true if <paramref name="cmd"/> is allowed while a session is in
+        /// pub/sub subscription mode (RESP2). Per the Redis protocol, only
+        /// (P|S)SUBSCRIBE, (P|S)UNSUBSCRIBE, PING, and QUIT are valid in this state.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static bool IsAllowedInSubscriptionMode(this RespCommand cmd)
+        {
+            return cmd is RespCommand.SUBSCRIBE
+                or RespCommand.UNSUBSCRIBE
+                or RespCommand.PSUBSCRIBE
+                or RespCommand.PUNSUBSCRIBE
+                or RespCommand.SSUBSCRIBE
+                or RespCommand.PING
+                or RespCommand.QUIT;
+        }
     }
 
     /// <summary>

libs/server/Resp/PubSubCommands.cs

@@ -3,7 +3,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using Garnet.common;
 using Tsavorite.core;
 
@@ -21,9 +20,15 @@ internal sealed unsafe partial class RespServerSession : ServerSessionBase
         /// <inheritdoc />
         public override unsafe void Publish(PinnedSpanByte key, PinnedSpanByte value)
         {
+            // When a session publishes to a channel it is itself subscribed to, the Broadcast
+            // callback is invoked on the same thread that already holds the network sender lock
+            // via TryConsumeMessages. Detect this reentrant case and write directly to the
+            // existing output buffer instead of re-entering the lock.
+            var reentrant = commandProcessingThreadId == Environment.CurrentManagedThreadId;
             try
             {
-                networkSender.EnterAndGetResponseObject(out dcurr, out dend);
+                if (!reentrant)
+                    networkSender.EnterAndGetResponseObject(out dcurr, out dend);
 
                 WritePushLength(3);
 
@@ -35,7 +40,7 @@ public override unsafe void Publish(PinnedSpanByte key, PinnedSpanByte value)
                 WriteDirectLargeRespString(value.ReadOnlySpan);
 
                 // Flush the publish message for this subscriber
-                if (dcurr > networkSender.GetResponseObjectHead())
+                if (!reentrant && dcurr > networkSender.GetResponseObjectHead())
                     Send(networkSender.GetResponseObjectHead());
             }
             catch
@@ -44,16 +49,19 @@ public override unsafe void Publish(PinnedSpanByte key, PinnedSpanByte value)
             }
             finally
             {
-                networkSender.ExitAndReturnResponseObject();
+                if (!reentrant)
+                    networkSender.ExitAndReturnResponseObject();
             }
         }
 
         /// <inheritdoc />
         public override unsafe void PatternPublish(PinnedSpanByte pattern, PinnedSpanByte key, PinnedSpanByte value)
         {
+            var reentrant = commandProcessingThreadId == Environment.CurrentManagedThreadId;
             try
             {
-                networkSender.EnterAndGetResponseObject(out dcurr, out dend);
+                if (!reentrant)
+                    networkSender.EnterAndGetResponseObject(out dcurr, out dend);
 
                 WritePushLength(4);
 
@@ -65,7 +73,7 @@ public override unsafe void PatternPublish(PinnedSpanByte pattern, PinnedSpanByt
                 WriteDirectLargeRespString(key.ReadOnlySpan);
                 WriteDirectLargeRespString(value.ReadOnlySpan);
 
-                if (dcurr > networkSender.GetResponseObjectHead())
+                if (!reentrant && dcurr > networkSender.GetResponseObjectHead())
                     Send(networkSender.GetResponseObjectHead());
             }
             catch
@@ -74,7 +82,8 @@ public override unsafe void PatternPublish(PinnedSpanByte pattern, PinnedSpanByt
             }
             finally
             {
-                networkSender.ExitAndReturnResponseObject();
+                if (!reentrant)
+                    networkSender.ExitAndReturnResponseObject();
             }
         }
 
@@ -100,7 +109,6 @@ private bool NetworkPUBLISH(RespCommand cmd)
                 return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_CLUSTER_DISABLED);
             }
 
-            Debug.Assert(isSubscriptionSession == false);
             // PUBLISH channel message => [*3\r\n$7\r\nPUBLISH\r\n$]7\r\nchannel\r\n$7\r\message\r\n
 
             var key = parseState.GetArgSliceByRef(0);

libs/server/Resp/RespServerSession.cs

@@ -80,6 +80,10 @@ internal sealed unsafe partial class RespServerSession : ServerSessionBase
 
         int opCount;
 
+        // Thread ID of the thread currently processing commands (used by Publish/PatternPublish
+        // callbacks to detect reentrant calls when the same session publishes to a self-subscribed channel)
+        int commandProcessingThreadId;
+
         /// <summary>
         /// Current database session items
         /// </summary>
@@ -441,7 +445,9 @@ public override int TryConsumeMessages(byte* reqBuffer, int bytesReceived)
                 clusterSession?.AcquireCurrentEpoch();
                 recvBufferPtr = reqBuffer;
                 networkSender.EnterAndGetResponseObject(out dcurr, out dend);
+                commandProcessingThreadId = Environment.CurrentManagedThreadId;
                 ProcessMessages();
+                commandProcessingThreadId = 0;
                 recvBufferPtr = null;
             }
             catch (RespParsingException ex)
@@ -575,7 +581,14 @@ private void ProcessMessages()
 
                     if (CheckACLPermissions(cmd) && (noScriptPassed = CheckScriptPermissions(cmd)))
                     {
-                        if (txnManager.state != TxnState.None)
+                        // In RESP2, only a small set of commands are allowed while in subscription mode.
+                        // RESP3 uses distinct push types for subscription messages, so all commands are valid.
+                        if (isSubscriptionSession && respProtocolVersion == 2 && !cmd.IsAllowedInSubscriptionMode())
+                        {
+                            while (!RespWriteUtils.TryWriteError(string.Format(CmdStrings.GenericPubSubCommandNotAllowed, cmd.ToString()), ref dcurr, dend))
+                                SendAndReset();
+                        }
+                        else if (txnManager.state != TxnState.None)
                         {
                             if (txnManager.state == TxnState.Running)
                             {

test/Garnet.test/RespPubSubTests.cs

@@ -6,6 +6,7 @@
 using System.Security.Cryptography;
 using System.Threading;
 using Allure.NUnit;
+using Garnet.common;
 using NUnit.Framework;
 using NUnit.Framework.Legacy;
 using StackExchange.Redis;
@@ -230,5 +231,140 @@ private void SubscribeAndPublish(ISubscriber sub, IDatabase db, RedisChannel cha
                 ClassicAssert.IsTrue(repeat != 0, "Timeout waiting for subscription receive");
             }
         }
+
+        /// <summary>
+        /// Verifies that disallowed commands (GET, SET, PUBLISH, MULTI) are rejected
+        /// with an appropriate error when a RESP2 session is in subscription mode.
+        /// </summary>
+        [Test]
+        public void PubSubModeRejectsDisallowedCommandsInResp2()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest(countResponseType: CountResponseType.Bytes);
+
+            // Subscribe to enter subscription mode
+            var subscribeResp = "*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:1\r\n";
+            var response = lightClientRequest.Execute("SUBSCRIBE foo", subscribeResp.Length);
+            ClassicAssert.AreEqual(subscribeResp, response);
+
+            // GET should be rejected
+            var errorResp = "-ERR Can't execute 'GET': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT are allowed in this context\r\n";
+            response = lightClientRequest.Execute("GET bar", errorResp.Length);
+            ClassicAssert.AreEqual(errorResp, response);
+
+            // SET should be rejected
+            errorResp = "-ERR Can't execute 'SET': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT are allowed in this context\r\n";
+            response = lightClientRequest.Execute("SET bar value", errorResp.Length);
+            ClassicAssert.AreEqual(errorResp, response);
+
+            // PUBLISH should be rejected
+            errorResp = "-ERR Can't execute 'PUBLISH': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT are allowed in this context\r\n";
+            response = lightClientRequest.Execute("PUBLISH foo bar", errorResp.Length);
+            ClassicAssert.AreEqual(errorResp, response);
+
+            // MULTI should be rejected
+            errorResp = "-ERR Can't execute 'MULTI': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT are allowed in this context\r\n";
+            response = lightClientRequest.Execute("MULTI", errorResp.Length);
+            ClassicAssert.AreEqual(errorResp, response);
+        }
+
+        /// <summary>
+        /// Verifies that allowed commands (PING, SUBSCRIBE, PSUBSCRIBE, UNSUBSCRIBE,
+        /// PUNSUBSCRIBE, QUIT) work correctly in RESP2 subscription mode.
+        /// </summary>
+        [Test]
+        public void PubSubModeAllowsValidCommandsInResp2()
+        {
+            using var lightClientRequest = TestUtils.CreateRequest(countResponseType: CountResponseType.Bytes);
+
+            // Enter subscription mode
+            var subscribeResp = "*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:1\r\n";
+            var response = lightClientRequest.Execute("SUBSCRIBE foo", subscribeResp.Length);
+            ClassicAssert.AreEqual(subscribeResp, response);
+
+            // PING should return subscription-mode PONG
+            var pongResp = "*2\r\n$4\r\npong\r\n$0\r\n\r\n";
+            response = lightClientRequest.Execute("PING", pongResp.Length);
+            ClassicAssert.AreEqual(pongResp, response);
+
+            // Another SUBSCRIBE should work (channel count goes to 2)
+            subscribeResp = "*3\r\n$9\r\nsubscribe\r\n$3\r\nbar\r\n:2\r\n";
+            response = lightClientRequest.Execute("SUBSCRIBE bar", subscribeResp.Length);
+            ClassicAssert.AreEqual(subscribeResp, response);
+
+            // PSUBSCRIBE should work (channel count goes to 3)
+            var psubResp = "*3\r\n$10\r\npsubscribe\r\n$4\r\nbaz*\r\n:3\r\n";
+            response = lightClientRequest.Execute("PSUBSCRIBE baz*", psubResp.Length);
+            ClassicAssert.AreEqual(psubResp, response);
+
+            // UNSUBSCRIBE should work (channel count goes to 2)
+            var unsubResp = "*3\r\n$11\r\nunsubscribe\r\n$3\r\nbar\r\n:2\r\n";
+            response = lightClientRequest.Execute("UNSUBSCRIBE bar", unsubResp.Length);
+            ClassicAssert.AreEqual(unsubResp, response);
+
+            // PUNSUBSCRIBE should work (channel count goes to 1)
+            var punsubResp = "*3\r\n$12\r\npunsubscribe\r\n$4\r\nbaz*\r\n:1\r\n";
+            response = lightClientRequest.Execute("PUNSUBSCRIBE baz*", punsubResp.Length);
+            ClassicAssert.AreEqual(punsubResp, response);
+
+            // Still in subscription mode - GET should be rejected
+            var errorResp = "-ERR Can't execute 'GET': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT are allowed in this context\r\n";
+            response = lightClientRequest.Execute("GET bar", errorResp.Length);
+            ClassicAssert.AreEqual(errorResp, response);
+
+            // UNSUBSCRIBE last channel to exit subscription mode (channel count goes to 0)
+            unsubResp = "*3\r\n$11\r\nunsubscribe\r\n$3\r\nfoo\r\n:0\r\n";
+            response = lightClientRequest.Execute("UNSUBSCRIBE foo", unsubResp.Length);
+            ClassicAssert.AreEqual(unsubResp, response);
+
+            // No longer in subscription mode - GET should work (key doesn't exist = null)
+            var getResp = "$-1\r\n";
+            response = lightClientRequest.Execute("GET bar", getResp.Length);
+            ClassicAssert.AreEqual(getResp, response);
+        }
+
+        /// <summary>
+        /// Verifies that a RESP3 session in subscription mode can execute PUBLISH
+        /// (including self-publish to its own subscribed channel) without a
+        /// SynchronizationLockException crash. This was the core lock bug in issue #1615.
+        /// </summary>
+        [Test]
+        public void PubSubSelfPublishResp3NoLockError()
+        {
+            // Use Newlines counting for HELLO 3 (variable-length map response)
+            using var lightClientRequest = TestUtils.CreateRequest(countResponseType: CountResponseType.Newlines);
+
+            // Switch to RESP3
+            var response = lightClientRequest.Execute("HELLO 3", 30);
+            ClassicAssert.IsTrue(response.Contains("proto"));
+
+            // Switch to Bytes counting for precise response control
+            lightClientRequest.countResponseType = CountResponseType.Bytes;
+
+            // Subscribe to a channel
+            var subscribeResp = "*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:1\r\n";
+            response = lightClientRequest.Execute("SUBSCRIBE foo", subscribeResp.Length);
+            ClassicAssert.AreEqual(subscribeResp, response);
+
+            // Self-publish: this triggers the reentrant Publish() callback on the same session.
+            // Before the fix, this would crash with SynchronizationLockException.
+            //
+            // Expected response consists of:
+            //   1. Push message (self-notification): >3\r\n$7\r\nmessage\r\n$3\r\nfoo\r\n$3\r\nbar\r\n
+            //   2. PUBLISH response: :1\r\n
+            var pushMsg = ">3\r\n$7\r\nmessage\r\n$3\r\nfoo\r\n$3\r\nbar\r\n";
+            var publishResp = ":1\r\n";
+            var expectedPublishTotal = pushMsg + publishResp;
+            response = lightClientRequest.Execute("PUBLISH foo bar", expectedPublishTotal.Length);
+            ClassicAssert.AreEqual(expectedPublishTotal, response);
+
+            // Unsubscribe and verify server is still responsive
+            var unsubResp = "*3\r\n$11\r\nunsubscribe\r\n$3\r\nfoo\r\n:0\r\n";
+            response = lightClientRequest.Execute("UNSUBSCRIBE foo", unsubResp.Length);
+            ClassicAssert.AreEqual(unsubResp, response);
+
+            // PING to confirm server health
+            response = lightClientRequest.Execute("PING", "+PONG\r\n".Length);
+            ClassicAssert.AreEqual("+PONG\r\n", response);
+        }
     }
 }