Skip to content

Commit 5897200

Browse files
vazoisvazois
authored
Harden Offset Validation for Bitmap Operations (#1610)
* bound validation for max 512MB strings * add BITPOS validation * add BITCOUNT offset validation * addressing comments * fix slot verification test * addressing comments * making check in backend bitfield call Debug.Assert * fix formatting * fix Release build
1 parent 7c10720 commit 5897200

File tree

8 files changed

+354
-61
lines changed

8 files changed

+354
-61
lines changed

‎libs/server/Resp/Bitmap/BitmapCommands.cs‎

Lines changed: 81 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ private bool NetworkStringSetBit<TGarnetApi>(ref TGarnetApi storageApi)
138138
var sbKey = parseState.GetArgSliceByRef(0).SpanByte;
139139

140140
// Validate offset
141-
if (!parseState.TryGetLong(1, out var offset) || (offset < 0))
141+
if (!parseState.TryGetLong(1, out var offset) || (offset < 0) || !BitmapManager.IsValidBitOffset(offset))
142142
{
143143
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BITOFFSET_IS_NOT_INTEGER);
144144
}
@@ -178,7 +178,7 @@ private bool NetworkStringGetBit<TGarnetApi>(ref TGarnetApi storageApi)
178178
var sbKey = parseState.GetArgSliceByRef(0).SpanByte;
179179

180180
// Validate offset
181-
if (!parseState.TryGetLong(1, out var offset) || (offset < 0))
181+
if (!parseState.TryGetLong(1, out var offset) || (offset < 0) || !BitmapManager.IsValidBitOffset(offset))
182182
{
183183
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BITOFFSET_IS_NOT_INTEGER);
184184
}
@@ -205,29 +205,43 @@ private bool NetworkStringBitCount<TGarnetApi>(ref TGarnetApi storageApi)
205205
where TGarnetApi : IGarnetApi
206206
{
207207
var count = parseState.Count;
208-
if (count < 1 || count > 4)
208+
if (count is not 1 and not 3 and not 4)
209209
{
210210
return AbortWithWrongNumberOfArguments(nameof(RespCommand.BITCOUNT));
211211
}
212212

213213
// <[Get Key]>
214214
var sbKey = parseState.GetArgSliceByRef(0).SpanByte;
215215

216-
// Validate start & end offsets, if exist
217-
if (parseState.Count > 1)
216+
// Extract parameters in command order:
217+
// start, end, [BIT|BYTE]
218+
var useBitIndex = true;
219+
if (count > 1)
218220
{
219-
if (!parseState.TryGetInt(1, out _) || (parseState.Count > 2 && !parseState.TryGetInt(2, out _)))
221+
if (!parseState.TryGetLong(1, out _) || !parseState.TryGetLong(2, out _))
220222
{
221223
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_VALUE_IS_NOT_INTEGER);
222224
}
223-
}
224225

225-
var input = new RawStringInput(RespCommand.BITCOUNT, ref parseState, startIdx: 1);
226+
if (count > 3)
227+
{
228+
var spanOffsetType = parseState.GetArgSliceByRef(3).ReadOnlySpan;
229+
if (spanOffsetType.EqualsUpperCaseSpanIgnoringCase("BIT"u8))
230+
useBitIndex = true;
231+
else if (spanOffsetType.EqualsUpperCaseSpanIgnoringCase("BYTE"u8))
232+
useBitIndex = false;
233+
else
234+
{
235+
return AbortWithErrorMessage(CmdStrings.RESP_SYNTAX_ERROR);
236+
}
226237

227-
var o = new SpanByteAndMemory(dcurr, (int)(dend - dcurr));
228238

229-
var status = storageApi.StringBitCount(ref sbKey, ref input, ref o);
239+
}
240+
}
230241

242+
var input = new RawStringInput(RespCommand.BITCOUNT, ref parseState, startIdx: 1, arg1: useBitIndex ? 1 : 0);
243+
var o = new SpanByteAndMemory(dcurr, (int)(dend - dcurr));
244+
var status = storageApi.StringBitCount(ref sbKey, ref input, ref o);
231245
if (status == GarnetStatus.OK)
232246
{
233247
if (!o.IsSpanByte)
@@ -251,7 +265,7 @@ private bool NetworkStringBitPosition<TGarnetApi>(ref TGarnetApi storageApi)
251265
where TGarnetApi : IGarnetApi
252266
{
253267
var count = parseState.Count;
254-
if (count < 2 || count > 5)
268+
if (count is < 2 or > 5)
255269
{
256270
return AbortWithWrongNumberOfArguments(nameof(RespCommand.BITPOS));
257271
}
@@ -266,31 +280,57 @@ private bool NetworkStringBitPosition<TGarnetApi>(ref TGarnetApi storageApi)
266280
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BIT_IS_NOT_INTEGER);
267281
}
268282

