Fixed authorization with query using resource token for a specific partition key

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs

@@ -123,6 +123,7 @@ private static async Task<TryCatch<CosmosQueryExecutionContext>> TryCreateCoreCo
                         cosmosQueryContext.QueryClient,
                         inputParameters.SqlQuerySpec,
                         cosmosQueryContext.ResourceLink,
+                        inputParameters.PartitionKey,
                         cancellationToken);
                 }
                 else

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs

@@ -61,6 +61,7 @@ internal abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsy
             Documents.ResourceType resourceType,
             Documents.OperationType operationType,
             SqlQuerySpec sqlQuerySpec,
+            PartitionKey? partitionKey,
             string supportedQueryFeatures,
             CancellationToken cancellationToken);
 

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

@@ -74,6 +74,7 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
             CosmosQueryClient client,
             SqlQuerySpec sqlQuerySpec,
             Uri resourceLink,
+            PartitionKey? partitionKey,
             CancellationToken cancellationToken = default)
         {
             if (client == null)
@@ -98,6 +99,7 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
                 ResourceType.Document,
                 OperationType.QueryPlan,
                 sqlQuerySpec,
+                partitionKey,
                 QueryPlanRetriever.SupportedQueryFeaturesString,
                 cancellationToken);
         }

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -169,6 +169,7 @@ internal override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequ
             ResourceType resourceType,
             OperationType operationType,
             SqlQuerySpec sqlQuerySpec,
+            PartitionKey? partitionKey,
             string supportedQueryFeatures,
             CancellationToken cancellationToken)
         {
@@ -178,7 +179,7 @@ internal override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequ
                 resourceType: resourceType,
                 operationType: operationType,
                 requestOptions: null,
-                partitionKey: null,
+                partitionKey: partitionKey,
                 cosmosContainerCore: this.cosmosContainerCore,
                 streamPayload: this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType),
                 requestEnricher: (requestMessage) =>

Microsoft.Azure.Cosmos/src/Resource/Settings/PermissionProperties.cs

@@ -58,7 +58,7 @@ public PermissionProperties(string id,
             this.ResourceUri = ((ContainerInlineCore)container).ClientContext.CreateLink(
                     parentLink: ((ContainerInlineCore)container).LinkUri.OriginalString,
                     uriPathSegment: Paths.DocumentsPathSegment,
-                    id: id).OriginalString;
+                    id: itemId).OriginalString;
             this.InternalResourcePartitionKey = resourcePartitionKey.InternalKey;
         }
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosPermissionTests.cs

@@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
     using System.Collections.Generic;
     using System.Net;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Fluent;
     using Microsoft.Azure.Cosmos.Scripts;
     using Microsoft.Azure.Cosmos.Utils;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -164,6 +165,125 @@ public async Task ContainerResourcePermissionTest()
             }
         }
 
