fix(dynamodb): enforce primary key size limits

src/main/java/io/github/hectorvent/floci/services/dynamodb/DynamoDbItemSize.java

@@ -38,7 +38,7 @@ static int calculateItemSize(JsonNode item) {
         return total;
     }
 
-    private static int attributeValueSize(JsonNode attr) {
+    static int attributeValueSize(JsonNode attr) {
         if (attr == null) return 0;
         if (attr.has("S")) return utf8Length(attr.get("S").asText());
         if (attr.has("N")) return attr.get("N").asText().length();

src/main/java/io/github/hectorvent/floci/services/dynamodb/DynamoDbService.java

@@ -3059,6 +3059,7 @@ String buildItemKey(TableDefinition table, JsonNode item, KeySurface surface) {
             throw missingKeyException(surface);
         }
         validateKeyAttributeValue(table, pkAttr, pkName, surface);
+        validateKeySize(pkAttr, true);
 
         String pk = extractScalarValue(pkAttr);
         String skName = table.getSortKeyName();
@@ -3068,11 +3069,22 @@ String buildItemKey(TableDefinition table, JsonNode item, KeySurface surface) {
                 throw missingKeyException(surface);
             }
             validateKeyAttributeValue(table, skAttr, skName, surface);
+            validateKeySize(skAttr, false);
             return pk + "#" + extractScalarValue(skAttr);
         }
         return pk;
     }
 
