Skip to content

Commit 7cc05c7

Browse files
prvykprvyk
committed
Check type in blocking functions.
Add test.
1 parent 3efcc05 commit 7cc05c7

File tree

6 files changed

+152
-40
lines changed

6 files changed

+152
-40
lines changed

‎libs/server/Resp/CmdStrings.cs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,6 +303,13 @@ static partial class CmdStrings
303303
public static ReadOnlySpan<byte> STRING => "STRING"u8;
304304
public static ReadOnlySpan<byte> stringt => "string"u8;
305305

306+
public const string listTyprStr = "list";
307+
public const string setTypeStr = "set";
308+
public const string zsetTypeStr = "zset";
309+
public const string hashTypeStr = "hash";
310+
public const string stringTypeStr = "string";
311+
public const string noneTypeStr = "none";
312+
306313
/// <summary>
307314
/// Register object types
308315
/// </summary>

‎libs/server/Resp/Objects/ListCommands.cs

Lines changed: 67 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,8 @@ private unsafe bool ListPopMultiple<TGarnetApi>(ref TGarnetApi storageApi)
278278
return true;
279279
}
280280

281-
private bool ListBlockingPop(RespCommand command)
281+
private bool ListBlockingPop<TGarnetApi>(RespCommand command, ref TGarnetApi storageApi)
282+
where TGarnetApi : IGarnetApi
282283
{
283284
if (parseState.Count < 2)
284285
{
@@ -289,14 +290,24 @@ private bool ListBlockingPop(RespCommand command)
289290

290291
for (var i = 0; i < keysBytes.Length; i++)
291292
{
292-
keysBytes[i] = parseState.GetArgSliceByRef(i).SpanByte.ToByteArray();
293+
var key = parseState.GetArgSliceByRef(i);
294+
var status = storageApi.GetKeyType(key, out var typeName);
295+
if (status == GarnetStatus.OK && typeName != CmdStrings.listTyprStr)
296+
{
297+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_TYPE);
298+
}
299+
300+
keysBytes[i] = key.SpanByte.ToByteArray();
293301
}
294302

295303
if (!parseState.TryGetDouble(parseState.Count - 1, out var timeout))
296304
{
297-
while (!RespWriteUtils.TryWriteError(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT, ref dcurr, dend))
298-
SendAndReset();
299-
return true;
305+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT);
306+
}
307+
308+
if (timeout < 0)
309+
{
310+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_IS_NEGATIVE);
300311
}
301312

302313
if (storeWrapper.itemBroker == null)
@@ -331,7 +342,8 @@ private bool ListBlockingPop(RespCommand command)
331342
return true;
332343
}
333344

334-
private unsafe bool ListBlockingMove()
345+
private unsafe bool ListBlockingMove<TGarnetApi>(ref TGarnetApi storageApi)
346+
where TGarnetApi : IGarnetApi
335347
{
336348
if (parseState.Count != 5)
337349
{
@@ -345,19 +357,23 @@ private unsafe bool ListBlockingMove()
345357

346358
if (!parseState.TryGetDouble(4, out var timeout))
347359
{
348-
while (!RespWriteUtils.TryWriteError(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT, ref dcurr, dend))
349-
SendAndReset();
350-
return true;
360+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT);
361+
}
362+
363+
if (timeout < 0)
364+
{
365+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_IS_NEGATIVE);
351366
}
352367

353-
return ListBlockingMove(srcKey, dstKey, srcDir, dstDir, timeout);
368+
return ListBlockingMove(ref storageApi, srcKey, dstKey, srcDir, dstDir, timeout);
354369
}
355370

356371
/// <summary>
357372
/// BRPOPLPUSH
358373
/// </summary>
359374
/// <returns></returns>
360-
private bool ListBlockingPopPush()
375+
private bool ListBlockingPopPush<TGarnetApi>(ref TGarnetApi storageApi)
376+
where TGarnetApi : IGarnetApi
361377
{
362378
if (parseState.Count != 3)
363379
{
@@ -371,15 +387,20 @@ private bool ListBlockingPopPush()
371387

372388
if (!parseState.TryGetDouble(2, out var timeout))
373389
{
374-
while (!RespWriteUtils.TryWriteError(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT, ref dcurr, dend))
375-
SendAndReset();
376-
return true;
390+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT);
391+
}
392+
393+
if (timeout < 0)
394+
{
395+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_IS_NEGATIVE);
377396
}
378397

379-
return ListBlockingMove(srcKey, dstKey, rightOption, leftOption, timeout);
398+
return ListBlockingMove(ref storageApi, srcKey, dstKey, rightOption, leftOption, timeout);
380399
}
381400

