fix(cosmos): enforce partition isolation for point reads and deletes

compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs

@@ -0,0 +1,64 @@
+using System.Net;
+using Microsoft.Azure.Cosmos;
+
+namespace FlociAz.Compatibility;
+
+[NotInParallel]
+public sealed class CosmosPointPartitionCompatibilityTests
+{
+    [Test]
+    [Timeout(60_000)]
+    public async Task PointReadsAndDeletesStayWithinTheRequestedPartition(CancellationToken cancellationToken)
+    {
+        string endpoint = Environment.GetEnvironmentVariable("FLOCI_AZ_ENDPOINT") ?? "http://localhost:4577";
+        using var client = new CosmosClient($"{endpoint}/devstoreaccount1-cosmos/",
+            "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
+            new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway, LimitToEndpoint = true });
+        Database database = await client.CreateDatabaseAsync(
+            $"dotnet-point-partition-{Guid.NewGuid():N}", cancellationToken: cancellationToken);
+        try
+        {
+            Container container = await database.CreateContainerAsync(
+                "items", "/pk", cancellationToken: cancellationToken);
+            await container.CreateItemAsync(new { id = "same-id", pk = "alice" },
+                new PartitionKey("alice"), cancellationToken: cancellationToken);
+
+            foreach (PartitionKey missing in new[] { new PartitionKey("bob"), new PartitionKey(""), PartitionKey.Null, PartitionKey.None })
+            {
+                using ResponseMessage read = await container.ReadItemStreamAsync("same-id", missing,
+                    cancellationToken: cancellationToken);
+                await Assert.That(read.StatusCode).IsEqualTo(HttpStatusCode.NotFound);
+                using ResponseMessage delete = await container.DeleteItemStreamAsync("same-id", missing,
+                    cancellationToken: cancellationToken);
+                await Assert.That(delete.StatusCode).IsEqualTo(HttpStatusCode.NotFound);
+            }
+
+            await container.CreateItemAsync(new { id = "same-id", pk = "bob" },
+                new PartitionKey("bob"), cancellationToken: cancellationToken);
+            using ResponseMessage bobRead = await container.ReadItemStreamAsync("same-id", new PartitionKey("bob"),
+                cancellationToken: cancellationToken);
+            await Assert.That(bobRead.StatusCode).IsEqualTo(HttpStatusCode.OK);
+            using ResponseMessage bobDelete = await container.DeleteItemStreamAsync("same-id", new PartitionKey("bob"),
+                cancellationToken: cancellationToken);
+            await Assert.That(bobDelete.StatusCode).IsEqualTo(HttpStatusCode.NoContent);
+            using ResponseMessage bobMissing = await container.ReadItemStreamAsync("same-id", new PartitionKey("bob"),
+                cancellationToken: cancellationToken);
+            await Assert.That(bobMissing.StatusCode).IsEqualTo(HttpStatusCode.NotFound);
+            using ResponseMessage alice = await container.ReadItemStreamAsync("same-id", new PartitionKey("alice"),
+                cancellationToken: cancellationToken);
+            await Assert.That(alice.StatusCode).IsEqualTo(HttpStatusCode.OK);
+            await container.CreateItemAsync(new { id = "empty", pk = "" }, new PartitionKey(""),
+                cancellationToken: cancellationToken);
+            using ResponseMessage nullRead = await container.ReadItemStreamAsync("empty", PartitionKey.Null,
+                cancellationToken: cancellationToken);
+            await Assert.That(nullRead.StatusCode).IsEqualTo(HttpStatusCode.NotFound);
+            using ResponseMessage nullDelete = await container.DeleteItemStreamAsync("empty", PartitionKey.Null,
+                cancellationToken: cancellationToken);
+            await Assert.That(nullDelete.StatusCode).IsEqualTo(HttpStatusCode.NotFound);
+        }
+        finally
+        {
+            await database.DeleteAsync(cancellationToken: cancellationToken);
+        }
+    }
+}

src/main/java/io/floci/az/services/cosmos/CosmosHandler.java

