Fix bug of scan aggregate index returning empty non-end continuation

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/KeyValueCursorBase.java

@@ -147,18 +147,17 @@
         private final SerializationMode serializationMode;
         /*
         how we chose this "magic number":
-        The goal is to make sure an old continuation won't be accidentally parsed as: {magic_number = 1234567890L, inner_continuation = some byte array}
-        An example that can be parsed as above is:
-        byte[] data = new byte[] {0x11, (byte) 0xD2, 0x02, (byte) 0x96, 0x49, 0x00, 0x00, 0x00, 0x00, 0x0A, 0x01, 0x14};
-        where 0x11 is wire tag for sfixed64,
-        [(byte) 0xD2, 0x02, (byte) 0x96, 0x49, 0x00, 0x00, 0x00, 0x00] is 1234567890L
-        0x0A is wire tag for bytes
-        0x01 represents that the byte array contains 1 byte
-        0x14 is the value of the byte array
-
-        Because 0xD2 is a negative number in java, it is not possible for a continuation to reach this number.
+        The goal is to make sure an old continuation won't be accidentally parsed as: {magic_number = 6773487359078157740L, inner_continuation = some byte array}
+        In little endian, MAGIC_NUMBER: 6773487359078157740L is:
+        new byte[]{ (byte) 0xAC, (byte) 0xCD, 0x73, (byte) 0x98, (byte) 0xDD, 0x42, 0x00, 0x5E };
+        Note that none of those bytes are valid Tuple codes (except for the 0x00--which is also deliberate).
+        That includes the byte at position 4, that is \x98, which is the relevant byte if we had a Tuple that began with \x11.
+        The choice of 0x00 as the penultimate byte is there to protect against the only case where we don't get a valid Tuple back from a scan,
+        namely PREFIX_STRING scan. In that case, the byte string will begin with some String suffix, which theoretically could be a valid Protobuf value.
+        But that String suffix will have any \x00 bytes escaped by following it with an \xff byte.
+        So the sequence \x00\x5e would mean "end-of-string" followed by the beginning of a new Tuple value with code \x5e, which again, is invalid.
          */
-        private static final long MAGIC_NUMBER = 1234567890L;
+        private static final long MAGIC_NUMBER = 677_348_735_907_815_774_0L;
 
         public Continuation(@Nullable final byte[] lastKey, final int prefixLength, final SerializationMode serializationMode) {
             // Note that doing this without a full copy is dangerous if the array is ever mutated.
@@ -201,7 +200,7 @@
             return byteString.isEmpty() ? new byte[0] : byteString.toByteArray();
         }
 
-        public static byte[] fromRawBytes(@Nullable byte[] rawBytes) {
+        public static byte[] getInnerContinuation(@Nullable byte[] rawBytes) {
             if (rawBytes == null) {
                 return null;
             }
@@ -229,7 +228,7 @@
                 return builder.setMagicNumber(MAGIC_NUMBER).build();
             } else {
                 ByteString base = ZeroCopyByteString.wrap(Objects.requireNonNull(lastKey));
                 // even when prefixLength = lastKey.length, proto.getInnerContinuation() = ByteString.EMPTY, proto.hasContinuation() = true
                 return builder.setInnerContinuation(base.substring(prefixLength, lastKey.length)).setMagicNumber(MAGIC_NUMBER).build();
             }
         }
