Skip to content

Commit c94689b

Browse files
aulmeaulme
committed
Replace existing key backfill code with the new framework
1 parent 31edcfb commit c94689b

File tree

9 files changed

+168
-420
lines changed

9 files changed

+168
-420
lines changed

‎src/main/java/com/uid2/admin/Main.java‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import com.uid2.admin.auth.OktaAuthProvider;
66
import com.uid2.admin.auth.AuthProvider;
77
import com.uid2.admin.auth.TokenRefreshHandler;
8-
import com.uid2.admin.cloudencryption.CloudKeyRotationStrategy;
8+
import com.uid2.admin.cloudencryption.CloudKeyStatePlanner;
99
import com.uid2.admin.cloudencryption.ExpiredKeyCountRetentionStrategy;
1010
import com.uid2.admin.job.JobDispatcher;
1111
import com.uid2.admin.job.jobsync.EncryptedFilesSyncJob;
@@ -248,9 +248,9 @@ public void run() {
248248
ClientSideKeypairService clientSideKeypairService = new ClientSideKeypairService(config, auth, writeLock, clientSideKeypairStoreWriter, clientSideKeypairProvider, siteProvider, keysetManager, keypairGenerator, clock);
249249

250250
var cloudEncryptionSecretGenerator = new CloudSecretGenerator(keyGenerator);
251-
var cloudEncryptionKeyManager = new CloudEncryptionKeyManager(rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, cloudEncryptionSecretGenerator);
252251
var cloudEncryptionKeyRetentionStrategy = new ExpiredKeyCountRetentionStrategy(clock, 5);
253-
var cloudEncryptionKeyRotationStrategy = new CloudKeyRotationStrategy(cloudEncryptionSecretGenerator, clock, cloudEncryptionKeyRetentionStrategy);
252+
var cloudEncryptionKeyRotationStrategy = new CloudKeyStatePlanner(cloudEncryptionSecretGenerator, clock, cloudEncryptionKeyRetentionStrategy);
253+
var cloudEncryptionKeyManager = new CloudEncryptionKeyManager(rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, operatorKeyProvider, cloudEncryptionKeyRotationStrategy);
254254

255255
IService[] services = {
256256
new ClientKeyService(config, auth, writeLock, clientKeyStoreWriter, clientKeyProvider, siteProvider, keysetManager, keyGenerator, keyHasher),
@@ -268,7 +268,7 @@ public void run() {
268268
new PrivateSiteDataRefreshService(auth, jobDispatcher, writeLock, config, rotatingCloudEncryptionKeyProvider),
269269
new JobDispatcherService(auth, jobDispatcher),
270270
new SearchService(auth, clientKeyProvider, operatorKeyProvider),
271-
new CloudEncryptionKeyService(auth, rotatingCloudEncryptionKeyProvider, cloudEncryptionKeyStoreWriter, operatorKeyProvider, cloudEncryptionKeyRotationStrategy)
271+
new CloudEncryptionKeyService(auth, cloudEncryptionKeyManager)
272272
};
273273

274274

@@ -293,11 +293,7 @@ public void run() {
293293
}
294294

295295
synchronized (writeLock) {
296-
cloudEncryptionKeyManager.generateKeysForOperators(
297-
operatorKeyProvider.getAll(),
298-
config.getLong("cloud_encryption_key_activates_in_seconds"),
299-
config.getInteger("cloud_encryption_key_count_per_site")
300-
);
296+
cloudEncryptionKeyManager.backfillKeys();
301297
rotatingCloudEncryptionKeyProvider.loadContent();
302298
}
303299

‎src/main/java/com/uid2/admin/cloudencryption/CloudEncryptionKeyManager.java‎

Lines changed: 44 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -1,109 +1,71 @@
11
package com.uid2.admin.cloudencryption;
22

3+
import com.uid2.admin.model.CloudEncryptionKeySummary;
34
import com.uid2.admin.store.writer.CloudEncryptionKeyStoreWriter;
45
import com.uid2.shared.auth.OperatorKey;
6+
import com.uid2.shared.auth.RotatingOperatorKeyProvider;
57
import com.uid2.shared.model.CloudEncryptionKey;
68
import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
79
import org.slf4j.Logger;
810
import org.slf4j.LoggerFactory;
911

10-
import java.time.Instant;
1112
import java.util.*;
13+
import java.util.function.Function;
14+
import java.util.stream.Collectors;
1215

1316
public class CloudEncryptionKeyManager {
14-
15-
private static final Logger LOGGER = LoggerFactory.getLogger(CloudEncryptionKeyManager.class);
16-
17-
private final RotatingCloudEncryptionKeyProvider RotatingCloudEncryptionKeyProvider;
18-
private final CloudEncryptionKeyStoreWriter cloudEncryptionKeyStoreWriter;
19-
private final CloudSecretGenerator secretGenerator;
17+
private final RotatingCloudEncryptionKeyProvider keyProvider;
18+
private final RotatingOperatorKeyProvider operatorKeyProvider;
19+
private final CloudEncryptionKeyStoreWriter keyWriter;
20+
private final CloudKeyStatePlanner planner;
21+
private Collection<OperatorKey> operatorKeys;
22+
private Collection<CloudEncryptionKey> existingKeys;
2023

2124
public CloudEncryptionKeyManager(
22-
RotatingCloudEncryptionKeyProvider RotatingCloudEncryptionKeyProvider,
23-
CloudEncryptionKeyStoreWriter cloudEncryptionKeyStoreWriter,
24-
CloudSecretGenerator keyGenerator
25-
) {
26-
this.RotatingCloudEncryptionKeyProvider = RotatingCloudEncryptionKeyProvider;
27-
this.cloudEncryptionKeyStoreWriter = cloudEncryptionKeyStoreWriter;
28-
this.secretGenerator = keyGenerator;
29-
}
30-
31-
// Ensures there are `keyCountPerSite` sites for each site corresponding of operatorKeys. If there are less - create new ones.
32-
// Give all new keys for each site `activationInterval` seconds between activations, starting now
33-
public void generateKeysForOperators(
34-
Collection<OperatorKey> operatorKeys,
35-
long activationInterval,
36-
int keyCountPerSite
37-
) throws Exception {
38-
this.RotatingCloudEncryptionKeyProvider.loadContent();
39-
40-
if (operatorKeys == null || operatorKeys.isEmpty()) {
41-
throw new IllegalArgumentException("Operator keys collection must not be null or empty");
42-
}
43-
if (activationInterval <= 0) {
44-
throw new IllegalArgumentException("Key activate interval must be greater than zero");
45-
}
46-
if (keyCountPerSite <= 0) {
47-
throw new IllegalArgumentException("Key count per site must be greater than zero");
48-
}
49-
50-
for (Integer siteId : uniqueSiteIdsForOperators(operatorKeys)) {
51-
ensureEnoughKeysForSite(activationInterval, keyCountPerSite, siteId);
52-
}
25+
RotatingCloudEncryptionKeyProvider keyProvider,
26+
CloudEncryptionKeyStoreWriter keyWriter,
27+
RotatingOperatorKeyProvider operatorKeyProvider,
28+
CloudKeyStatePlanner planner) {
29+
this.keyProvider = keyProvider;
30+
this.operatorKeyProvider = operatorKeyProvider;
31+
this.keyWriter = keyWriter;
32+
this.planner = planner;
5333
}
5434

55-
private void ensureEnoughKeysForSite(long activationInterval, int keyCountPerSite, Integer siteId) throws Exception {
56-
// Check if the site ID already exists in the S3 key provider and has fewer than the required number of keys
57-
int currentKeyCount = countKeysForSite(siteId);
58-
if (currentKeyCount >= keyCountPerSite) {
59-
LOGGER.info("Site ID {} already has the required number of keys. Skipping key generation.", siteId);
60-
return;
61-
}
62-
63-
int keysToGenerate = keyCountPerSite - currentKeyCount;
64-
for (int i = 0; i < keysToGenerate; i++) {
65-
addKey(activationInterval, siteId, i);
66-
}
67-
LOGGER.info("Generated {} keys for site ID {}", keysToGenerate, siteId);
68-
}
69-
70-
private void addKey(long keyActivateInterval, Integer siteId, int keyIndex) throws Exception {
71-
long created = Instant.now().getEpochSecond();
72-
long activated = created + (keyIndex * keyActivateInterval);
73-
CloudEncryptionKey cloudEncryptionKey = generateCloudEncryptionKey(siteId, activated, created);
74-
addCloudEncryptionKey(cloudEncryptionKey);
75-
}
35+
private static final Logger logger = LoggerFactory.getLogger(CloudEncryptionKeyManager.class);
7636

77-
private static Set<Integer> uniqueSiteIdsForOperators(Collection<OperatorKey> operatorKeys) {
78-
Set<Integer> uniqueSiteIds = new HashSet<>();
79-
for (OperatorKey operatorKey : operatorKeys) {
80-
uniqueSiteIds.add(operatorKey.getSiteId());
81-
}
82-
return uniqueSiteIds;
37+
// For any site that has an operator create a new key activating now
38+
// Keep up to 5 most recent old keys per site, delete the rest
39+
public void rotateKeys() throws Exception {
40+
refreshCloudData();
41+
var desiredKeys = planner.planRotation(existingKeys, operatorKeys);
42+
writeKeys(desiredKeys);
8343
}
8444

85-
CloudEncryptionKey generateCloudEncryptionKey(int siteId, long activates, long created) throws Exception {
86-
int newKeyId = getNextKeyId();
87-
String secret = secretGenerator.generate();
88-
return new CloudEncryptionKey(newKeyId, siteId, activates, created, secret);
45+
// For any site that has an operator, if there are no keys, create a key activating now
46+
public void backfillKeys() throws Exception {
47+
refreshCloudData();
48+
var desiredKeys = planner.planBackfill(existingKeys, operatorKeys);
49+
writeKeys(desiredKeys);
8950
}
9051

91-
void addCloudEncryptionKey(CloudEncryptionKey cloudEncryptionKey) throws Exception {
92-
Map<Integer, CloudEncryptionKey> cloudEncryptionKeys = new HashMap<>(RotatingCloudEncryptionKeyProvider.getAll());
93-
cloudEncryptionKeys.put(cloudEncryptionKey.getId(), cloudEncryptionKey);
94-
cloudEncryptionKeyStoreWriter.upload(cloudEncryptionKeys, null);
52+
public List<CloudEncryptionKeySummary> getKeySummaries() throws Exception {
53+
refreshCloudData();
54+
return existingKeys.stream().map(CloudEncryptionKeySummary::fromFullKey).toList();
9555
}
9656

97-
int getNextKeyId() {
98-
Map<Integer, CloudEncryptionKey> cloudEncryptionKeys = RotatingCloudEncryptionKeyProvider.getAll();
99-
if (cloudEncryptionKeys == null || cloudEncryptionKeys.isEmpty()) {
100-
return 1;
101-
}
102-
return cloudEncryptionKeys.keySet().stream().max(Integer::compareTo).orElse(0) + 1;
57+
private void writeKeys(Set<CloudEncryptionKey> desiredKeys) throws Exception {
58+
var keysForWriting = desiredKeys.stream().collect(Collectors.toMap(
59+
CloudEncryptionKey::getId,
60+
Function.identity())
61+
);
62+
keyWriter.upload(keysForWriting, null);
10363
}
10464

105-
int countKeysForSite(int siteId) {
106-
Map<Integer, CloudEncryptionKey> allKeys = RotatingCloudEncryptionKeyProvider.getAll();
107-
return (int) allKeys.values().stream().filter(key -> key.getSiteId() == siteId).count();
65+
private void refreshCloudData() throws Exception {
66+
keyProvider.loadContent();
67+
operatorKeyProvider.loadContent(operatorKeyProvider.getMetadata());
68+
operatorKeys = operatorKeyProvider.getAll();
69+
existingKeys = keyProvider.getAll().values();
10870
}
10971
}

src/main/java/com/uid2/admin/cloudencryption/CloudKeyRotationStrategy.java renamed to src/main/java/com/uid2/admin/cloudencryption/CloudKeyStatePlanner.java

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,12 @@
1111
import java.util.stream.Collectors;
1212
import java.util.stream.Stream;
1313

14-
public class CloudKeyRotationStrategy {
14+
public class CloudKeyStatePlanner {
1515
private final CloudSecretGenerator secretGenerator;
1616
private final Clock clock;
1717
private final CloudKeyRetentionStrategy keyRetentionStrategy;
1818

19-
public CloudKeyRotationStrategy(
19+
public CloudKeyStatePlanner(
2020
CloudSecretGenerator secretGenerator,
2121
Clock clock,
2222
CloudKeyRetentionStrategy keyRetentionStrategy
@@ -26,24 +26,32 @@ public CloudKeyRotationStrategy(
2626
this.keyRetentionStrategy = keyRetentionStrategy;
2727
}
2828

29-
public Set<CloudEncryptionKey> computeDesiredKeys(
30-
Collection<CloudEncryptionKey> existingCloudKeys,
29+
public Set<CloudEncryptionKey> planRotation(
30+
Collection<CloudEncryptionKey> existingKeys,
3131
Collection<OperatorKey> operatorKeys
3232
) {
33-
var keyGenerator = new CloudEncryptionKeyGenerator(clock, secretGenerator, existingCloudKeys);
34-
Map<Integer, Set<CloudEncryptionKey>> existingKeysBySite = existingCloudKeys
33+
var keyGenerator = new CloudEncryptionKeyGenerator(clock, secretGenerator, existingKeys);
34+
Map<Integer, Set<CloudEncryptionKey>> existingKeysBySite = existingKeys
3535
.stream()
3636
.collect(Collectors.groupingBy(CloudEncryptionKey::getSiteId, Collectors.toSet()));
3737

38-
return operatorKeys
39-
.stream()
40-
.map(OperatorKey::getSiteId)
41-
.distinct()
42-
.flatMap(siteId -> desiredKeysForSite(siteId, keyGenerator, existingKeysBySite.getOrDefault(siteId, Set.of())))
38+
return siteIdsWithOperators(operatorKeys)
39+
.flatMap(siteId -> planRotationForSite(siteId, keyGenerator, existingKeysBySite.getOrDefault(siteId, Set.of())))
4340
.collect(Collectors.toSet());
4441
}
4542

46-
private Stream<CloudEncryptionKey> desiredKeysForSite(
43+
public Set<CloudEncryptionKey> planBackfill(
44+
Collection<CloudEncryptionKey> existingKeys,
45+
Collection<OperatorKey> operatorKeys
46+
) {
47+
var keyGenerator = new CloudEncryptionKeyGenerator(clock, secretGenerator, existingKeys);
48+
var siteIdsWithKeys = existingKeys.stream().map(CloudEncryptionKey::getSiteId).collect(Collectors.toSet());
49+
var sitesWithoutKeys = siteIdsWithOperators(operatorKeys).filter(siteId -> !siteIdsWithKeys.contains(siteId));
50+
var newKeys = sitesWithoutKeys.map(keyGenerator::makeNewKey);
51+
return Streams.concat(existingKeys.stream(), newKeys).collect(Collectors.toSet());
52+
}
53+
54+
private Stream<CloudEncryptionKey> planRotationForSite(
4755
Integer siteId,
4856
CloudEncryptionKeyGenerator keyGenerator,
4957
Set<CloudEncryptionKey> existingKeys
@@ -52,4 +60,11 @@ private Stream<CloudEncryptionKey> desiredKeysForSite(
5260
var newKey = keyGenerator.makeNewKey(siteId);
5361
return Streams.concat(existingKeysToRetain.stream(), Stream.of(newKey));
5462
}
63+
64+
private static Stream<Integer> siteIdsWithOperators(Collection<OperatorKey> operatorKeys) {
65+
return operatorKeys
66+
.stream()
67+
.map(OperatorKey::getSiteId)
68+
.distinct();
69+
}
5570
}

‎src/main/java/com/uid2/admin/vertx/service/CloudEncryptionKeyService.java‎

Lines changed: 10 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -3,44 +3,31 @@
33
import com.fasterxml.jackson.core.JsonProcessingException;
44
import com.fasterxml.jackson.databind.ObjectMapper;
55
import com.uid2.admin.auth.AdminAuthMiddleware;
6-
import com.uid2.admin.cloudencryption.CloudKeyRotationStrategy;
6+
import com.uid2.admin.cloudencryption.CloudEncryptionKeyManager;
77
import com.uid2.admin.model.CloudEncryptionKeyListResponse;
8-
import com.uid2.admin.model.CloudEncryptionKeySummary;
9-
import com.uid2.admin.store.writer.CloudEncryptionKeyStoreWriter;
108
import com.uid2.admin.vertx.Endpoints;
119
import com.uid2.shared.auth.Role;
12-
import com.uid2.shared.auth.RotatingOperatorKeyProvider;
13-
import com.uid2.shared.model.CloudEncryptionKey;
14-
import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
1510
import com.uid2.shared.util.Mapper;
1611
import io.vertx.core.http.HttpHeaders;
1712
import io.vertx.ext.web.Router;
1813
import io.vertx.ext.web.RoutingContext;
1914

20-
import java.util.Set;
21-
import java.util.function.Function;
22-
import java.util.stream.Collectors;
15+
// We need a slightly different thing to a rotation for on startup and on operator update
16+
// There we only want to fill in missing keys, not create new immediately active keys for sites which have keys
17+
// We could still
18+
2319

2420
public class CloudEncryptionKeyService implements IService {
2521
private final AdminAuthMiddleware auth;
26-
private final RotatingCloudEncryptionKeyProvider keyProvider;
27-
private final RotatingOperatorKeyProvider operatorKeyProvider;
28-
private final CloudEncryptionKeyStoreWriter keyWriter;
29-
private final CloudKeyRotationStrategy rotationStrategy;
22+
private final CloudEncryptionKeyManager keyManager;
3023
private static final ObjectMapper OBJECT_MAPPER = Mapper.getInstance();
3124

3225
public CloudEncryptionKeyService(
3326
AdminAuthMiddleware auth,
34-
RotatingCloudEncryptionKeyProvider keyProvider,
35-
CloudEncryptionKeyStoreWriter keyWriter,
36-
RotatingOperatorKeyProvider operatorKeyProvider,
37-
CloudKeyRotationStrategy rotationStrategy
27+
CloudEncryptionKeyManager keyManager
3828
) {
3929
this.auth = auth;
40-
this.keyProvider = keyProvider;
41-
this.operatorKeyProvider = operatorKeyProvider;
42-
this.keyWriter = keyWriter;
43-
this.rotationStrategy = rotationStrategy;
30+
this.keyManager = keyManager;
4431
}
4532

4633
@Override
@@ -56,38 +43,16 @@ public void setupRoutes(Router router) {
5643

5744
private void handleRotate(RoutingContext rc) {
5845
try {
59-
keyProvider.loadContent();
60-
operatorKeyProvider.loadContent(operatorKeyProvider.getMetadata());
61-
var operatorKeys = operatorKeyProvider.getAll();
62-
var existingKeys = keyProvider.getAll().values();
63-
64-
var desiredKeys = rotationStrategy.computeDesiredKeys(existingKeys, operatorKeys);
65-
writeKeys(desiredKeys);
66-
46+
keyManager.rotateKeys();
6747
rc.response().end();
6848
} catch (Exception e) {
6949
rc.fail(500, e);
7050
}
7151
}
7252

73-
private void writeKeys(Set<CloudEncryptionKey> desiredKeys) throws Exception {
74-
var keysForWriting = desiredKeys.stream().collect(Collectors.toMap(
75-
CloudEncryptionKey::getId,
76-
Function.identity())
77-
);
78-
keyWriter.upload(keysForWriting, null);
79-
}
80-
8153
private void handleList(RoutingContext rc) {
8254
try {
83-
keyProvider.loadContent();
84-
85-
var keySummaries = keyProvider.getAll()
86-
.values()
87-
.stream()
88-
.map(CloudEncryptionKeySummary::fromFullKey)
89-
.toList();
90-
CloudEncryptionKeyListResponse response = new CloudEncryptionKeyListResponse(keySummaries);
55+
var response = new CloudEncryptionKeyListResponse(keyManager.getKeySummaries());
9156
respondWithJson(rc, response);
9257
} catch (Exception e) {
9358
rc.fail(500, e);

‎src/main/java/com/uid2/admin/vertx/service/OperatorKeyService.java‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,8 +47,6 @@ public class OperatorKeyService implements IService {
4747
private final KeyHasher keyHasher;
4848
private final String operatorKeyPrefix;
4949
private final CloudEncryptionKeyManager cloudEncryptionKeyManager;
50-
private final long cloudEncryptionKeyActivatesInSeconds;
51-
private final int cloudEncryptionKeyCountPerSite;
5250

5351
public OperatorKeyService(JsonObject config,
5452
AdminAuthMiddleware auth,
@@ -69,8 +67,6 @@ public OperatorKeyService(JsonObject config,
6967
this.cloudEncryptionKeyManager = cloudEncryptionKeyManager;
7068

7169
this.operatorKeyPrefix = config.getString("operator_key_prefix");
72-
this.cloudEncryptionKeyActivatesInSeconds = config.getLong("cloud_encryption_key_activates_in_seconds",0L);
73-
this.cloudEncryptionKeyCountPerSite = config.getInteger("cloud_encryption_key_count_per_site",0);
7470
}
7571

7672
@Override
@@ -274,7 +270,8 @@ private void handleOperatorAdd(RoutingContext rc) {
274270
// upload to storage
275271
operatorKeyStoreWriter.upload(operators);
276272

277-
cloudEncryptionKeyManager.generateKeysForOperators(Collections.singletonList(newOperator), cloudEncryptionKeyActivatesInSeconds, cloudEncryptionKeyCountPerSite);
273+
// generate cloud encryption keys as needed
274+
cloudEncryptionKeyManager.backfillKeys();
278275

279276
// respond with new key
280277
rc.response().end(JSON_WRITER.writeValueAsString(new RevealedKey<>(newOperator, key)));
@@ -413,7 +410,7 @@ private void handleOperatorUpdate(RoutingContext rc) {
413410
operatorKeyStoreWriter.upload(operators);
414411

415412
if (siteIdChanged) {
416-
cloudEncryptionKeyManager.generateKeysForOperators(Collections.singletonList(existingOperator), cloudEncryptionKeyActivatesInSeconds, cloudEncryptionKeyCountPerSite);
413+
cloudEncryptionKeyManager.backfillKeys();
417414
}
418415

419416
// return the updated client

0 commit comments

Comments
 (0)