@@ -1270,6 +1270,34 @@ private String encodeContinuationToken(int skip) {
     private StoredObject findDoc(AzureRequest req, String dbId, String collId, String docId) {
         Optional<StoredObject> collFound = store.get(collKey(req.accountName(), dbId, collId));
         Object defaultTtl = containerDefaultTtl(collFound);
+        String partitionHeader = req.headers().getHeaderString("x-ms-documentdb-partitionkey");
+        if (partitionHeader != null) {
+            if (collFound.isEmpty()) {
+                return null;
+            }
+            try {
+                var partition = CosmosQueryPartition.parse(partitionHeader, parseData(collFound.get()));
+                String exact = docKey(req.accountName(), dbId, collId,
+                        encodeKey(extractPartitionKeyValue(req)), docId);
+                Optional<StoredObject> found = liveDoc(store.get(exact), defaultTtl)
+                        .filter(object -> partition.test(parseData(object)));
+                if (found.isPresent()) {
+                    return found.get();
+                }
+                // Legacy storage keys stringify partition values. Compare logical values too:
+                // numbers may have equivalent encodings, while null and strings remain distinct.
+                String prefix = req.accountName() + K_DOC + dbId + "|" + collId + "|";
+                return store.scan(key -> key.startsWith(prefix) && key.endsWith("|" + docId)).stream()
+                        .filter(object -> {
+                            Map<String, Object> document = parseData(object);
+                            return docId.equals(document.get("id")) && partition.test(document);
+                        })
+                        .map(object -> liveDoc(Optional.of(object), defaultTtl))
+                        .flatMap(Optional::stream).findFirst().orElse(null);
+            } catch (IllegalArgumentException e) {
+                throw new jakarta.ws.rs.WebApplicationException(errorResponse(400, "BadRequest", e.getMessage()));
+            }
+        }
         // Fast path: construct exact key using partition key from header
         if (collFound.isPresent()) {
             String pk    = extractPartitionKeyValue(req);
@@ -1278,7 +1306,7 @@ private StoredObject findDoc(AzureRequest req, String dbId, String collId, Strin
             Optional<StoredObject> found = liveDoc(store.get(exact), defaultTtl);
             if (found.isPresent()) return found.get();
         }
-        // Fallback: scan (handles missing PK header or cross-partition reads)
+        // Only requests without a partition header may use the legacy unscoped lookup.
         String prefix = req.accountName() + K_DOC + dbId + "|" + collId + "|";
         return liveDoc(store.scan(k -> k.startsWith(prefix) && k.endsWith("|" + docId))
                 .stream().findFirst(), defaultTtl).orElse(null);

src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java

@@ -0,0 +1,102 @@
+package io.floci.az.services.cosmos;
+
+import io.quarkus.test.junit.QuarkusTest;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Map;
+
+import static io.restassured.RestAssured.given;
+import static org.hamcrest.Matchers.is;
+
+@QuarkusTest
+class CosmosPointPartitionTest {
+    private static final String BASE = "/pointpartition-cosmos/dbs/db";
+    private static final String DOCS = BASE + "/colls/items/docs";
+
+    @BeforeEach
+    void setup() {
+        given().post("/_admin/reset").then().statusCode(204);
+        given().contentType("application/json").body(Map.of("id", "db"))
+                .post("/pointpartition-cosmos/dbs").then().statusCode(201);
+        given().contentType("application/json")
+                .body("{\"id\":\"items\",\"partitionKey\":{\"paths\":[\"/pk\"],\"kind\":\"Hash\"}}")
+                .post(BASE + "/colls").then().statusCode(201);
+        create("alice");
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"[\"bob\"]", "[\"\"]", "[null]", "[{}]"})
+    void readWithMissingPartitionDoesNotFindAnotherPartitionsDocument(String partition) {
+        given().header("x-ms-documentdb-partitionkey", partition)
+                .get(DOCS + "/same-id").then().statusCode(404);
+        assertAliceExists();
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"[\"bob\"]", "[\"\"]", "[null]", "[{}]"})
+    void deleteWithMissingPartitionDoesNotDeleteAnotherPartitionsDocument(String partition) {
+        given().header("x-ms-documentdb-partitionkey", partition)
+                .delete(DOCS + "/same-id").then().statusCode(404);
+        assertAliceExists();
+    }
+
+    @Test
+    void sameIdInDifferentPartitionsRemainsIndependent() {
+        create("bob");
+        given().header("x-ms-documentdb-partitionkey", "[\"bob\"]")
+                .get(DOCS + "/same-id").then().statusCode(200).body("pk", is("bob"));
+        given().header("x-ms-documentdb-partitionkey", "[\"bob\"]")
+                .delete(DOCS + "/same-id").then().statusCode(204);
+        given().header("x-ms-documentdb-partitionkey", "[\"bob\"]")
+                .get(DOCS + "/same-id").then().statusCode(404);
+        assertAliceExists();
+    }
+
+    private void create(String partition) {
+        given().contentType("application/json")
+                .body(Map.of("id", "same-id", "pk", partition))
+                .post(DOCS).then().statusCode(201);
+    }
+
+    @Test
+    void typedPartitionValuesCannotReadOrDeleteStringPartitions() {
+        create("");
+        create("42");
+        for (String partition : new String[] {"[null]", "[42]"}) {
+            given().header("x-ms-documentdb-partitionkey", partition)
+                    .get(DOCS + "/same-id").then().statusCode(404);
+            given().header("x-ms-documentdb-partitionkey", partition)
+                    .delete(DOCS + "/same-id").then().statusCode(404);
+        }
+        given().header("x-ms-documentdb-partitionkey", "[\"\"]")
+                .get(DOCS + "/same-id").then().statusCode(200);
+        given().header("x-ms-documentdb-partitionkey", "[\"42\"]")
+                .get(DOCS + "/same-id").then().statusCode(200);
+    }
+
+    @Test
+    void equivalentNumericPartitionRepresentationsStillMatch() {
+        given().contentType("application/json").body(Map.of("id", "numeric", "pk", 42))
+                .post(DOCS).then().statusCode(201);
+        given().header("x-ms-documentdb-partitionkey", "[42.0]")
+                .get(DOCS + "/numeric").then().statusCode(200);
+    }
+
+    @Test
+    void scopedFallbackRequiresTheWholeDocumentId() {
+        given().contentType("application/json").body(Map.of("id", "prefix|missing", "pk", "alice"))
+                .post(DOCS).then().statusCode(201);
+        given().header("x-ms-documentdb-partitionkey", "[\"alice\"]")
+                .get(DOCS + "/missing").then().statusCode(404);
+        given().header("x-ms-documentdb-partitionkey", "[\"alice\"]")
+                .delete(DOCS + "/missing").then().statusCode(404);
+    }
+
+    private void assertAliceExists() {
+        given().header("x-ms-documentdb-partitionkey", "[\"alice\"]")
+                .get(DOCS + "/same-id").then().statusCode(200).body("pk", is("alice"));
+    }
+}