382-
private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey, ArgSlice srcDir, ArgSlice dstDir, double timeout)
401+
private bool ListBlockingMove<TGarnetApi>(ref TGarnetApi storageApi, ArgSlice srcKey, ArgSlice dstKey,
402+
ArgSlice srcDir, ArgSlice dstDir, double timeout)
403+
where TGarnetApi : IGarnetApi
383404
{
384405
var cmdArgs = new ArgSlice[] { default, default, default };
385406

@@ -394,14 +415,26 @@ private bool ListBlockingMove(ArgSlice srcKey, ArgSlice dstKey, ArgSlice srcDir,
394415
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_SYNTAX_ERROR);
395416
}
396417

418+
if (storeWrapper.itemBroker == null)
419+
throw new GarnetException("Object store is disabled");
420+
421+
if (storageApi.GetKeyType(srcKey, out var keyType) == GarnetStatus.OK &&
422+
keyType != CmdStrings.listTyprStr)
423+
{
424+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_TYPE);
425+
}
426+
427+
if (storageApi.GetKeyType(dstKey, out keyType) == GarnetStatus.OK &&
428+
keyType != CmdStrings.listTyprStr)
429+
{
430+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_TYPE);
431+
}
432+
397433
var pSrcDir = (byte*)&sourceDirection;
398434
var pDstDir = (byte*)&destinationDirection;
399435
cmdArgs[1] = new ArgSlice(pSrcDir, 1);
400436
cmdArgs[2] = new ArgSlice(pDstDir, 1);
401437

402-
if (storeWrapper.itemBroker == null)
403-
throw new GarnetException("Object store is disabled");
404-
405438
var result = storeWrapper.itemBroker.MoveCollectionItemAsync(RespCommand.BLMOVE, srcKey.ToArray(), this, timeout,
406439
cmdArgs).Result;
407440

@@ -910,7 +943,8 @@ public bool ListSet<TGarnetApi>(ref TGarnetApi storageApi)
910943
/// BLMPOP timeout numkeys key [key ...] LEFT|RIGHT [COUNT count]
911944
/// </summary>
912945
/// <returns></returns>
913-
private unsafe bool ListBlockingPopMultiple()
946+
private unsafe bool ListBlockingPopMultiple<TGarnetApi>(ref TGarnetApi storageApi)
947+
where TGarnetApi : IGarnetApi
914948
{
915949
if (parseState.Count < 4)
916950
{
@@ -925,6 +959,11 @@ private unsafe bool ListBlockingPopMultiple()
925959
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT);
926960
}
927961

962+
if (timeout < 0)
963+
{
964+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_IS_NEGATIVE);
965+
}
966+
928967
// Read count of keys
929968
if (!parseState.TryGetInt(currTokenId++, out var numKeys))
930969
{
@@ -941,7 +980,14 @@ private unsafe bool ListBlockingPopMultiple()
941980
var keysBytes = new byte[numKeys][];
942981
for (var i = 0; i < keysBytes.Length; i++)
943982
{
944-
keysBytes[i] = parseState.GetArgSliceByRef(currTokenId++).SpanByte.ToByteArray();
983+
var key = parseState.GetArgSliceByRef(currTokenId++);
984+
var status = storageApi.GetKeyType(key, out var typeName);
985+
if (status == GarnetStatus.OK && typeName != CmdStrings.listTyprStr)
986+
{
987+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_TYPE);
988+
}
989+
990+
keysBytes[i] = key.SpanByte.ToByteArray();
945991
}
946992

947993
var cmdArgs = new ArgSlice[2];

‎libs/server/Resp/Objects/SortedSetCommands.cs

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1541,7 +1541,8 @@ private unsafe bool SortedSetUnionStore<TGarnetApi>(ref TGarnetApi storageApi)
15411541
/// <summary>
15421542
/// BZPOPMIN/BZPOPMAX key [key ...] timeout
15431543
/// </summary>
1544-
private unsafe bool SortedSetBlockingPop(RespCommand command)
1544+
private unsafe bool SortedSetBlockingPop<TGarnetApi>(RespCommand command, ref TGarnetApi storageApi)
1545+
where TGarnetApi : IGarnetApi
15451546
{
15461547
if (storeWrapper.itemBroker == null)
15471548
throw new GarnetException("Object store is disabled");
@@ -1551,16 +1552,28 @@ private unsafe bool SortedSetBlockingPop(RespCommand command)
15511552
return AbortWithWrongNumberOfArguments(command.ToString());
15521553
}
15531554