@@ -323,7 +322,7 @@
             reverse = scanProperties.isReverse();
 
             if (continuation != null) {
-                byte[] realContinuation = KeyValueCursorBase.Continuation.fromRawBytes(continuation);
+                byte[] realContinuation = KeyValueCursorBase.Continuation.getInnerContinuation(continuation);
                 final byte[] continuationBytes = new byte[prefixLength + realContinuation.length];
                 System.arraycopy(lowBytes, 0, continuationBytes, 0, prefixLength);
                 System.arraycopy(realContinuation, 0, continuationBytes, prefixLength, realContinuation.length);
@@ -365,7 +364,7 @@
 
         @SpotBugsSuppressWarnings(value = "EI2", justification = "copies are expensive")
         public T setContinuation(@Nullable byte[] continuation) {
-            this.continuation = continuation;
+            this.continuation = KeyValueCursorBase.Continuation.getInnerContinuation(continuation);
             return self();
         }
 

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java

@@ -133,7 +133,7 @@
  */
 @API(API.Status.INTERNAL)
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 public class RecordQueryIndexPlan implements RecordQueryPlanWithNoChildren,
                                              RecordQueryPlanWithComparisons,
                                              RecordQueryPlanWithIndex,
                                              PlannerGraphRewritable,
@@ -180,7 +180,7 @@
         this(indexName, commonPrimaryKey, scanParameters, useIndexPrefetch, fetchIndexRecords, reverse, strictlySorted, Optional.empty(), new Type.Any(), QueryPlanConstraint.noConstraint());
     }
 
     public RecordQueryIndexPlan(@Nonnull final String indexName,
                                 @Nullable final KeyExpression commonPrimaryKey,
                                 @Nonnull final IndexScanParameters scanParameters,
                                 @Nonnull final IndexFetchMethod indexFetchMethod,
@@ -222,7 +222,7 @@
     }
 
     @VisibleForTesting
     public RecordQueryIndexPlan(@Nonnull final String indexName,
                                 @Nullable final KeyExpression commonPrimaryKey,
                                 @Nonnull final IndexScanParameters scanParameters,
                                 @Nonnull final IndexFetchMethod indexFetchMethod,
@@ -310,9 +310,7 @@
         final RecordMetaData metaData = store.getRecordMetaData();
         final Index index = metaData.getIndex(indexName);
         final IndexScanBounds scanBounds = scanParameters.bind(store, index, context);
-        byte[] innerContinuation = KeyValueCursorBase.Continuation.fromRawBytes(continuation);
-
-        return store.scanIndexRemoteFetch(index, scanBounds, innerContinuation, executeProperties.asScanProperties(isReverse()), IndexOrphanBehavior.ERROR)
+        return store.scanIndexRemoteFetch(index, scanBounds, continuation, executeProperties.asScanProperties(isReverse()), IndexOrphanBehavior.ERROR)
                 .map(store::queriedRecord)
                 .map(QueryResult::fromQueriedRecord);
     }
@@ -324,15 +322,13 @@
         final RecordMetaData metaData = store.getRecordMetaData();
         final Index index = metaData.getIndex(indexName);
         final IndexScanBounds scanBounds = scanParameters.bind(store, index, context);
-        byte[] innerContinuation = KeyValueCursorBase.Continuation.fromRawBytes(continuation);
-
         if (!IndexScanType.BY_VALUE_OVER_SCAN.equals(getScanType())) {
-            return store.scanIndex(index, scanBounds, innerContinuation, executeProperties.asScanProperties(reverse));
+            return store.scanIndex(index, scanBounds, continuation, executeProperties.asScanProperties(reverse));
         }
 
         // Evaluate the scan bounds. Again, this optimization can only be done if we have a scan range
         if (!(scanBounds instanceof IndexScanRange)) {
-            return store.scanIndex(index, scanBounds, innerContinuation, executeProperties.asScanProperties(reverse));
+            return store.scanIndex(index, scanBounds, continuation, executeProperties.asScanProperties(reverse));
         }
 
         // Try to widen the scan range to include everything up
@@ -341,7 +337,7 @@
         TupleRange widenedScanRange = widenRange(tupleScanRange);
         if (widenedScanRange == null) {
             // Unable to widen the range. Fall back to the original execution.
-            return store.scanIndex(index, scanBounds, innerContinuation, executeProperties.asScanProperties(reverse));
+            return store.scanIndex(index, scanBounds, continuation, executeProperties.asScanProperties(reverse));
         }
 
         return executeEntriesWithOverScan(tupleScanRange, widenedScanRange, store, index, continuation, executeProperties);
@@ -793,16 +789,16 @@
                 return null;
             }
             // Add the prefix back to the inner continuation
