Query : Fixes OFFSET, LIMIT, TOP data types to match those coming from query plan (#4939)

## Query : Fixes OFFSET, LIMIT, TOP data types to match those coming
from query plan

This change ensures that the data types of OFFSET, LIMIT and TOP fields
on QueryInfo and HybridSearchQueryInfo objects match those coming from
the query plan. Currently the queryplan/service uses ULONG which has a
range of 0x00000000 to 0xFFFFFFFF. Since QueryInfo fields use int?, this
causes an overflow for a value that's larger than 0x7FFFFFFF for any of
these fields.

## Type of change

Please delete options that are not relevant.

- [x] Bug fix (non-breaking change which fixes an issue)
- [] New feature (non-breaking change which adds functionality)
- [] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [] This change requires a documentation update

Author: adityasa

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs

@@ -108,19 +108,27 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
                     maxConcurrency: maxConcurrency);
 
                 if (hybridSearchQueryInfo.Skip != null)
-                {
+                {
+                    Debug.Assert(hybridSearchQueryInfo.Skip.Value <= int.MaxValue, "PipelineFactory Assert!", "Skip value must be <= int.MaxValue");
+
+                    int skipCount = (int)hybridSearchQueryInfo.Skip.Value;
+
                     MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage;
                     monadicCreatePipelineStage = (continuationToken) => SkipQueryPipelineStage.MonadicCreate(
-                        hybridSearchQueryInfo.Skip.Value,
+                        skipCount,
                         continuationToken,
                         monadicCreateSourceStage);
                 }
 
                 if (hybridSearchQueryInfo.Take != null)
-                {
+                {
+                    Debug.Assert(hybridSearchQueryInfo.Take.Value <= int.MaxValue, "PipelineFactory Assert!", "Take value must be <= int.MaxValue");
+
+                    int takeCount = (int)hybridSearchQueryInfo.Take.Value;
+
                     MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage;
                     monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateLimitStage(
-                        hybridSearchQueryInfo.Take.Value,
+                        takeCount,
                         requestContinuationToken,
                         monadicCreateSourceStage);
                 }
@@ -150,7 +158,7 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
             long optimalPageSize = maxItemCount;
             if (queryInfo.HasOrderBy)
             {
-                int top;
+                uint top;
                 if (queryInfo.HasTop && (queryInfo.Top.Value > 0))
                 {
                     top = queryInfo.Top.Value;
@@ -162,6 +170,11 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
                 else
                 {
                     top = 0;
+                }
+
+                if (top > int.MaxValue)
+                {
+                    throw new ArgumentOutOfRangeException(nameof(queryInfo.Top.Value));
                 }
 
                 if (top > 0)
@@ -257,27 +270,39 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
 
             if (queryInfo.HasOffset)
             {
+                Debug.Assert(queryInfo.Offset.Value <= int.MaxValue, "PipelineFactory Assert!", "Offset value must be <= int.MaxValue");
+
+                int offsetCount = (int)queryInfo.Offset.Value;
+
                 MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage;
                 monadicCreatePipelineStage = (continuationToken) => SkipQueryPipelineStage.MonadicCreate(
-                    queryInfo.Offset.Value,
+                    offsetCount,
                     continuationToken,
                     monadicCreateSourceStage);
             }
 
             if (queryInfo.HasLimit)
             {
+                Debug.Assert(queryInfo.Limit.Value <= int.MaxValue, "PipelineFactory Assert!", "Limit value must be <= int.MaxValue");
+
+                int limitCount = (int)queryInfo.Limit.Value;
+
                 MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage;
                 monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateLimitStage(
-                    queryInfo.Limit.Value,
+                    limitCount,
                     continuationToken,
                     monadicCreateSourceStage);
             }
 
             if (queryInfo.HasTop)
             {
+                Debug.Assert(queryInfo.Top.Value <= int.MaxValue, "PipelineFactory Assert!", "Top value must be <= int.MaxValue");
+
+                int topCount = (int)queryInfo.Top.Value;
+
                 MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage;
                 monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateTopStage(
-                    queryInfo.Top.Value,
+                    topCount,
                     continuationToken,
                     monadicCreateSourceStage);
             }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs

@@ -5,7 +5,8 @@
 namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip
 {
     using System;
-    using System.Collections.Generic;
+    using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
@@ -32,7 +33,7 @@ private ClientSkipQueryPipelineStage(
                 int offsetCount,
                 CosmosElement continuationToken,
                 MonadicCreatePipelineStage monadicCreatePipelineStage)
-            {
+            {
                 if (monadicCreatePipelineStage == null)
                 {
                     throw new ArgumentNullException(nameof(monadicCreatePipelineStage));
@@ -120,7 +121,9 @@ public override async ValueTask<bool> MoveNextAsync(ITrace trace, CancellationTo
                 IReadOnlyList<CosmosElement> documentsAfterSkip = sourcePage.Documents.Skip(this.skipCount).ToList();
 
                 int numberOfDocumentsSkipped = sourcePage.Documents.Count - documentsAfterSkip.Count;
-                this.skipCount -= numberOfDocumentsSkipped;
+                this.skipCount -= numberOfDocumentsSkipped;
+
+                Debug.Assert(this.skipCount >= 0, $"{nameof(SkipQueryPipelineStage)} Assert!", "this.skipCount should be greater than or equal to 0");
 
                 QueryState state;
                 if ((sourcePage.State != null) && (sourcePage.DisallowContinuationTokenMessage == null))

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs

@@ -22,7 +22,7 @@ protected SkipQueryPipelineStage(
             long skipCount)
             : base(source)
         {
-            if (skipCount > int.MaxValue)
+            if (skipCount > int.MaxValue || skipCount < 0)
             {
                 throw new ArgumentOutOfRangeException(nameof(skipCount));
             }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs

@@ -38,7 +38,7 @@ private ClientTakeQueryPipelineStage(
             {
                 if (limitCount < 0)
                 {
-                    throw new ArgumentException($"{nameof(limitCount)}: {limitCount} must be a non negative number.");
+                    throw new ArgumentException($"{nameof(limitCount)}: {limitCount} must be a non negative number.");
                 }
 
                 if (monadicCreatePipelineStage == null)
@@ -256,7 +256,7 @@ private sealed class LimitContinuationToken : TakeContinuationToken
                 /// <param name="limit">The limit to the number of document drained for the remainder of the query.</param>
                 /// <param name="sourceToken">The continuation token for the source component of the query.</param>
                 public LimitContinuationToken(int limit, string sourceToken)
-                {
+                {
                     if (limit < 0)
                     {
                         throw new ArgumentException($"{nameof(limit)} must be a non negative number.");
@@ -330,7 +330,7 @@ private sealed class TopContinuationToken : TakeContinuationToken
                 /// <param name="top">The limit to the number of document drained for the remainder of the query.</param>
                 /// <param name="sourceToken">The continuation token for the source component of the query.</param>
                 public TopContinuationToken(int top, string sourceToken)
-                {
+                {
                     this.Top = top;
                     this.SourceToken = sourceToken;
                 }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs

@@ -19,7 +19,7 @@ protected TakeQueryPipelineStage(
             IQueryPipelineStage source,
             int takeCount)
             : base(source)
-        {
+        {
             this.takeCount = takeCount;
         }
 

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/HybridSearchQueryInfo.cs

@@ -38,14 +38,14 @@ public QueryInfo ProjectionQueryInfo
         }
 
         [JsonProperty("skip")]
-        public int? Skip
+        public uint? Skip
         {
             get;
             set;
         }
 
         [JsonProperty("take")]
-        public int? Take
+        public uint? Take
         {
             get;
             set;

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs

@@ -24,21 +24,21 @@ public DistinctQueryType DistinctType
         }
 
         [JsonProperty("top")]
-        public int? Top
+        public uint? Top
         {
             get;
             set;
         }
 
         [JsonProperty("offset")]
-        public int? Offset
+        public uint? Offset
         {
             get;
             set;
         }
 
         [JsonProperty("limit")]
-        public int? Limit
+        public uint? Limit
         {
             get;
             set;

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs

@@ -318,9 +318,56 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
                    {
                        DateParseHandling = DateParseHandling.None,
                        MaxDepth = 64, // https://github.com/advisories/GHSA-5crp-9r3c-p9vr
-                   });
+                   });
+
+            if (!this.ValidateQueryExecutionInfo(queryInfoInternal, out ArgumentException innerException))
+            {
+                return TryCatch<PartitionedQueryExecutionInfoInternal>.FromException(
+                    new ExpectedQueryPartitionProviderException(
+                        serializedQueryExecutionInfo,
+                        innerException));
+            }
 
             return TryCatch<PartitionedQueryExecutionInfoInternal>.FromResult(queryInfoInternal);
+        }
+
+        private bool ValidateQueryExecutionInfo(PartitionedQueryExecutionInfoInternal queryExecutionInfo, out ArgumentException innerException)
+        {
+            if (queryExecutionInfo.QueryInfo?.Limit.HasValue == true &&
+                queryExecutionInfo.QueryInfo.Limit.Value > int.MaxValue)
+            {
+                innerException = new ArgumentOutOfRangeException("QueryInfo.Limit");
+                return false;
+            }
+
+            if (queryExecutionInfo.QueryInfo?.Offset.HasValue == true &&
+                queryExecutionInfo.QueryInfo.Offset.Value > int.MaxValue)
+            {
+                innerException = new ArgumentOutOfRangeException("QueryInfo.Offset");
+                return false;
+            }
+
+            if (queryExecutionInfo.QueryInfo?.Top.HasValue == true &&
+                queryExecutionInfo.QueryInfo.Top.Value > int.MaxValue)
+            {
+                innerException = new ArgumentOutOfRangeException("QueryInfo.Top");
+                return false;
+            }
+
+            if ((queryExecutionInfo.HybridSearchQueryInfo?.Skip ?? 0) > int.MaxValue)
+            {
+                innerException = new ArgumentOutOfRangeException("HybridSearchQueryInfo.Skip");
+                return false;
+            }
+
+            if ((queryExecutionInfo.HybridSearchQueryInfo?.Take ?? 0) > int.MaxValue)
+            {
+                innerException = new ArgumentOutOfRangeException("HybridSearchQueryInfo.Take");
+                return false;
+            }
+
+            innerException = null;
+            return true;
         }
 
         internal static TryCatch<IntPtr> TryCreateServiceProvider(string queryEngineConfiguration)

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NegativeQueryTests.cs

@@ -101,5 +101,47 @@ await this.CreateIngestQueryDeleteAsync(
                 QueryTestsBase.NoDocuments,
                 ImplementationAsync);
         }
+
+        [TestMethod]
+        public async Task TestTopOffsetLimitClientRanges()
+        {
+            async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject> documents)
+            {
+                await QueryTestsBase.NoOp();
+
+                foreach((string parameterName, string query) in new[]
+                    {
+                        ("QueryInfo.Offset", "SELECT c.name FROM c OFFSET 2147483648 LIMIT 10"),
+                        ("QueryInfo.Limit",  "SELECT c.name FROM c OFFSET 10 LIMIT 2147483648"),
+                        ("QueryInfo.Top",    "SELECT TOP 2147483648 c.name FROM c"),
+                    })
+                try
+                {
+                    List<Document> expectedValues = new List<Document>();
+                    FeedIterator<Document> resultSetIterator = container.GetItemQueryIterator<Document>(
+                        query,
+                        requestOptions: new QueryRequestOptions() { MaxConcurrency = 0 });
+
+                    while (resultSetIterator.HasMoreResults)
+                    {
+                        expectedValues.AddRange(await resultSetIterator.ReadNextAsync());
+                    }
+
+                    Assert.Fail("Expected to get an exception for this query.");
+                }
+                catch (CosmosException e)
+                {
+                    Assert.IsTrue(e.StatusCode == HttpStatusCode.BadRequest);
+                    Assert.IsTrue(e.InnerException?.InnerException is ArgumentException ex &&
+                        ex.Message.Contains(parameterName));
+                }
+            }
+
+            await this.CreateIngestQueryDeleteAsync(
+                ConnectionModes.Direct | ConnectionModes.Gateway,
+                CollectionTypes.MultiPartition,
+                QueryTestsBase.NoDocuments,
+                ImplementationAsync);
+        }
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SkipTakeQueryTests.cs

@@ -185,5 +185,43 @@ ORDER BY c.guid
                 }
             }
         }
+
+        [TestMethod]
+        public async Task TestTopOffsetLimitClientRanges()
+        {
+            async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject> documents)
+            {
+                await QueryTestsBase.NoOp();
+
+                foreach (string query in new[]
+                    {
+                        "SELECT c.name FROM c OFFSET 0 LIMIT 10",
+                        "SELECT c.name FROM c OFFSET 2147483647 LIMIT 10",
+                        "SELECT c.name FROM c OFFSET 10 LIMIT 0",
+                        "SELECT c.name FROM c OFFSET 10 LIMIT 2147483647",
+                        "SELECT TOP 0 c.name FROM c",
+                        "SELECT TOP 2147483647 c.name FROM c",
+                    })
+                {
+                    List<CosmosElement> expectedValues = new List<CosmosElement>();
+                    FeedIterator<CosmosElement> resultSetIterator = container.GetItemQueryIterator<CosmosElement>(
+                        query,
+                        requestOptions: new QueryRequestOptions() { MaxConcurrency = 0 });
+
+                    while (resultSetIterator.HasMoreResults)
+                    {
+                        expectedValues.AddRange(await resultSetIterator.ReadNextAsync());
+                    }
+
+                    Assert.AreEqual(0, expectedValues.Count);
+                }
+            }
+
+            await this.CreateIngestQueryDeleteAsync(
+                ConnectionModes.Direct | ConnectionModes.Gateway,
+                CollectionTypes.MultiPartition,
+                QueryTestsBase.NoDocuments,
+                ImplementationAsync);
+        }
     }
 }