1554-
if (!parseState.TryGetDouble(parseState.Count - 1, out var timeout) || (timeout < 0))
1555+
if (!parseState.TryGetDouble(parseState.Count - 1, out var timeout))
15551556
{
15561557
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_NOT_VALID_FLOAT);
15571558
}
15581559

1560+
if (timeout < 0)
1561+
{
1562+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_TIMEOUT_IS_NEGATIVE);
1563+
}
1564+
15591565
var keysBytes = new byte[parseState.Count - 1][];
15601566

15611567
for (var i = 0; i < keysBytes.Length; i++)
15621568
{
1563-
keysBytes[i] = parseState.GetArgSliceByRef(i).SpanByte.ToByteArray();
1569+
var key = parseState.GetArgSliceByRef(i);
1570+
var status = storageApi.GetKeyType(key, out var typeName);
1571+
if (status == GarnetStatus.OK && typeName != CmdStrings.zsetTypeStr)
1572+
{
1573+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_TYPE);
1574+
}
1575+
1576+
keysBytes[i] = key.SpanByte.ToByteArray();
15641577
}
15651578

15661579
var result = storeWrapper.itemBroker.GetCollectionItemAsync(command, keysBytes, this, timeout).Result;
@@ -1590,7 +1603,8 @@ private unsafe bool SortedSetBlockingPop(RespCommand command)
15901603
/// <summary>
15911604
/// BZMPOP timeout numkeys key [key ...] &lt;MIN | MAX&gt; [COUNT count]
15921605
/// </summary>
1593-
private unsafe bool SortedSetBlockingMPop()
1606+
private unsafe bool SortedSetBlockingMPop<TGarnetApi>(ref TGarnetApi storageApi)
1607+
where TGarnetApi : IGarnetApi
15941608
{
15951609
if (storeWrapper.itemBroker == null)
15961610
throw new GarnetException("Object store is disabled");
@@ -1629,7 +1643,14 @@ private unsafe bool SortedSetBlockingMPop()
16291643
var keysBytes = new byte[numKeys][];
16301644
for (var i = 0; i < keysBytes.Length; i++)
16311645
{
1632-
keysBytes[i] = parseState.GetArgSliceByRef(currTokenId++).SpanByte.ToByteArray();
1646+
var key = parseState.GetArgSliceByRef(currTokenId++);
1647+
var status = storageApi.GetKeyType(key, out var typeName);
1648+
if (status == GarnetStatus.OK && typeName != CmdStrings.zsetTypeStr)
1649+
{
1650+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_WRONG_TYPE);
1651+
}
1652+
1653+
keysBytes[i] = key.SpanByte.ToByteArray();
16331654
}
16341655

16351656
var cmdArgs = new ArgSlice[2];