-            byte[] innerContinuation = KeyValueCursorBase.Continuation.fromRawBytes(continuation);
-            return ByteArrayUtil.join(prefixBytes, innerContinuation);
+            byte[] innerContinuation = KeyValueCursorBase.Continuation.getInnerContinuation(continuation);
+            return new KeyValueCursorBase.Continuation(ByteArrayUtil.join(prefixBytes, innerContinuation), 0, serializationMode).toBytes();
         }
 
         @Override
         public RecordCursorContinuation wrapContinuation(@Nonnull final RecordCursorContinuation continuation) {
             if (continuation.isEnd()) {
                 return continuation;
             }
-            byte[] continuationBytes = KeyValueCursorBase.Continuation.fromRawBytes(continuation.toBytes());
+            byte[] continuationBytes = KeyValueCursorBase.Continuation.getInnerContinuation(continuation.toBytes());
             if (continuationBytes != null && ByteArrayUtil.startsWith(continuationBytes, prefixBytes)) {
                 // Strip away the prefix. Note that ByteStrings re-use the underlying ByteArray, so this can
                 // save a copy.
@@ -821,9 +817,10 @@
             @SuppressWarnings("squid:S3077") // array immutable once initialized, so AtomicByteArray not necessary
             @Nullable
             private volatile byte[] bytes;
-            private final KeyValueCursorBase.SerializationMode serializationMode;
+            @Nonnull
+            private KeyValueCursorBase.SerializationMode serializationMode;
 
-            private PrefixRemovingContinuation(RecordCursorContinuation baseContinuation, int prefixLength, KeyValueCursorBase.SerializationMode serializationMode) {
+            private PrefixRemovingContinuation(RecordCursorContinuation baseContinuation, int prefixLength, @Nonnull KeyValueCursorBase.SerializationMode serializationMode) {
                 this.baseContinuation = baseContinuation;
                 this.prefixLength = prefixLength;
                 this.serializationMode = serializationMode;
@@ -835,11 +832,7 @@
                 if (bytes == null) {
                     synchronized (this) {
                         if (bytes == null) {
-                            byte[] baseContinuationBytes = KeyValueCursorBase.Continuation.fromRawBytes(baseContinuation.toBytes());
-                            if (baseContinuationBytes == null) {
-                                return null;
-                            }
-                            this.bytes = Arrays.copyOfRange(baseContinuationBytes, prefixLength, baseContinuationBytes.length);
+                            bytes = KeyValueCursorBase.Continuation.getInnerContinuation(new KeyValueCursorBase.Continuation(baseContinuation.toBytes(), prefixLength, serializationMode).toBytes());
                         }
                     }
                 }
@@ -849,8 +842,8 @@
             @Nonnull
             @Override
             public ByteString toByteString() {
-                byte[] result = KeyValueCursorBase.Continuation.fromRawBytes(baseContinuation.toBytes());
-                return result == null ? ByteString.EMPTY : ByteString.copyFrom(result).substring(prefixLength);
+                byte[] bytes1 = toBytes();
+                return bytes1 == null ? ByteString.EMPTY : ByteString.copyFrom(bytes1);
             }
 
             @Override

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryScanPlan.java

@@ -36,7 +36,6 @@
 import com.apple.foundationdb.record.provider.common.StoreTimer;
 import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase;
 import com.apple.foundationdb.record.provider.foundationdb.FDBStoreTimer;
-import com.apple.foundationdb.record.provider.foundationdb.KeyValueCursorBase;
 import com.apple.foundationdb.record.query.plan.AvailableFields;
 import com.apple.foundationdb.record.query.plan.ScanComparisons;
 import com.apple.foundationdb.record.query.plan.cascades.AliasMap;
@@ -180,10 +179,8 @@ public <M extends Message> RecordCursor<QueryResult> executePlan(@Nonnull final
                                                                      @Nullable final byte[] continuation,
                                                                      @Nonnull final ExecuteProperties executeProperties) {
         final TupleRange range = comparisons.toTupleRange(store, context);
-        byte[] innerContinuation = KeyValueCursorBase.Continuation.fromRawBytes(continuation);
-
         return store.scanRecords(
-                range.getLow(), range.getHigh(), range.getLowEndpoint(), range.getHighEndpoint(), innerContinuation,
+                range.getLow(), range.getHigh(), range.getLowEndpoint(), range.getHighEndpoint(), continuation,
                 executeProperties.asScanProperties(reverse))
                 .map(store::queriedRecord)
                 .map(QueryResult::fromQueriedRecord);

fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/KeyValueCursorTest.java

@@ -26,6 +26,7 @@
 import com.apple.foundationdb.record.ExecuteState;
 import com.apple.foundationdb.record.RecordCoreException;
 import com.apple.foundationdb.record.RecordCursor;
+import com.apple.foundationdb.record.RecordCursorContinuation;
 import com.apple.foundationdb.record.RecordCursorIterator;
 import com.apple.foundationdb.record.RecordCursorProto;
 import com.apple.foundationdb.record.RecordCursorResult;
@@ -51,8 +52,10 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
 
+import java.security.SecureRandom;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Objects;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
@@ -99,48 +102,125 @@
     @ParameterizedTest
     @EnumSource(KeyValueCursorBase.SerializationMode.class)
     public void all(KeyValueCursorBase.SerializationMode serializationMode) {
         fdb.run(context -> {
             byte[] continuation = null;
             KeyValueCursor cursor = null;
             for (int i = 0; i < 5; i++) {
                 for (int j = 0; j < 5; j++) {
                     cursor = KeyValueCursor.Builder.withSubspace(subspace)
                             .setContext(context)
                             .setRange(TupleRange.ALL)
                             .setContinuation(continuation)
                             .setScanProperties(ScanProperties.FORWARD_SCAN)
                             .setSerializationMode(serializationMode)
                             .build();
                     RecordCursorResult<KeyValue> cursorResult = cursor.getNext();
                     KeyValue kv = cursorResult.get();
                     continuation = cursorResult.getContinuation().toBytes();
                     assertArrayEquals(subspace.pack(Tuple.from(i, j)), kv.getKey());
                     assertArrayEquals(Tuple.from(i, j).pack(), kv.getValue());
                 }
             }
             assertThat(cursor.getNext().hasNext(), is(false));
 
             cursor = KeyValueCursor.Builder.withSubspace(subspace)
                     .setContext(context)
                     .setRange(TupleRange.ALL)
                     .setContinuation(null)
                     .setScanProperties(new ScanProperties(ExecuteProperties.newBuilder().setReturnedRowLimit(10).build()))
                     .setSerializationMode(serializationMode)
                     .build();
             assertEquals(10, (int)cursor.getCount().join());
             cursor = KeyValueCursor.Builder.withSubspace(subspace)
                     .setContext(context)
                     .setRange(TupleRange.ALL)
                     .setContinuation(cursor.getNext().getContinuation().toBytes())
                     .setScanProperties(ScanProperties.FORWARD_SCAN)
                     .setSerializationMode(serializationMode)
                     .build();
             assertEquals(15, (int)cursor.getCount().join());
 
             return null;
         });
     }
 
+    @ParameterizedTest
+    @EnumSource(KeyValueCursorBase.SerializationMode.class)
+    public void allInARange(KeyValueCursorBase.SerializationMode serializationMode) throws InvalidProtocolBufferException {
+        // pick 2 examples that can be serialized as RecordCursorProto.KeyValueCursorContinuation, but was correctly rejected by the magic number check
+        byte[] lowBytes = new byte[]{ 0x11, (byte) 0xac,  (byte) 0xcd, (byte) 0x73, 0x01, (byte) 0xdd, 0x42, (byte) 0x98, 0x5e, 0x0A, 0x04, 0x0f, (byte) 0xdb, 0x00, 0x14 };
+        byte[] highBytes = new byte[]{ 0x18, 0x01, 0x0A, 0x02, 0x01, 0x14 };
+        RecordCursorProto.KeyValueCursorContinuation lowProto = RecordCursorProto.KeyValueCursorContinuation.parseFrom(lowBytes);
+        Assertions.assertEquals(0x5e9842dd0173cdacL, lowProto.getMagicNumber());
+        Assertions.assertEquals(ByteString.copyFrom(new byte[] { 0x0F, (byte)0xDB, 0x00, 0x14 }), lowProto.getInnerContinuation());
+        RecordCursorProto.KeyValueCursorContinuation highProto = RecordCursorProto.KeyValueCursorContinuation.parseFrom(highBytes);
+        Assertions.assertEquals(ByteString.copyFrom(new byte[] {(byte) 1, (byte) 20}), highProto.getInnerContinuation());
+
+        Tuple low = Tuple.fromBytes(lowBytes);  // should set the magic_number to 0x5e9842dd0173cdacL and the inner continuation to `\x0f\xdb\x00\x14`
+        Tuple high = Tuple.fromBytes(highBytes); // should set an unknown field 3 to 1 and the inner_continuation to `\x01\x14`
+        TupleRange tupleRange = new TupleRange(low, high, EndpointType.RANGE_INCLUSIVE, EndpointType.RANGE_INCLUSIVE);
+
+        // clear data populated by BeforeEach
+        fdb.database().run(transaction -> {
+            transaction.clear(subspace.range());
+            return null;
+        });
+
+        // Populate with data.
+        fdb.database().run(tr -> {
+            tr.set(subspace.pack(low), low.pack());
+            tr.set(subspace.pack(high), high.pack());
+            return null;
+        });
+
+
+        fdb.run(context -> {
+            byte[] continuation = null;
+            KeyValueCursor cursor = null;
+            for (int k = 0; k < 2; k++) {
+                cursor = KeyValueCursor.Builder.withSubspace(subspace)
+                        .setContext(context)
+                        .setRange(tupleRange)
+                        .setContinuation(continuation)
+                        .setScanProperties(ScanProperties.FORWARD_SCAN)
+                        .setSerializationMode(serializationMode)
+                        .build();
+                RecordCursorResult<KeyValue> cursorResult = cursor.getNext();
+                KeyValue kv = cursorResult.get();
+                continuation = cursorResult.getContinuation().toBytes();
+                if (k == 0) {
+                    assertArrayEquals(subspace.pack(low), kv.getKey());
+                    assertArrayEquals(low.pack(), kv.getValue());
+                } else {
+                    assertArrayEquals(subspace.pack(high), kv.getKey());
+                    assertArrayEquals(high.pack(), kv.getValue());
+                }
+            }
+
+            assertThat(cursor.getNext().hasNext(), is(false));
+
+
+            cursor = KeyValueCursor.Builder.withSubspace(subspace)
+                    .setContext(context)
+                    .setRange(tupleRange)
+                    .setContinuation(null)
+                    .setScanProperties(new ScanProperties(ExecuteProperties.newBuilder().setReturnedRowLimit(1).build()))
+                    .setSerializationMode(serializationMode)
+                    .build();
+            assertEquals(1, (int)cursor.getCount().join());
+            cursor = KeyValueCursor.Builder.withSubspace(subspace)
+                    .setContext(context)
+                    .setRange(tupleRange)
+                    .setContinuation(cursor.getNext().getContinuation().toBytes())
+                    .setScanProperties(ScanProperties.FORWARD_SCAN)
+                    .setSerializationMode(serializationMode)
+                    .build();
+            assertEquals(1, (int)cursor.getCount().join());
+
+            return null;
+        });
+    }
+
     @ParameterizedTest
     @EnumSource(KeyValueCursorBase.SerializationMode.class)
     public void beginsWith(KeyValueCursorBase.SerializationMode serializationMode) {
@@ -587,10 +667,28 @@
         RecordCursorProto.KeyValueCursorContinuation continuationProto = RecordCursorProto.KeyValueCursorContinuation.parseFrom(continuation);
         Assertions.assertFalse(continuationProto.hasInnerContinuation());
 
         // setting continuation = ByteString.EMPTY -> continuation.hasContinuation() = true
         byte[] continuation2 = RecordCursorProto.KeyValueCursorContinuation.newBuilder().setInnerContinuation(ByteString.EMPTY).build().toByteArray();
         RecordCursorProto.KeyValueCursorContinuation continuationProto2 = RecordCursorProto.KeyValueCursorContinuation.parseFrom(continuation2);
         Assertions.assertTrue(continuationProto2.hasInnerContinuation());
         Assertions.assertEquals(ByteString.EMPTY, continuationProto2.getInnerContinuation());
     }
+
+    @Test
+    void emptyInnerContinuationTest() {
+        // ensure that when innerContinuation is empty, the wrapped continuation is not at beginning
+        int prefixLength = 50;
+        byte[] randomBytes = new byte[prefixLength];
+        new SecureRandom().nextBytes(randomBytes);
+        RecordCursorContinuation oldContinuation = new KeyValueCursorBase.Continuation(randomBytes, prefixLength, KeyValueCursorBase.SerializationMode.TO_OLD);
+        Assertions.assertFalse(oldContinuation.isEnd());
+        Assertions.assertEquals(ByteString.EMPTY, oldContinuation.toByteString());
+
+        KeyValueCursorBase.Continuation newContinuation = new KeyValueCursorBase.Continuation(randomBytes, prefixLength, KeyValueCursorBase.SerializationMode.TO_NEW);
+        Assertions.assertFalse(newContinuation.isEnd());
+        // inner continuation is empty
+        Assertions.assertEquals(0, Objects.requireNonNull(KeyValueCursorBase.Continuation.getInnerContinuation(newContinuation.toBytes())).length);
+        // wrapped continuation is not empty
+        Assertions.assertNotEquals(ByteString.EMPTY, newContinuation.toByteString());
+    }
 }