269-
// Validate start & end offsets, if exist
270-
if (parseState.Count > 2)
283+
long startOffset = 0;
284+
long endOffset = -1;
285+
byte offsetType = 0x0;
286+
var hasStartOffset = false;
287+
var hasEndOffset = false;
288+
289+
// Extract parameters in command order, consistent with PrivateMethods BITPOS decoding:
290+
// bit, start, end, [BIT|BYTE]
291+
if (count > 2)
271292
{
272-
if (!parseState.TryGetInt(2, out _) ||
273-
(parseState.Count > 3 && !parseState.TryGetInt(3, out _)))
293+
// start
294+
if (!parseState.TryGetLong(2, out startOffset))
274295
{
275296
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_VALUE_IS_NOT_INTEGER);
276297
}
277-
}
298+
hasStartOffset = true;
278299

279-
// Validate offset range type (BIT / BYTE), if exists
280-
if (parseState.Count > 4)
281-
{
282-
var sbOffsetType = parseState.GetArgSliceByRef(4).ReadOnlySpan;
283-
if (!sbOffsetType.EqualsUpperCaseSpanIgnoringCase("BIT"u8) &&
284-
!sbOffsetType.EqualsUpperCaseSpanIgnoringCase("BYTE"u8))
300+
if (count > 3)
285301
{
286-
return AbortWithErrorMessage(CmdStrings.RESP_SYNTAX_ERROR);
302+
// end
303+
if (!parseState.TryGetLong(3, out endOffset))
304+
{
305+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_VALUE_IS_NOT_INTEGER);
306+
}
307+
hasEndOffset = true;
308+
309+
if (count > 4)
310+
{
311+
// Optional index type (default BYTE)
312+
var sbOffsetType = parseState.GetArgSliceByRef(4).ReadOnlySpan;
313+
if (sbOffsetType.EqualsUpperCaseSpanIgnoringCase("BIT"u8))
314+
{
315+
offsetType = 0x1;
316+
}
317+
else if (!sbOffsetType.EqualsUpperCaseSpanIgnoringCase("BYTE"u8))
318+
{
319+
return AbortWithErrorMessage(CmdStrings.RESP_SYNTAX_ERROR);
320+
}
321+
}
287322
}
288323
}
289324

290-
var input = new RawStringInput(RespCommand.BITPOS, ref parseState, startIdx: 1);
325+
if (BitmapManager.TryValidateBitPosOffsets(startOffset, endOffset, offsetType, hasStartOffset, hasEndOffset))
326+
{
327+
while (!RespWriteUtils.TryWriteInt64(-1, ref dcurr, dend))
328+
SendAndReset();
329+
return true;
330+
}
291331

332+
var input = new RawStringInput(RespCommand.BITPOS, ref parseState, startIdx: 1);
292333
var o = new SpanByteAndMemory(dcurr, (int)(dend - dcurr));
293-
294334
var status = storageApi.StringBitPosition(ref sbKey, ref input, ref o);
295335

296336
if (status == GarnetStatus.OK)
@@ -385,17 +425,23 @@ private bool StringBitField<TGarnetApi>(ref TGarnetApi storageApi)
385425

386426
// [GET <encoding> <offset>] [SET <encoding> <offset> <value>] [INCRBY <encoding> <offset> <increment>]
387427
// Process encoding argument
388-
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldEncoding(currTokenIdx, out _, out _))
428+
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldEncoding(currTokenIdx, out var bitCount, out _))
389429
{
390430
return AbortWithErrorMessage(CmdStrings.RESP_ERR_INVALID_BITFIELD_TYPE);
391431
}
392432
var encodingSlice = parseState.GetArgSliceByRef(currTokenIdx++);
393433

394434
// Process offset argument
395-
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldOffset(currTokenIdx, out _, out _))
435+
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldOffset(currTokenIdx, out var offset, out var multiplyOffset))
436+
{
437+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BITOFFSET_IS_NOT_INTEGER);
438+
}
439+
440+
if (!BitmapManager.TryValidateBitfieldOffset(offset, (byte)bitCount, multiplyOffset, out _, out _))
396441
{
397442
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BITOFFSET_IS_NOT_INTEGER);
398443
}
444+
399445
var offsetSlice = parseState.GetArgSliceByRef(currTokenIdx++);
400446

401447
// GET Subcommand takes 2 args, encoding and offset
@@ -467,17 +513,23 @@ private bool StringBitFieldReadOnly<TGarnetApi>(ref TGarnetApi storageApi)
467513
// GET Subcommand takes 2 args, encoding and offset
468514