‎libs/server/Resp/RespServerSession.cs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -709,9 +709,9 @@ private bool ProcessArrayCommands<TGarnetApi>(RespCommand cmd, ref TGarnetApi st
709709
RespCommand.ZRANDMEMBER => SortedSetRandomMember(ref storageApi),
710710
RespCommand.ZDIFF => SortedSetDifference(ref storageApi),
711711
RespCommand.ZDIFFSTORE => SortedSetDifferenceStore(ref storageApi),
712-
RespCommand.BZMPOP => SortedSetBlockingMPop(),
713-
RespCommand.BZPOPMAX => SortedSetBlockingPop(cmd),
714-
RespCommand.BZPOPMIN => SortedSetBlockingPop(cmd),
712+
RespCommand.BZMPOP => SortedSetBlockingMPop(ref storageApi),
713+
RespCommand.BZPOPMAX => SortedSetBlockingPop(cmd, ref storageApi),
714+
RespCommand.BZPOPMIN => SortedSetBlockingPop(cmd, ref storageApi),
715715
RespCommand.ZREVRANGE => SortedSetRange(cmd, ref storageApi),
716716
RespCommand.ZREVRANGEBYLEX => SortedSetRange(cmd, ref storageApi),
717717
RespCommand.ZREVRANGEBYSCORE => SortedSetRange(cmd, ref storageApi),
@@ -757,11 +757,11 @@ private bool ProcessArrayCommands<TGarnetApi>(RespCommand cmd, ref TGarnetApi st
757757
RespCommand.LMOVE => ListMove(ref storageApi),
758758
RespCommand.LMPOP => ListPopMultiple(ref storageApi),
759759
RespCommand.LSET => ListSet(ref storageApi),
760-
RespCommand.BLPOP => ListBlockingPop(cmd),
761-
RespCommand.BRPOP => ListBlockingPop(cmd),
762-
RespCommand.BLMOVE => ListBlockingMove(),
763-
RespCommand.BRPOPLPUSH => ListBlockingPopPush(),
764-
RespCommand.BLMPOP => ListBlockingPopMultiple(),
760+
RespCommand.BLPOP => ListBlockingPop(cmd, ref storageApi),
761+
RespCommand.BRPOP => ListBlockingPop(cmd, ref storageApi),
762+
RespCommand.BLMOVE => ListBlockingMove(ref storageApi),
763+
RespCommand.BRPOPLPUSH => ListBlockingPopPush(ref storageApi),
764+
RespCommand.BLMPOP => ListBlockingPopMultiple(ref storageApi),
765765
// Hash Commands
766766
RespCommand.HSET => HashSet(cmd, ref storageApi),
767767
RespCommand.HMSET => HashSet(cmd, ref storageApi),

‎libs/server/Storage/Session/MainStore/MainStoreOps.cs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1196,7 +1196,7 @@ public GarnetStatus GetKeyType<TContext, TObjectContext>(ArgSlice key, out strin
11961196
where TContext : ITsavoriteContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator>
11971197
where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions, ObjectStoreFunctions, ObjectStoreAllocator>
11981198
{
1199-
keyType = "string";
1199+
keyType = CmdStrings.stringTypeStr;
12001200
// Check if key exists in Main store
12011201
var status = EXISTS(key, StoreType.Main, ref context, ref objectContext);
12021202

@@ -1208,24 +1208,24 @@ public GarnetStatus GetKeyType<TContext, TObjectContext>(ArgSlice key, out strin
12081208
{
12091209
if ((output.GarnetObject as SortedSetObject) != null)
12101210
{
1211-
keyType = "zset";
1211+
keyType = CmdStrings.zsetTypeStr;
12121212
}
12131213
else if ((output.GarnetObject as ListObject) != null)
12141214
{
1215-
keyType = "list";
1215+
keyType = CmdStrings.listTyprStr;
12161216
}
12171217
else if ((output.GarnetObject as SetObject) != null)
12181218
{
1219-
keyType = "set";
1219+
keyType = CmdStrings.setTypeStr;
12201220
}
12211221
else if ((output.GarnetObject as HashObject) != null)
12221222
{
1223-
keyType = "hash";
1223+
keyType = CmdStrings.hashTypeStr;
12241224
}
12251225
}
12261226
else
12271227
{
1228-
keyType = "none";
1228+
keyType = CmdStrings.noneTypeStr;
12291229
status = GarnetStatus.NOTFOUND;
12301230
}
12311231
}

‎test/Garnet.test/RespBlockingCollectionTests.cs

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -571,5 +571,43 @@ public void BzpopMinMaxTimeoutTest(string command)
571571
var expectedResponse = "$-1\r\n";
572572
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
573573
}
574+
575+
[Test]
576+
public void PopWrongTypeTest()
577+
{
578+
using var lightClientRequest = TestUtils.CreateRequest();
579+
lightClientRequest.SendCommand("SET key 0");
580+
581+
var response = lightClientRequest.SendCommand($"BLMPOP 0 1 key RIGHT");
582+
var expectedResponse = "-WRONGTYPE Operation against a key holding the wrong kind of value.\r\n";
583+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
584+
585+
response = lightClientRequest.SendCommand($"BZMPOP 0 1 key MAX");
586+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
587+
588+
response = lightClientRequest.SendCommand($"BZPOPMIN key 0");
589+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
590+
591+
response = lightClientRequest.SendCommand($"BZPOPMAX key 0");
592+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
593+
594+
response = lightClientRequest.SendCommand($"BLMOVE foo key LEFT LEFT 0");
595+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
596+
597+
response = lightClientRequest.SendCommand($"BLMOVE key foo RIGHT RIGHT 0");
598+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
599+
600+
response = lightClientRequest.SendCommand($"BRPOPLPUSH key foo 0");
601+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
602+
603+
response = lightClientRequest.SendCommand($"BRPOPLPUSH foo key 0");
604+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
605+
606+
response = lightClientRequest.SendCommand($"BLPOP key 0");
607+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
608+
609+
response = lightClientRequest.SendCommand($"BRPOP key 0");
610+
TestUtils.AssertEqualUpToExpectedLength(expectedResponse, response);
611+
}
574612
}
575613
}

0 commit comments

Comments
 (0)