+    private void validateKeySize(JsonNode attr, boolean partitionKey) {
+        int limit = partitionKey ? 2048 : 1024;
+        if ((attr.has("S") || attr.has("B")) && DynamoDbItemSize.attributeValueSize(attr) > limit) {
+            throw new AwsException("ValidationException",
+                    "One or more parameter values were invalid: Size of "
+                    + (partitionKey ? "hashkey" : "rangekey")
+                    + " has exceeded the maximum size limit of " + limit + " bytes", 400);
+        }
+    }
+
     private AwsException missingKeyException(KeySurface surface) {
         if (surface == KeySurface.ITEM_BODY) {
             return new AwsException("ValidationException",

src/test/java/io/github/hectorvent/floci/services/dynamodb/DynamoDbKeySizeIntegrationTest.java

@@ -0,0 +1,73 @@
+package io.github.hectorvent.floci.services.dynamodb;
+
+import io.github.hectorvent.floci.testing.RestAssuredJsonUtils;
+import io.quarkus.test.junit.QuarkusTest;
+import io.restassured.response.ValidatableResponse;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import java.util.Base64;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import static io.restassured.RestAssured.given;
+import static org.hamcrest.Matchers.equalTo;
+
+@QuarkusTest
+class DynamoDbKeySizeIntegrationTest {
+
+    @BeforeAll
+    static void configureRestAssured() {
+        RestAssuredJsonUtils.configureAwsContentTypes();
+    }
+
+    @ParameterizedTest
+    @CsvSource({"S,pk,2048", "S,sk,1024", "B,pk,2048", "B,sk,1024"})
+    void oversizedPrimaryKeyReturnsValidationExceptionWithoutStoringItem(String type, String keyName, int limit) {
+        String table = "KeySizes-" + UUID.randomUUID();
+        request("CreateTable", Map.of(
+                "TableName", table,
+                "KeySchema", List.of(
+                        Map.of("AttributeName", "pk", "KeyType", "HASH"),
+                        Map.of("AttributeName", "sk", "KeyType", "RANGE")),
+                "AttributeDefinitions", List.of(
+                        Map.of("AttributeName", "pk", "AttributeType", type),
+                        Map.of("AttributeName", "sk", "AttributeType", type)),
+                "BillingMode", "PAY_PER_REQUEST")).statusCode(200);
+        try {
+            request("PutItem", Map.of("TableName", table, "Item", Map.of(
+                    "pk", value(type, "pk".equals(keyName) ? limit : 1),
+                    "sk", value(type, "sk".equals(keyName) ? limit : 1)))).statusCode(200);
+
+            request("PutItem", Map.of("TableName", table, "Item", Map.of(
+                    "pk", value(type, "pk".equals(keyName) ? limit + 1 : 1),
+                    "sk", value(type, "sk".equals(keyName) ? limit + 1 : 1))))
+                    .statusCode(400)
+                    .body("__type", equalTo("ValidationException"))
+                    .body("message", equalTo("One or more parameter values were invalid: Size of "
+                            + ("pk".equals(keyName) ? "hashkey" : "rangekey")
+                            + " has exceeded the maximum size limit of " + limit + " bytes"));
+
+            request("Scan", Map.of("TableName", table, "Select", "COUNT"))
+                    .statusCode(200).body("Count", equalTo(1));
+        } finally {
+            request("DeleteTable", Map.of("TableName", table)).statusCode(200);
+        }
+    }
+
+    private static Map<String, String> value(String type, int size) {
+        return Map.of(type, "B".equals(type)
+                ? Base64.getEncoder().encodeToString(new byte[size]) : "x".repeat(size));
+    }
+
+    private static ValidatableResponse request(String action, Map<String, ?> body) {
+        return given()
+                .header("X-Amz-Target", "DynamoDB_20120810." + action)
+                .contentType("application/x-amz-json-1.0")
+                .body(body)
+                .post("/")
+                .then();
+    }
+}

src/test/java/io/github/hectorvent/floci/services/dynamodb/DynamoDbKeySizeServiceTest.java

@@ -0,0 +1,144 @@
+package io.github.hectorvent.floci.services.dynamodb;
+
+import com.fasterxml.jackson.databind.node.JsonNodeFactory;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import io.github.hectorvent.floci.core.common.AwsException;
+import io.github.hectorvent.floci.core.storage.InMemoryStorage;
+import io.github.hectorvent.floci.services.dynamodb.model.AttributeDefinition;
+import io.github.hectorvent.floci.services.dynamodb.model.KeySchemaElement;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class DynamoDbKeySizeServiceTest {
+
+    private static final String REGION = "us-east-1";
+    private DynamoDbService service;
+
+    @BeforeEach
+    void setUp() {
+        service = new DynamoDbService(new InMemoryStorage<>());
+    }
+
+    @ParameterizedTest
+    @CsvSource({"S,pk,2048", "S,sk,1024", "B,pk,2048", "B,sk,1024"})
+    void acceptsLimitAndRejectsOneByteOver(String type, String keyName, int limit) {
+        createTable(type);
+        ObjectNode item = key(type, "a", "b");
+        item.set(keyName, value(type, "x".repeat(limit)));
+        service.putItem("KeySizes", item, REGION);
+        assertNotNull(service.getItem("KeySizes", item, REGION));
+
+        ObjectNode oversized = item.deepCopy();
+        oversized.set(keyName, value(type, "x".repeat(limit + 1)));
+        assertSizeError(assertThrows(AwsException.class,
+                () -> service.putItem("KeySizes", oversized, REGION)), keyName, limit);
+    }
+
+    @ParameterizedTest
+    @CsvSource({"pk,2048", "sk,1024"})
+    void measuresStringsInUtf8Bytes(String keyName, int limit) {
+        createTable("S");
+        ObjectNode item = key("S", "a", "b");
+        item.set(keyName, value("S", "\u00e9".repeat(limit / 2)));
+        service.putItem("KeySizes", item, REGION);
+        assertNotNull(service.getItem("KeySizes", item, REGION));
+
+        item.set(keyName, value("S", "\u00e9".repeat(limit / 2) + "x"));
+        assertSizeError(assertThrows(AwsException.class,
+                () -> service.putItem("KeySizes", item, REGION)), keyName, limit);
+    }
+
+    @Test
+    void oversizedKeyArgumentsAreRejected() {
+        createTable("S");
+        ObjectNode oversized = key("S", "x".repeat(2049), "b");
+        assertSizeError(assertThrows(AwsException.class,
+                () -> service.getItem("KeySizes", oversized, REGION)), "pk", 2048);
+        assertSizeError(assertThrows(AwsException.class,
+                () -> service.deleteItem("KeySizes", oversized, REGION)), "pk", 2048);
+        assertSizeError(assertThrows(AwsException.class,
+                () -> service.updateItem("KeySizes", oversized, null,
+                        "SET data = :v", null,
+                        JsonNodeFactory.instance.objectNode().set(":v", value("S", "value")),
+                        null, REGION)), "pk", 2048);
+    }
+
+    @Test
+    void oversizedBatchKeyDoesNotPartiallyWrite() {
+        createTable("S");
+        ObjectNode valid = key("S", "valid", "b");
+        ObjectNode oversized = key("S", "x".repeat(2049), "b");
+        ObjectNode first = JsonNodeFactory.instance.objectNode();
+        first.putObject("PutRequest").set("Item", valid);
+        ObjectNode second = JsonNodeFactory.instance.objectNode();
+        second.putObject("PutRequest").set("Item", oversized);
+
+        assertSizeError(assertThrows(AwsException.class,
+                () -> service.batchWriteItem(Map.of("KeySizes", List.of(first, second)), REGION)), "pk", 2048);
+        assertNull(service.getItem("KeySizes", valid, REGION));
+    }
+
+    @Test
+    void oversizedTransactionKeyDoesNotPartiallyWrite() {
+        createTable("S");
+        ObjectNode valid = key("S", "valid", "b");
+        ObjectNode oversized = key("S", "a", "x".repeat(1025));
+        ObjectNode first = JsonNodeFactory.instance.objectNode();
+        first.putObject("Put").put("TableName", "KeySizes").set("Item", valid);
+        ObjectNode second = JsonNodeFactory.instance.objectNode();
+        second.putObject("Put").put("TableName", "KeySizes").set("Item", oversized);
+
+        assertThrows(AwsException.class,
+                () -> service.transactWriteItems(List.of(first, second), REGION));
+        assertNull(service.getItem("KeySizes", valid, REGION));
+    }
+
+    @Test
+    void nonKeyValuesCanExceedKeyLimits() {
+        createTable("S");
+        ObjectNode item = key("S", "a", "b");
+        item.set("data", value("S", "x".repeat(2049)));
+        service.putItem("KeySizes", item, REGION);
+        assertEquals(item, service.getItem("KeySizes", key("S", "a", "b"), REGION));
+    }
+
+    private void createTable(String type) {
+        service.createTable("KeySizes",
+                List.of(new KeySchemaElement("pk", "HASH"), new KeySchemaElement("sk", "RANGE")),
+                List.of(new AttributeDefinition("pk", type), new AttributeDefinition("sk", type)),
+                5L, 5L, REGION);
+    }
+
+    private ObjectNode key(String type, String pk, String sk) {
+        ObjectNode item = JsonNodeFactory.instance.objectNode();
+        item.set("pk", value(type, pk));
+        item.set("sk", value(type, sk));
+        return item;
+    }
+
+    private ObjectNode value(String type, String text) {
+        String encoded = "B".equals(type)
+                ? Base64.getEncoder().encodeToString(text.getBytes(StandardCharsets.UTF_8)) : text;
+        return JsonNodeFactory.instance.objectNode().put(type, encoded);
+    }
+
+    private void assertSizeError(AwsException error, String keyName, int limit) {
+        assertEquals("ValidationException", error.getErrorCode());
+        assertEquals(400, error.getHttpStatus());
+        assertEquals("One or more parameter values were invalid: Size of "
+                + ("pk".equals(keyName) ? "hashkey" : "rangekey")
+                + " has exceeded the maximum size limit of " + limit + " bytes", error.getMessage());
+    }
+}