fix: Fix squashing of empty exclusiveStartKey (#218)

Author: lavaleri

DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/DynamoDbEncryptionInterceptor.java

@@ -125,18 +125,40 @@ public SdkRequest modifyRequest(Context.ModifyRequest context, ExecutionAttribut
                 outgoingRequest = copyOverrideConfig((PutItemRequest) originalRequest, transformedRequest);
                 break;
             } case "Query": {
+                QueryRequest queryRequest = (QueryRequest) originalRequest;
                 QueryRequest transformedRequest = transformer.QueryInputTransform(
                         QueryInputTransformInput.builder()
-                                .sdkInput((QueryRequest) originalRequest)
+                                .sdkInput(queryRequest)
                                 .build()).transformedInput();
-                outgoingRequest = copyOverrideConfig((QueryRequest) originalRequest, transformedRequest);
+
+                // Our current Java->Dafny conversion squashes empty maps into the "None" type.
+                // In order to avoid gray failures for invalid `exclusiveStartKey`,
+                // and because our transforms do not act on or modify this value currently,
+                // copy over the original `exclusiveStartKey`
+                // so that the server can correctly reject it as invalid if it is empty.
+                transformedRequest = transformedRequest.toBuilder()
+                        .exclusiveStartKey(queryRequest.exclusiveStartKey())
+                        .build();
+
+                outgoingRequest = copyOverrideConfig(queryRequest, transformedRequest);
                 break;
             } case "Scan": {
+                ScanRequest scanRequest = (ScanRequest) originalRequest;
                 ScanRequest transformedRequest = transformer.ScanInputTransform(
                         ScanInputTransformInput.builder()
-                                .sdkInput((ScanRequest) originalRequest)
+                                .sdkInput(scanRequest)
                                 .build()).transformedInput();
-                outgoingRequest = copyOverrideConfig((ScanRequest) originalRequest, transformedRequest);
+
+                // Our current Java->Dafny conversion squashes empty maps into the "None" type.
+                // In order to avoid gray failures for invalid `exclusiveStartKey`,
+                // and because our transforms do not act on or modify this value currently,
+                // copy over the original `exclusiveStartKey`
+                // so that the server can correctly reject it as invalid if it is empty.
+                transformedRequest = transformedRequest.toBuilder()
+                        .exclusiveStartKey(scanRequest.exclusiveStartKey())
+                        .build();
+
+                outgoingRequest = copyOverrideConfig(scanRequest, transformedRequest);
                 break;
             } case "TransactGetItems": {
                 TransactGetItemsRequest transformedRequest = transformer.TransactGetItemsInputTransform(

DynamoDbEncryption/runtimes/java/src/test/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/DynamoDbEncryptionInterceptorTest.java

@@ -4,6 +4,7 @@
 import software.amazon.awssdk.core.ClientType;
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.interceptor.*;
+import software.amazon.awssdk.core.util.DefaultSdkAutoConstructMap;
 import software.amazon.awssdk.services.dynamodb.model.*;
 
 import java.util.*;
@@ -550,4 +551,46 @@ public void TestPutMissingSort() {
 
         interceptor.modifyRequest(context, attributes);
     }
+
+    @Test
+    public void TestQueryPreservesEmptyExclusiveStartKey() {
+        QueryRequest oldRequest = QueryRequest.builder()
+                .tableName(TEST_TABLE_NAME)
+                .exclusiveStartKey(new HashMap<>())
+                .build();
+
+        Context.ModifyRequest context = InterceptorContext.builder()
+                .request(oldRequest)
+                .build();
+
+        ExecutionAttributes attributes = validAttributes.toBuilder()
+                .put(SdkExecutionAttribute.OPERATION_NAME, "Query")
+                .build();
+
+        SdkRequest newRequest = interceptor.modifyRequest(context, attributes);
+        assertTrue(newRequest instanceof QueryRequest);
+        assertFalse(((QueryRequest) newRequest).exclusiveStartKey() instanceof DefaultSdkAutoConstructMap);
+        assertEquals(((QueryRequest) newRequest).exclusiveStartKey().size(), 0);
+    }
+
+    @Test
+    public void TestScanPreservesEmptyExclusiveStartKey() {
+        ScanRequest oldRequest = ScanRequest.builder()
+                .tableName(TEST_TABLE_NAME)
+                .exclusiveStartKey(new HashMap<>())
+                .build();
+
+        Context.ModifyRequest context = InterceptorContext.builder()
+                .request(oldRequest)
+                .build();
+
+        ExecutionAttributes attributes = validAttributes.toBuilder()
+                .put(SdkExecutionAttribute.OPERATION_NAME, "Scan")
+                .build();
+
+        SdkRequest newRequest = interceptor.modifyRequest(context, attributes);
+        assertTrue(newRequest instanceof ScanRequest);
+        assertFalse(((ScanRequest) newRequest).exclusiveStartKey() instanceof DefaultSdkAutoConstructMap);
+        assertEquals(((ScanRequest) newRequest).exclusiveStartKey().size(), 0);
+    }
 }