469515
// Process encoding argument
470-
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldEncoding(currTokenIdx, out _, out _))
516+
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldEncoding(currTokenIdx, out var bitCount, out _))
471517
{
472518
return AbortWithErrorMessage(CmdStrings.RESP_ERR_INVALID_BITFIELD_TYPE);
473519
}
474520
var encodingSlice = parseState.GetArgSliceByRef(currTokenIdx++);
475521

476522
// Process offset argument
477-
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldOffset(currTokenIdx, out _, out _))
523+
if ((currTokenIdx >= parseState.Count) || !parseState.TryGetBitfieldOffset(currTokenIdx, out var offset, out var multiplyOffset))
478524
{
479525
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BITOFFSET_IS_NOT_INTEGER);
480526
}
527+
528+
if (!BitmapManager.TryValidateBitfieldOffset(offset, (byte)bitCount, multiplyOffset, out _, out _))
529+
{
530+
return AbortWithErrorMessage(CmdStrings.RESP_ERR_GENERIC_BITOFFSET_IS_NOT_INTEGER);
531+
}
532+
481533
var offsetSlice = parseState.GetArgSliceByRef(currTokenIdx++);
482534

483535
secondaryCommandArgs.Add((RespCommand.GET, [commandSlice, encodingSlice, offsetSlice]));

‎libs/server/Resp/Bitmap/BitmapManager.cs‎

Lines changed: 97 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT license.
33

4+
using System;
45
using System.Runtime.CompilerServices;
6+
using Garnet.common;
57

68
#pragma warning disable IDE0054 // Use compound assignment
79
#pragma warning disable IDE1006 // Naming Styles
@@ -14,12 +16,103 @@ namespace Garnet.server
1416
public unsafe partial class BitmapManager
1517
{
1618
static readonly byte[] lookup = [0x0, 0x8, 0x4, 0xc, 0x2, 0xa, 0x6, 0xe, 0x1, 0x9, 0x5, 0xd, 0x3, 0xb, 0x7, 0xf];
19+
internal const int MaxBitmapPayloadBytes = 512 * 1024 * 1024;
20+
internal const long MaxOffsetForBitmapLength = (MaxBitmapPayloadBytes * 8L) - 1;
1721

1822
[MethodImpl(MethodImplOptions.AggressiveInlining)]
19-
private static int Index(long offset) => (int)(offset >> 3);
23+
internal static bool IsValidBitOffset(long offset)
24+
=> (ulong)offset <= (ulong)MaxOffsetForBitmapLength;
2025

2126
[MethodImpl(MethodImplOptions.AggressiveInlining)]
22-
private static int LengthInBytes(long offset) => (int)((offset >> 3) + 1);
27+
internal static bool TryValidateLengthInBytes(long offset, out int lengthInBytes)
28+
{
29+
lengthInBytes = default;
30+
if (!IsValidBitOffset(offset))
31+
return false;
32+
33+
lengthInBytes = (int)((offset >> 3) + 1);
34+
return true;
35+
}
36+
37+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
38+
internal static bool TryValidateBitfieldOffset(long offset, byte bitCount, bool multiplyOffset, out long normalizedOffset, out long endOffset)
39+
{
40+
normalizedOffset = default;
41+
endOffset = default;
42+
43+
if (bitCount == 0)
44+
return false;
45+
46+
try
47+
{
48+
normalizedOffset = multiplyOffset
49+
? checked(offset * bitCount)
50+
: offset;
51+
52+
if (normalizedOffset < 0)
53+
return false;
54+
55+
endOffset = checked(normalizedOffset + bitCount - 1);
56+
}
57+
catch (OverflowException)
58+
{
59+
return false;
60+
}
61+
62+
return IsValidBitOffset(endOffset);
63+
}
64+
65+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
66+
internal static bool TryValidateBitPosOffsets(long startOffset, long endOffset, byte offsetType, bool hasStartOffset, bool hasEndOffset)
67+
{
68+
// BYTE mode uses byte index bounds; BIT mode uses bit index bounds.
69+
var maxOffset = offsetType == 0x1
70+
? MaxOffsetForBitmapLength
71+
: MaxBitmapPayloadBytes - 1L;
72+
73+
if (hasStartOffset && (startOffset < -maxOffset || startOffset > maxOffset))
74+
return true;
75+
76+
if (hasEndOffset && (endOffset < -maxOffset || endOffset > maxOffset))
77+
return true;
78+
79+
return false;
80+
}
81+
82+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
83+
internal static void NormalizeBitCountOffsets(ref long startOffset, ref long endOffset, byte offsetType)
84+
{
85+
// BYTE mode uses byte index bounds; BIT mode uses bit index bounds.
86+
var maxOffset = offsetType == 0x1
87+
? MaxOffsetForBitmapLength
88+
: MaxBitmapPayloadBytes - 1L;
89+
90+
if (startOffset < -maxOffset)
91+
startOffset = -maxOffset;
92+
else if (startOffset > maxOffset)
93+
startOffset = maxOffset;
94+
95+
if (endOffset < -maxOffset)
96+
endOffset = -maxOffset;
97+
else if (endOffset > maxOffset)
98+
endOffset = maxOffset;
99+
}
100+
101+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
102+
private static int Index(long offset)
103+
{
104+
if (!IsValidBitOffset(offset))
105+
throw new GarnetException("ERR value is not an integer or out of range.");
106+
return (int)(offset >> 3);
107+
}
108+
109+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
110+
private static int LengthInBytes(long offset)
111+
{
112+
if (!TryValidateLengthInBytes(offset, out var lengthInBytes))
113+
throw new GarnetException("ERR value is not an integer or out of range.");
114+
return lengthInBytes;
115+
}
23116

24117
/// <summary>
25118
/// Check to see if offset contained by value size
@@ -101,8 +194,8 @@ public static byte GetBit(long offset, byte* value, int valLen)
101194
}
102195

