Skip to content

Commit b652997

Browse files
MMcMMMcM
authored
Add a test helper to generate secret keys … (#3683)
… from a test random seed to make them more reproducible. Also fix the test that sometimes failed due to not recognizing enough error cases.
1 parent 7b23ace commit b652997

File tree

8 files changed

+189
-110
lines changed

8 files changed

+189
-110
lines changed

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/common/FixedZeroKeyManager.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,18 +32,18 @@
3232
public class FixedZeroKeyManager implements SerializationKeyManager {
3333
private final Key encryptionKey;
3434
private final String cipherName;
35-
private final SecureRandom secureRandom;
35+
private final Random random;
3636

37-
public FixedZeroKeyManager(@Nonnull Key encryptionKey, @Nullable String cipherName, @Nullable SecureRandom secureRandom) {
37+
public FixedZeroKeyManager(@Nonnull Key encryptionKey, @Nullable String cipherName, @Nullable Random random) {
3838
if (cipherName == null) {
3939
cipherName = CipherPool.DEFAULT_CIPHER;
4040
}
41-
if (secureRandom == null) {
42-
secureRandom = new SecureRandom();
41+
if (random == null) {
42+
random = new SecureRandom();
4343
}
4444
this.encryptionKey = encryptionKey;
4545
this.cipherName = cipherName;
46-
this.secureRandom = secureRandom;
46+
this.random = random;
4747
}
4848

4949
@Override
@@ -72,6 +72,6 @@ public Random getRandom(int keyNumber) {
7272
if (keyNumber != 0) {
7373
throw new RecordSerializationException("only provide key number 0");
7474
}
75-
return secureRandom;
75+
return random;
7676
}
7777
}

‎fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/common/RollingTestKeyManager.java‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,10 @@
2121
package com.apple.foundationdb.record.provider.common;
2222

2323
import com.apple.foundationdb.record.RecordCoreArgumentException;
24+
import com.apple.foundationdb.record.util.RandomSecretUtil;
2425

25-
import javax.crypto.KeyGenerator;
2626
import javax.crypto.SecretKey;
2727
import java.security.Key;
28-
import java.security.NoSuchAlgorithmException;
2928
import java.util.HashMap;
3029
import java.util.Map;
3130
import java.util.Random;
@@ -34,13 +33,10 @@
3433
* A {@link SerializationKeyManager} that gives out lots of different keys.
3534
*/
3635
public class RollingTestKeyManager implements SerializationKeyManager {
37-
private final KeyGenerator keyGenerator;
3836
private final Map<Integer, SecretKey> keys;
3937
private final Random random;
4038

41-
public RollingTestKeyManager(long seed) throws NoSuchAlgorithmException {
42-
keyGenerator = KeyGenerator.getInstance("AES");
43-
keyGenerator.init(128);
39+
public RollingTestKeyManager(long seed) {
4440
keys = new HashMap<>();
4541
random = new Random(seed);
4642
}
@@ -49,7 +45,7 @@ public RollingTestKeyManager(long seed) throws NoSuchAlgorithmException {
4945
public int getSerializationKey() {
5046
int newKey = random.nextInt();
5147
if (!keys.containsKey(newKey)) {
52-
keys.put(newKey, keyGenerator.generateKey());
48+
keys.put(newKey, RandomSecretUtil.randomSecretKey(random));
5349
}
5450
return newKey;
5551
}

‎fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/common/TransformedRecordSerializerTest.java‎

Lines changed: 47 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@
2828
import com.apple.foundationdb.record.logging.KeyValueLogMessage;
2929
import com.apple.foundationdb.record.logging.LogMessageKeys;
3030
import com.apple.foundationdb.record.metadata.RecordType;
31+
import com.apple.foundationdb.record.util.RandomSecretUtil;
3132
import com.apple.foundationdb.tuple.Tuple;
32-
import com.apple.test.BooleanSource;
3333
import com.apple.test.ParameterizedTestUtils;
3434
import com.apple.test.RandomSeedSource;
3535
import com.apple.test.RandomizedTestUtils;
@@ -48,13 +48,11 @@
4848

4949
import javax.annotation.Nonnull;
5050
import javax.annotation.Nullable;
51-
import javax.crypto.KeyGenerator;
5251
import javax.crypto.SecretKey;
5352
import java.nio.ByteBuffer;
5453
import java.nio.ByteOrder;
5554
import java.security.InvalidKeyException;
5655
import java.security.Key;
57-
import java.security.SecureRandom;
5856
import java.util.ArrayList;
5957
import java.util.Arrays;
6058
import java.util.List;
@@ -374,11 +372,9 @@ void unrecognizedEncoding() {
374372
}
375373

376374
@ParameterizedTest
377-
@BooleanSource
378-
void encryptWhenSerializing(boolean compressToo) throws Exception {
379-
KeyGenerator keyGen = KeyGenerator.getInstance("AES");
380-
keyGen.init(128);
381-
SecretKey key = keyGen.generateKey();
375+
@MethodSource("randomAndCompressed")
376+
void encryptWhenSerializing(long seed, boolean compressToo) {
377+
SecretKey key = RandomSecretUtil.randomSecretKey(seed);
382378
TransformedRecordSerializer<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
383379
.setEncryptWhenSerializing(true)
384380
.setEncryptionKey(key)
@@ -517,7 +513,7 @@ public static Stream<Arguments> randomAndCompressed() {
517513

518514
@ParameterizedTest
519515
@MethodSource("randomAndCompressed")
520-
void encryptRollingKeys(long seed, boolean compressToo) throws Exception {
516+
void encryptRollingKeys(long seed, boolean compressToo) {
521517
RollingTestKeyManager keyManager = new RollingTestKeyManager(seed);
522518
TransformedRecordSerializer<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
523519
.setEncryptWhenSerializing(true)
@@ -550,12 +546,11 @@ void encryptRollingKeys(long seed, boolean compressToo) throws Exception {
550546
assertEquals(records, deserialized);
551547
}
552548

553-
@Test
554-
void cannotDecryptUnknownKey() throws Exception {
555-
KeyGenerator keyGen = KeyGenerator.getInstance("AES");
556-
keyGen.init(128);
557-
SecretKey key = keyGen.generateKey();
558-
SecureRandom random = new SecureRandom();
549+
@ParameterizedTest
550+
@RandomSeedSource
551+
void cannotDecryptUnknownKey(long seed) {
552+
Random random = new Random(seed);
553+
SecretKey key = RandomSecretUtil.randomSecretKey(random);
559554
TransformedRecordSerializer<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
560555
.setEncryptWhenSerializing(true)
561556
.setKeyManager(new SerializationKeyManager() {
@@ -592,14 +587,20 @@ public Random getRandom(final int keyNumber) {
592587
assertThat(e.getMessage(), containsString("only provide key number 0"));
593588
}
594589

590+
public static Stream<Arguments> randomAndJCE() {
591+
return ParameterizedTestUtils.cartesianProduct(
592+
RandomizedTestUtils.randomSeeds(0xC0DE6EEDL, 0x6EEDC0DEL),
593+
ParameterizedTestUtils.booleans("jce")
594+
);
595+
}
596+
595597
@ParameterizedTest
596-
@BooleanSource
597-
void cannotDecryptWithoutKey(boolean jce) throws Exception {
598-
KeyGenerator keyGen = KeyGenerator.getInstance("AES");
599-
keyGen.init(128);
598+
@MethodSource("randomAndJCE")
599+
void cannotDecryptWithoutKey(long seed, boolean jce) {
600+
SecretKey key = RandomSecretUtil.randomSecretKey(seed);
600601
TransformedRecordSerializer<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
601602
.setEncryptWhenSerializing(true)
602-
.setEncryptionKey(keyGen.generateKey())
603+
.setEncryptionKey(key)
603604
.setWriteValidationRatio(1.0)
604605
.build();
605606
MySimpleRecord simpleRecord = MySimpleRecord.newBuilder().setRecNo(PRIMARY_KEY_REC_NO).setStrValueIndexed("Hello").build();
@@ -621,7 +622,7 @@ void cannotDecryptWithoutKey(boolean jce) throws Exception {
621622
}
622623

623624
@Test
624-
void cannotEncryptAfterClearKey() throws Exception {
625+
void cannotEncryptAfterClearKey() {
625626
RollingTestKeyManager keyManager = new RollingTestKeyManager(0);
626627
TransformedRecordSerializerJCE.Builder<Message> builder = TransformedRecordSerializerJCE.newDefaultBuilder()
627628
.setEncryptWhenSerializing(true)
@@ -631,14 +632,15 @@ void cannotEncryptAfterClearKey() throws Exception {
631632
assertThat(e.getMessage(), containsString("cannot encrypt when serializing if encryption key is not set"));
632633
}
633634

634-
@Test
635-
void keyDoesNotMatchAlgorithm() throws Exception {
636-
KeyGenerator keyGen = KeyGenerator.getInstance("DES");
637-
keyGen.init(56);
635+
@ParameterizedTest
636+
@RandomSeedSource
637+
void keyDoesNotMatchAlgorithm(long seed) {
638+
Random random = new Random(seed);
639+
SecretKey key = RandomSecretUtil.randomSecretKey(random, "DES", 56);
638640
try {
639641
TransformedRecordSerializer<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
640642
.setEncryptWhenSerializing(true)
641-
.setEncryptionKey(keyGen.generateKey())
643+
.setEncryptionKey(key)
642644
.setWriteValidationRatio(1.0)
643645
.build();
644646
MySimpleRecord simpleRecord = MySimpleRecord.newBuilder().setRecNo(PRIMARY_KEY_REC_NO).setStrValueIndexed("Hello").build();
@@ -653,23 +655,23 @@ void keyDoesNotMatchAlgorithm() throws Exception {
653655
}
654656
}
655657

656-
@Test
657-
void changeAlgorithm() throws Exception {
658-
KeyGenerator keyGen = KeyGenerator.getInstance("AES");
659-
keyGen.init(128);
658+
@ParameterizedTest
659+
@RandomSeedSource
660+
void changeAlgorithm(long seed) {
661+
Random random = new Random(seed);
662+
SecretKey key = RandomSecretUtil.randomSecretKey(random, "AES", 128);
660663
TransformedRecordSerializer<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
661664
.setEncryptWhenSerializing(true)
662-
.setEncryptionKey(keyGen.generateKey())
665+
.setEncryptionKey(key)
663666
.setWriteValidationRatio(1.0)
664667
.build();
665668
MySimpleRecord simpleRecord = MySimpleRecord.newBuilder().setRecNo(PRIMARY_KEY_REC_NO).setStrValueIndexed("Hello").build();
666669
byte[] serialized = serialize(serializer, simpleRecord);
667-
KeyGenerator keyGen2 = KeyGenerator.getInstance("DES");
668-
keyGen2.init(56);
670+
SecretKey key2 = RandomSecretUtil.randomSecretKey(random, "DES", 56);
669671
TransformedRecordSerializer<Message> deserializer = TransformedRecordSerializerJCE.newDefaultBuilder()
670672
.setEncryptWhenSerializing(true)
671-
.setCipherName("DES")
672-
.setEncryptionKey(keyGen2.generateKey())
673+
.setCipherName("DES")
674+
.setEncryptionKey(key2)
673675
.setWriteValidationRatio(1.0)
674676
.build();
675677
RecordSerializationException e = assertThrows(RecordSerializationException.class,
@@ -679,13 +681,14 @@ void changeAlgorithm() throws Exception {
679681

680682
public static Stream<Arguments> compressedAndOrEncrypted() {
681683
return ParameterizedTestUtils.cartesianProduct(
684+
RandomizedTestUtils.randomSeeds(),
682685
ParameterizedTestUtils.booleans("compressed"),
683686
ParameterizedTestUtils.booleans("encrypted"));
684687
}
685688

686689
@ParameterizedTest
687690
@MethodSource("compressedAndOrEncrypted")
688-
void typed(boolean compressed, boolean encrypted) throws Exception {
691+
void typed(long seed, boolean compressed, boolean encrypted) {
689692
RecordSerializer<MySimpleRecord> typedSerializer = new TypedRecordSerializer<>(
690693
TestRecords1Proto.RecordTypeUnion.getDescriptor().findFieldByNumber(TestRecords1Proto.RecordTypeUnion._MYSIMPLERECORD_FIELD_NUMBER),
691694
TestRecords1Proto.RecordTypeUnion::newBuilder,
@@ -695,9 +698,7 @@ void typed(boolean compressed, boolean encrypted) throws Exception {
695698
MySimpleRecord rec = MySimpleRecord.newBuilder().setRecNo(PRIMARY_KEY_REC_NO).setStrValueIndexed(SONNET_108).build();
696699

697700
if (encrypted) {
698-
KeyGenerator keyGen = KeyGenerator.getInstance("AES");
699-
keyGen.init(128);
700-
SecretKey key = keyGen.generateKey();
701+
SecretKey key = RandomSecretUtil.randomSecretKey(seed);
701702
typedSerializer = TransformedRecordSerializerJCE.newBuilder(typedSerializer)
702703
.setEncryptWhenSerializing(true)
703704
.setEncryptionKey(key)
@@ -726,13 +727,13 @@ void typed(boolean compressed, boolean encrypted) throws Exception {
726727
assertEquals(rec, untypedDeserialized);
727728
}
728729

729-
@Test
730-
void defaultKeyManagerKey() throws Exception {
731-
KeyGenerator keyGen = KeyGenerator.getInstance("AES");
732-
keyGen.init(128);
730+
@ParameterizedTest
731+
@RandomSeedSource
732+
void defaultKeyManagerKey(long seed) {
733+
SecretKey key = RandomSecretUtil.randomSecretKey(seed);
733734
TransformedRecordSerializerJCE<Message> serializer = TransformedRecordSerializerJCE.newDefaultBuilder()
734735
.setEncryptWhenSerializing(true)
735-
.setEncryptionKey(keyGen.generateKey())
736+
.setEncryptionKey(key)
736737
.setWriteValidationRatio(1.0)
737738
.build();
738739
SerializationKeyManager keyManager = serializer.keyManager;
@@ -753,7 +754,7 @@ void defaultKeyManagerKey() throws Exception {
753754
}
754755

755756
@Test
756-
void invalidKeyManagerBuilder() throws Exception {
757+
void invalidKeyManagerBuilder() {
757758
TransformedRecordSerializerJCE.Builder<Message> builder = TransformedRecordSerializerJCE.newDefaultBuilder();
758759
builder.setEncryptWhenSerializing(true);
759760

@@ -774,7 +775,7 @@ void invalidKeyManagerBuilder() throws Exception {
774775
}
775776

776777
@Test
777-
void reuseBuilder() throws Exception {
778+
void reuseBuilder() {
778779
RollingTestKeyManager keyManager = new RollingTestKeyManager(0);
779780
TransformedRecordSerializerJCE.Builder<Message> builderWithEncryptionKey = TransformedRecordSerializerJCE
780781
.newDefaultBuilder()

‎fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/cursors/SortCursorTests.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import com.apple.foundationdb.record.sorting.MemorySortAdapter;
4040
import com.apple.foundationdb.record.sorting.MemorySortCursor;
4141
import com.apple.foundationdb.record.sorting.MemorySorter;
42+
import com.apple.foundationdb.record.util.RandomSecretUtil;
4243
import com.apple.foundationdb.tuple.Tuple;
4344
import com.apple.test.Tags;
4445
import com.beust.jcommander.internal.Lists;
@@ -52,7 +53,6 @@
5253

5354
import javax.annotation.Nonnull;
5455
import javax.annotation.Nullable;
55-
import javax.crypto.KeyGenerator;
5656
import javax.crypto.SecretKey;
5757
import java.io.File;
5858
import java.io.IOException;
@@ -393,9 +393,7 @@ public int getMaxRecordCountInMemory() {
393393

394394
private FileSortAdapterBase fileSortEncryptedAdapter() throws Exception {
395395
final SecureRandom secureRandom = new SecureRandom();
396-
final KeyGenerator keyGen = KeyGenerator.getInstance("AES");
397-
keyGen.init(128, secureRandom);
398-
final SecretKey secretKey = keyGen.generateKey();
396+
final SecretKey secretKey = RandomSecretUtil.randomSecretKey(secureRandom);
399397
return new FileSortAdapterBase() {
400398
@Override
401399
public int getMinFileRecordCount() {

‎fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/util/RandomSecretUtil.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*
2+
* RandomSecretUtil.java
3+
*
4+
* This source file is part of the FoundationDB open source project
5+
*
6+
* Copyright 2025 Apple Inc. and the FoundationDB project authors
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this file except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
21+
package com.apple.foundationdb.record.util;
22+
23+
import com.apple.foundationdb.record.RecordCoreException;
24+
import com.apple.test.RandomizedTestUtils;
25+
26+
import javax.crypto.SecretKey;
27+
import javax.crypto.SecretKeyFactory;
28+
import javax.crypto.spec.PBEKeySpec;
29+
import javax.crypto.spec.SecretKeySpec;
30+
import java.security.GeneralSecurityException;
31+
import java.security.spec.KeySpec;
32+
import java.util.Random;
33+
34+
/**
35+
* Utilities for working with reproducible encryption keys in tests.
36+
* Instead of using a {@code KeyGenerator} with a {@code SecureRandom}, whose seeding behavior is somewhat variable,
37+
* use a key derivation function from password and salt generated by an ordinary {@code Random} that the caller
38+
* seeds with {@link RandomizedTestUtils#randomSeeds} or the like.
39+
*/
40+
public class RandomSecretUtil {
41+
private RandomSecretUtil() {
42+
}
43+
44+
public static SecretKey randomSecretKey(Random r) {
45+
return randomSecretKey(r, "AES", 128);
46+
}
47+
48+
public static SecretKey randomSecretKey(long seed) {
49+
return randomSecretKey(new Random(seed));
50+
}
51+
52+
public static SecretKey randomSecretKey(Random r, String algorithm, int length) {
53+
try {
54+
final SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
55+
final KeySpec keySpec = new PBEKeySpec(RandomUtil.randomAlphanumericString(r, 10).toCharArray(),
56+
RandomUtil.randomBytes(r, 16), 8, length);
57+
final SecretKey key = keyFactory.generateSecret(keySpec);
58+
return new SecretKeySpec(key.getEncoded(), algorithm);
59+
} catch (GeneralSecurityException ex) {
60+
throw new RecordCoreException("Cannot generate secret key", ex);
61+
}
62+
}
63+
}

0 commit comments

Comments
 (0)