+        [TestMethod]
+        [DataRow(ConnectionMode.Gateway)]
+        [DataRow(ConnectionMode.Direct)]
+        public async Task ContainerPartitionResourcePermissionTest(ConnectionMode connectionMode)
+        {
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConnectionMode = connectionMode
+            };
+
+            CosmosClient cosmosClient = TestCommon.CreateCosmosClient(cosmosClientOptions);
+
+            Database database = await cosmosClient.CreateDatabaseIfNotExistsAsync("PermissionTest");
+
+            //create user
+            string userId = Guid.NewGuid().ToString();
+            UserResponse userResponse = await database.CreateUserAsync(userId);
+            Assert.AreEqual(HttpStatusCode.Created, userResponse.StatusCode);
+            Assert.AreEqual(userId, userResponse.Resource.Id);
+            User user = userResponse.User;
+
+            //create resource
+            string containerId = Guid.NewGuid().ToString();
+
+            ContainerResponse containerResponse = await database.CreateContainerAsync(
+                id: containerId,
+                partitionKeyPath: "/id");
+
+            Assert.AreEqual(HttpStatusCode.Created, containerResponse.StatusCode);
+            Container container = containerResponse.Container;
+
+            // Create items to read
+            ToDoActivity itemAccess = ToDoActivity.CreateRandomToDoActivity();
+            ToDoActivity itemNoAccess = ToDoActivity.CreateRandomToDoActivity();
+
+            await container.CreateItemAsync<ToDoActivity>(
+                itemAccess,
+                new PartitionKey(itemAccess.id));
+
+            await container.CreateItemAsync<ToDoActivity>(
+                itemNoAccess,
+                new PartitionKey(itemNoAccess.id));
+
+            //create permission
+            string permissionId = Guid.NewGuid().ToString();
+            PartitionKey partitionKey = new PartitionKey(itemAccess.id);
+            PermissionProperties permissionProperties = new PermissionProperties(
+                permissionId,
+                PermissionMode.Read,
+                container,
+                partitionKey);
+
+            PermissionResponse permissionResponse = await user.CreatePermissionAsync(permissionProperties);
+            PermissionProperties permission = permissionResponse.Resource;
+            Assert.AreEqual(HttpStatusCode.Created, userResponse.StatusCode);
+            Assert.AreEqual(permissionId, permission.Id);
+            Assert.AreEqual(permissionProperties.PermissionMode, permission.PermissionMode);
+
+            using (CosmosClient tokenCosmosClient = TestCommon.CreateCosmosClient(clientOptions: cosmosClientOptions, resourceToken: permission.Token))
+            {
+                Container tokenContainer = tokenCosmosClient.GetContainer(database.Id, containerId);
+                await tokenContainer.ReadItemAsync<ToDoActivity>(itemAccess.id, new PartitionKey(itemAccess.id));
+
+                try
+                {
+                    await tokenContainer.ReadItemAsync<ToDoActivity>(itemNoAccess.id, new PartitionKey(itemNoAccess.id));
+                    Assert.Fail();
+                }
+                catch (CosmosException ex)
+                {
+                    Assert.AreEqual(HttpStatusCode.Forbidden, ex.StatusCode);
+                }
+
+                QueryRequestOptions queryRequestOptions = new QueryRequestOptions()
+                {
+                    PartitionKey = new PartitionKey(itemAccess.id)
+                };
+
+                FeedIterator<ToDoActivity> feedIterator = tokenContainer.GetItemQueryIterator<ToDoActivity>(
+                    queryText: "select * from T",
+                    requestOptions: queryRequestOptions);
+
+                List<ToDoActivity> result = new List<ToDoActivity>();
+                while (feedIterator.HasMoreResults)
+                {
+                    FeedResponse<ToDoActivity> toDoActivities = await feedIterator.ReadNextAsync();
+                    result.AddRange(toDoActivities);
+                }
+
+                Assert.AreEqual(1, result.Count);
+
+                // Test query with no service interop via gateway query plan to replicate x32 app
+                ContainerCore containerCore = (ContainerInlineCore)tokenContainer;
+                CrossPartitionQueryTests.MockCosmosQueryClient mock = new CrossPartitionQueryTests.MockCosmosQueryClient(
+                    clientContext: containerCore.ClientContext,
+                    cosmosContainerCore: containerCore,
+                    forceQueryPlanGatewayElseServiceInterop: true);
+
+                Container tokenGatewayQueryPlan = new ContainerCore(
+                    containerCore.ClientContext,
+                    (DatabaseCore)containerCore.Database,
+                    containerCore.Id,
+                    mock);
+
+                FeedIterator<ToDoActivity> feedIteratorGateway = tokenGatewayQueryPlan.GetItemQueryIterator<ToDoActivity>(
+                    queryText: "select * from T",
+                    requestOptions: queryRequestOptions);
+
+                List<ToDoActivity> resultGateway = new List<ToDoActivity>();
+                while (feedIteratorGateway.HasMoreResults)
+                {
+                    FeedResponse<ToDoActivity> toDoActivities = await feedIteratorGateway.ReadNextAsync();
+                    resultGateway.AddRange(toDoActivities);
+                }
+
+                Assert.AreEqual(1, resultGateway.Count);
+            }
+        }
+
         [TestMethod]
         public async Task ItemResourcePermissionTest()
         {
@@ -196,12 +316,16 @@ public async Task ItemResourcePermissionTest()
             //delete resource with PermissionMode.Read
             using (CosmosClient tokenCosmosClient = TestCommon.CreateCosmosClient(clientOptions: null, resourceToken: permission.Token))
             {
+                Container tokenContainer = tokenCosmosClient.GetContainer(this.cosmosDatabase.Id, containerId);
+                ItemResponse<dynamic> readPermissionItem = await tokenContainer.ReadItemAsync<dynamic>(itemId, partitionKey);
+                Assert.AreEqual(itemId, readPermissionItem.Resource.id.ToString());
+
                 try
                 {
-                    ItemResponse<dynamic> response = await tokenCosmosClient
-                    .GetDatabase(this.cosmosDatabase.Id)
-                    .GetContainer(containerId)
-                    .DeleteItemAsync<dynamic>(itemId, partitionKey);
+                    ItemResponse<dynamic> response = await tokenContainer.DeleteItemAsync<dynamic>(
+                        itemId,
+                        partitionKey);
+
                     Assert.Fail();
                 }
                 catch (CosmosException ex)

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs

@@ -5022,7 +5022,7 @@ public Person(
         /// <summary>
         /// A helper that forces the SDK to use the gateway or the service interop for the query plan
         /// </summary>
-        private class MockCosmosQueryClient : CosmosQueryClientCore
+        internal class MockCosmosQueryClient : CosmosQueryClientCore
         {
             /// <summary>
             /// True it will use the gateway query plan.
@@ -5047,10 +5047,24 @@ internal override bool ByPassQueryParsing()
                 return this.forceQueryPlanGatewayElseServiceInterop;
             }
 
-            internal override Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync(Uri resourceUri, ResourceType resourceType, OperationType operationType, SqlQuerySpec sqlQuerySpec, string supportedQueryFeatures, CancellationToken cancellationToken)
+            internal override Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync(
+                Uri resourceUri,
+                ResourceType resourceType,
+                OperationType operationType,
+                SqlQuerySpec sqlQuerySpec,
+                Cosmos.PartitionKey? partitionKey,
+                string supportedQueryFeatures,
+                CancellationToken cancellationToken)
             {
                 this.QueryPlanCalls++;
-                return base.ExecuteQueryPlanRequestAsync(resourceUri, resourceType, operationType, sqlQuerySpec, supportedQueryFeatures, cancellationToken);
+                return base.ExecuteQueryPlanRequestAsync(
+                    resourceUri,
+                    resourceType,
+                    operationType,
+                    sqlQuerySpec,
+                    partitionKey,
+                    supportedQueryFeatures,
+                    cancellationToken);
             }
 
             internal override Task<QueryResponseCore> ExecuteItemQueryAsync<RequestOptionType>(

changelog.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - [#1105](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1105) Custom serializer no longer calls SDK owned types that would cause serialization exceptions
 - [#1116](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1116) Fixed a deadlock on scenarios with SynchronizationContext while executing async query operations
+- [#1143](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1143) Fixed permission resource link and authorization issue when doing a query with resource token for a specific partition key
 
 ## <a name="3.5.1"/> [3.5.1](https://www.nuget.org/packages/Microsoft.Azure.Cosmos/3.5.1) - 2019-12-11