103196
[MethodImpl(MethodImplOptions.AggressiveInlining)]
104-
private static long ProcessNegativeOffset(long offset, int valLen)
105-
=> (offset % valLen) + valLen;
197+
private static long ProcessNegativeOffset(long offset, long valLen)
198+
=> valLen <= 0 ? 0 : (offset % valLen) + valLen;
106199

107200
[MethodImpl(MethodImplOptions.AggressiveInlining)]
108201
private static byte reverse(byte n)

‎libs/server/Resp/Bitmap/BitmapManagerBitCount.cs‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ public static long BitCountDriver(long startOffset, long endOffset, byte offsetT
6565
{
6666
long count = 0;
6767

68+
NormalizeBitCountOffsets(ref startOffset, ref endOffset, offsetType);
69+
6870
// BYTE indexing
6971
if (offsetType == 0x0)
7072
{
@@ -80,8 +82,21 @@ public static long BitCountDriver(long startOffset, long endOffset, byte offsetT
8082
}
8183
else // BIT Flag
8284
{
83-
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen * 8) : startOffset;
84-
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen * 8) : endOffset;
85+
var bitLen = (long)valLen * 8;
86+
if (bitLen == 0)
87+
return 0;
88+
89+
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, bitLen) : startOffset;
90+
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, bitLen) : endOffset;
91+
92+
if (startOffset >= bitLen)
93+
return 0;
94+
95+
if (startOffset > endOffset)
96+
return 0;
97+
98+
endOffset = endOffset >= bitLen ? bitLen - 1 : endOffset;
99+
85100
count += BitIndexCount(value, startOffset, endOffset);
86101

87102
// Adjust offsets to skip first and last byte

‎libs/server/Resp/Bitmap/BitmapManagerBitPos.cs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,9 @@ public static long BitPosDriver(byte* input, int inputLen, long startOffset, lon
3737
}
3838
else
3939
{
40-
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen * 8) : startOffset;
41-
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen * 8) : endOffset;
40+
var bitLen = (long)inputLen * 8;
41+
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, bitLen) : startOffset;
42+
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, bitLen) : endOffset;
4243

4344
var startByteIndex = startOffset >> 3;
4445
var endByteIndex = endOffset >> 3;
@@ -49,7 +50,7 @@ public static long BitPosDriver(byte* input, int inputLen, long startOffset, lon
4950
if (startByteIndex > endByteIndex) // If start offset beyond endOffset return 0
5051
return -1;
5152

52-
endOffset = endByteIndex >= inputLen ? inputLen << 3 : endOffset;
53+
endOffset = endByteIndex >= inputLen ? bitLen : endOffset;
5354

5455
// BIT search
5556
return BitPosBitSearch(input, inputLen, startOffset, endOffset, searchFor);
@@ -69,11 +70,11 @@ private static long BitPosBitSearch(byte* input, long inputLen, long startBitOff
6970
{
7071
var searchBit = searchFor == 1;
7172
var invalidPayload = (byte)(searchBit ? 0x00 : 0xff);
72-
var currentBitOffset = (int)startBitOffset;
73+
var currentBitOffset = startBitOffset;
7374
while (currentBitOffset <= endBitOffset)
7475
{
75-
var byteIndex = currentBitOffset >> 3;
76-
var leftBitOffset = currentBitOffset & 7;
76+
var byteIndex = (int)(currentBitOffset >> 3);
77+
var leftBitOffset = (int)(currentBitOffset & 7);
7778
var boundary = 8 - leftBitOffset;
7879
var rightBitOffset = currentBitOffset + boundary <= endBitOffset ? leftBitOffset + boundary : (int)(endBitOffset & 7) + 1;
7980

0 commit comments

Comments
 (0)