Skip to content

Commit 9dedb2b

Browse files
gmsdelmundogmsdelmundo
committed
Added fast forward and updated token generate/refresh simulation
1 parent 83014cd commit 9dedb2b

File tree

4 files changed

+212
-35
lines changed

4 files changed

+212
-35
lines changed

‎src/main/java/com/uid2/admin/salt/SaltRotation.java‎

Lines changed: 73 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.time.temporal.ChronoUnit;
1616
import java.util.*;
1717
import java.util.stream.Collectors;
18+
import java.util.stream.Stream;
1819

1920
public class SaltRotation {
2021
private static final long THIRTY_DAYS_IN_MS = Duration.ofDays(30).toMillis();
@@ -76,11 +77,81 @@ public Result rotateSalts(
7677
return Result.fromSnapshot(nextSnapshot);
7778
}
7879

80+
public Result rotateSaltsFastForward(
81+
SaltSnapshot snapshot,
82+
Duration[] minAges,
83+
double fraction,
84+
TargetDate targetDate,
85+
int iterations) throws Exception {
86+
var currentSnapshot = snapshot;
87+
88+
for (int i = 0; i < iterations; i++) {
89+
var preRotationSalts = currentSnapshot.getAllRotatingSalts();
90+
91+
var currentTargetDate = targetDate.plusDays(i);
92+
var nextEffective = currentTargetDate.asInstant();
93+
var nextExpires = nextEffective.plus(7, ChronoUnit.DAYS);
94+
if (nextEffective.equals(currentSnapshot.getEffective()) || nextEffective.isBefore(currentSnapshot.getEffective())) {
95+
return Result.noSnapshot("cannot create a new salt snapshot with effective timestamp equal or prior to that of an existing snapshot");
96+
}
97+
98+
// Salts that can be rotated based on their refreshFrom being at target date
99+
var refreshableSalts = findRefreshableSalts(preRotationSalts, currentTargetDate);
100+
101+
var saltsToRotate = pickSaltsToRotate(
102+
refreshableSalts,
103+
currentTargetDate,
104+
minAges,
105+
getNumSaltsToRotate(preRotationSalts, fraction)
106+
);
107+
108+
if (saltsToRotate.isEmpty()) {
109+
return Result.noSnapshot("all refreshable salts are below min rotation age");
110+
}
111+
112+
var postRotationSalts = rotateSalts(preRotationSalts, saltsToRotate, currentTargetDate);
113+
114+
LOGGER.info("Salt rotation complete target_date={}", currentTargetDate);
115+
logSaltAges("refreshable-salts", currentTargetDate, refreshableSalts);
116+
logSaltAges("rotated-salts", currentTargetDate, saltsToRotate);
117+
logSaltAges("total-salts", currentTargetDate, Arrays.asList(postRotationSalts));
118+
logBucketFormatCount(currentTargetDate, postRotationSalts);
119+
120+
currentSnapshot = new SaltSnapshot(
121+
nextEffective,
122+
nextExpires,
123+
postRotationSalts,
124+
currentSnapshot.getFirstLevelSalt());
125+
}
126+
127+
Map<Long, SaltEntry> originalSalts = Stream.of(snapshot.getAllRotatingSalts()).collect(Collectors.toMap(salt -> salt.id(), salt -> salt));
128+
List<SaltEntry> salts = new ArrayList<>();
129+
for (SaltEntry salt : currentSnapshot.getAllRotatingSalts()) {
130+
SaltEntry originalSalt = originalSalts.get(salt.id());
131+
132+
salts.add(new SaltEntry(
133+
salt.id(),
134+
salt.hashedId(),
135+
originalSalt.lastUpdated(),
136+
salt.currentSalt(),
137+
originalSalt.refreshFrom(),
138+
salt.previousSalt(),
139+
salt.currentKeySalt(),
140+
salt.previousKeySalt()
141+
));
142+
}
143+
144+
return Result.fromSnapshot(new SaltSnapshot(
145+
snapshot.getEffective(),
146+
snapshot.getExpires(),
147+
salts.toArray(new SaltEntry[salts.size()]),
148+
snapshot.getFirstLevelSalt()));
149+
}
150+
79151
public Result rotateSaltsZero(
80152
ISaltSnapshot effectiveSnapshot,
81153
TargetDate targetDate,
82-
Instant nextEffective
83-
) throws Exception {
154+
Instant nextEffective) throws Exception {
84155
var preRotationSalts = effectiveSnapshot.getAllRotatingSalts();
85156
var nextExpires = nextEffective.plus(7, ChronoUnit.DAYS);
86157

‎src/main/java/com/uid2/admin/vertx/AdminVerticle.java‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,12 @@ private void handleHealthCheck(RoutingContext rc) {
109109
}
110110

111111
private void handleUserinfo(RoutingContext rc) {
112-
if (isAuthDisabled(config)) rc.response().setStatusCode(200).end(
113-
JsonObject.of("groups", JsonArray.of("developer", "developer-elevated", "infra-admin", "admin"), "email", "[email protected]").toString());
112+
if (isAuthDisabled(config)) {
113+
rc.response().setStatusCode(200).end(
114+
JsonObject.of("groups", JsonArray.of("developer", "developer-elevated", "infra-admin", "admin"), "email", "[email protected]").toString());
115+
return;
116+
}
117+
114118
try {
115119
Jwt idJwt = this.authProvider.getIdTokenVerifier().decode(rc.user().principal().getString("id_token"), null);
116120
JsonObject jo = new JsonObject();
@@ -122,19 +126,18 @@ private void handleUserinfo(RoutingContext rc) {
122126
JsonObject userDetails = new JsonObject();
123127
userDetails.put("email", idJwt.getClaims().get("email"));
124128
userDetails.put("sub", idJwt.getClaims().get("sub"));
125-
126-
LOGGER.info("Authenticated user accessing admin page - User: {}", userDetails.toString());
129+
130+
LOGGER.info("Authenticated user accessing admin page - User: {}", userDetails);
127131
rc.put("user_details", userDetails);
128132
this.audit.log(rc, new AuditParams());
129133
}
130-
134+
131135
rc.response().setStatusCode(200).end(jo.toString());
132136
} catch (Exception e) {
133-
if (rc.session() != null) {
137+
if (rc.session() != null) {
134138
rc.session().destroy();
135139
}
136140
rc.response().putHeader("REQUIRES_AUTH", "1").setStatusCode(401).end();
137141
}
138142
}
139-
140143
}

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

Lines changed: 78 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ public class SaltService implements IService {
7777
private static final String OPERATOR_URL = "http://proxy:80";
7878
private static final Map<String, String> OPERATOR_HEADERS = Map.of("Authorization", String.format("Bearer %s", CLIENT_API_KEY));
7979

80-
private static final int IDENTITY_COUNT = 10_000;
80+
private static final int IDENTITY_COUNT = 1_000;
8181
private static final int TIMESTAMP_LENGTH = 8;
8282
private static final int IV_LENGTH = 12;
8383

@@ -110,13 +110,19 @@ public void setupRoutes(Router router) {
110110
}
111111
}, new AuditParams(List.of("fraction", "target_date"), Collections.emptyList()), Role.SUPER_USER, Role.SECRET_ROTATION));
112112

113-
router.post("/api/salt/simulate/token").blockingHandler(auth.handle((ctx) -> {
113+
router.post("/api/salt/fastForward").blockingHandler(auth.handle(ctx -> {
114+
synchronized (writeLock) {
115+
this.handleSaltFastForward(ctx);
116+
}
117+
}, new AuditParams(List.of("fraction"), Collections.emptyList()), Role.MAINTAINER, Role.SECRET_ROTATION));
118+
119+
router.post("/api/salt/simulateToken").blockingHandler(auth.handle(ctx -> {
114120
synchronized (writeLock) {
115121
this.handleSaltSimulateToken(ctx);
116122
}
117-
}, new AuditParams(List.of("target_date"), Collections.emptyList()), Role.MAINTAINER, Role.SECRET_ROTATION));
123+
}, new AuditParams(List.of("fraction", "target_date"), Collections.emptyList()), Role.MAINTAINER, Role.SECRET_ROTATION));
118124

119-
router.post("/api/salt/simulate").blockingHandler(auth.handle((ctx) -> {
125+
router.post("/api/salt/simulate").blockingHandler(auth.handle(ctx -> {
120126
synchronized (writeLock) {
121127
this.handleSaltSimulate(ctx);
122128
}
@@ -221,6 +227,40 @@ private JsonObject toJson(RotatingSaltProvider.SaltSnapshot snapshot) {
221227
return jo;
222228
}
223229

230+
private void handleSaltFastForward(RoutingContext rc) {
231+
try {
232+
final double fraction = RequestUtil.getDouble(rc, "fraction").orElse(0.002740);
233+
final int iterations = RequestUtil.getDouble(rc, "fast_forward_iterations").orElse(0.0).intValue();
234+
final boolean enableV4 = RequestUtil.getBoolean(rc, "enable_v4", false).orElse(false);
235+
236+
TargetDate targetDate =
237+
RequestUtil.getDate(rc, "target_date", DateTimeFormatter.ISO_LOCAL_DATE)
238+
.map(TargetDate::new)
239+
.orElse(TargetDate.now().plusDays(1));
240+
241+
saltProvider.loadContent();
242+
storageManager.archiveSaltLocations();
243+
244+
var snapshot = saltProvider.getSnapshots().getLast();
245+
246+
saltRotation.setEnableV4RawUid(enableV4);
247+
var result = saltRotation.rotateSaltsFastForward(snapshot, SALT_ROTATION_AGE_THRESHOLDS, fraction, targetDate, iterations);
248+
if (!result.hasSnapshot()) {
249+
ResponseUtil.error(rc, 200, result.getReason());
250+
return;
251+
}
252+
253+
storageManager.upload(result.getSnapshot());
254+
255+
rc.response()
256+
.putHeader(HttpHeaders.CONTENT_TYPE, "application/json")
257+
.end(toJson(result.getSnapshot()).encode());
258+
} catch (Exception e) {
259+
LOGGER.error(e.getMessage(), e);
260+
rc.fail(500, e);
261+
}
262+
}
263+
224264
private void handleSaltSimulateToken(RoutingContext rc) {
225265
try {
226266
final double fraction = RequestUtil.getDouble(rc, "fraction").orElse(0.002740);
@@ -231,16 +271,25 @@ private void handleSaltSimulateToken(RoutingContext rc) {
231271
.map(TargetDate::new)
232272
.orElse(TargetDate.now().plusDays(1));
233273

234-
// Step 1. Run /v2/token/generate for 10,000 emails
274+
// Step 1. Run /v2/token/generate for all emails
235275
List<String> emails = new ArrayList<>();
276+
Map<String, Boolean> preRotationEmailToSaltMap = new HashMap<>();
277+
RotatingSaltProvider.SaltSnapshot snapshot = saltProvider.getSnapshots().getLast();
236278
for (int j = 0; j < IDENTITY_COUNT; j++) {
237279
String email = randomEmail();
238280
emails.add(email);
281+
282+
SaltEntry salt = getSalt(email, snapshot);
283+
boolean isV4 = salt.currentKeySalt() != null && salt.currentKeySalt().key() != null && salt.currentKeySalt().salt() != null;
284+
preRotationEmailToSaltMap.put(email, isV4);
239285
}
240286

241287
Map<String, String> emailToRefreshTokenMap = new HashMap<>();
242288
Map<String, String> emailToRefreshResponseKeyMap = new HashMap<>();
243-
for (String email : emails) {
289+
for (int i = 0; i < emails.size(); i++) {
290+
LOGGER.info("Step 2 - Token Generate {}/{}", i + 1, emails.size());
291+
String email = emails.get(i);
292+
244293
JsonNode tokens = v2TokenGenerate(email);
245294
String refreshToken = tokens.at("/body/refresh_token").asText();
246295
String refreshResponseKey = tokens.at("/body/refresh_response_key").asText();
@@ -251,22 +300,25 @@ private void handleSaltSimulateToken(RoutingContext rc) {
251300

252301
// Step 2. Rotate salts
253302
saltRotation.setEnableV4RawUid(true);
254-
RotatingSaltProvider.SaltSnapshot snapshot = null;
303+
snapshot = null;
255304
for (int i = 0; i < iterations; i++) {
256305
snapshot = rotateSalts(rc, fraction, targetDate, i);
257306
}
258307

259-
// Step 3. Count how many emails are v2 vs v4 salts
260-
Map<String, Boolean> emailToV4TokenMap = new HashMap<>();
308+
Map<String, Boolean> postRotationEmailToSaltMap = new HashMap<>();
261309
for (String email : emails) {
262310
SaltEntry salt = getSalt(email, snapshot);
263311
boolean isV4 = salt.currentKeySalt() != null && salt.currentKeySalt().key() != null && salt.currentKeySalt().salt() != null;
264-
emailToV4TokenMap.put(email, isV4);
312+
postRotationEmailToSaltMap.put(email, isV4);
265313
}
266314

267315
// Step 4. Run /v2/token/refresh for all emails
316+
v2LoadSalts();
268317
Map<String, Boolean> emailToRefreshSuccessMap = new HashMap<>();
269-
for (String email : emails) {
318+
for (int i = 0; i < emails.size(); i++) {
319+
LOGGER.info("Step 4 - Token Refresh {}/{}", i + 1, emails.size());
320+
String email = emails.get(i);
321+
270322
try {
271323
JsonNode response = v2TokenRefresh(emailToRefreshTokenMap.get(email), emailToRefreshResponseKeyMap.get(email));
272324
emailToRefreshSuccessMap.put(email, response != null);
@@ -277,11 +329,14 @@ private void handleSaltSimulateToken(RoutingContext rc) {
277329
}
278330

279331
LOGGER.info(
280-
"UID token simulation: success_count={}, failure_count={}, v4_count={}, v2_count={}",
332+
"UID token simulation: success_count={}, failure_count={}, " +
333+
"v2_to_v4_count={}, v4_to_v4_count={}, v4_to_v2_count={}, v2_to_v2_count={}",
281334
emailToRefreshSuccessMap.values().stream().filter(x -> x).count(),
282335
emailToRefreshSuccessMap.values().stream().filter(x -> !x).count(),
283-
emailToV4TokenMap.values().stream().filter(x -> x).count(),
284-
emailToV4TokenMap.values().stream().filter(x -> !x).count());
336+
emails.stream().filter(email -> !preRotationEmailToSaltMap.get(email) && postRotationEmailToSaltMap.get(email)).count(),
337+
emails.stream().filter(email -> preRotationEmailToSaltMap.get(email) && postRotationEmailToSaltMap.get(email)).count(),
338+
emails.stream().filter(email -> preRotationEmailToSaltMap.get(email) && !postRotationEmailToSaltMap.get(email)).count(),
339+
emails.stream().filter(email -> !preRotationEmailToSaltMap.get(email) && !postRotationEmailToSaltMap.get(email)).count());
285340

286341
rc.response()
287342
.putHeader(HttpHeaders.CONTENT_TYPE, "application/json")
@@ -296,9 +351,9 @@ private void handleSaltSimulate(RoutingContext rc) {
296351
try {
297352
final double fraction = RequestUtil.getDouble(rc, "fraction").orElse(0.002740);
298353

299-
final int preMigrationIterations = RequestUtil.getDouble(rc, "preMigrationIterations").orElse(0.0).intValue();
300-
final int migrationV4Iterations = RequestUtil.getDouble(rc, "migrationV4Iterations").orElse(0.0).intValue();
301-
final int migrationV2V3Iterations = RequestUtil.getDouble(rc, "migrationV2V3Iterations").orElse(0.0).intValue();
354+
final int preMigrationIterations = RequestUtil.getDouble(rc, "pre_migration_iterations").orElse(0.0).intValue();
355+
final int migrationV4Iterations = RequestUtil.getDouble(rc, "migration_v4_iterations").orElse(0.0).intValue();
356+
final int migrationV2Iterations = RequestUtil.getDouble(rc, "migration_v2_iterations").orElse(0.0).intValue();
302357

303358
TargetDate targetDate =
304359
RequestUtil.getDate(rc, "target_date", DateTimeFormatter.ISO_LOCAL_DATE)
@@ -331,8 +386,8 @@ private void handleSaltSimulate(RoutingContext rc) {
331386
}
332387

333388
saltRotation.setEnableV4RawUid(false);
334-
for (int i = 0; i < migrationV2V3Iterations; i++) {
335-
LOGGER.info("Step 3 - Migration V2/V3 Iteration {}/{}", i + 1, migrationV2V3Iterations);
389+
for (int i = 0; i < migrationV2Iterations; i++) {
390+
LOGGER.info("Step 3 - Migration V2 Iteration {}/{}", i + 1, migrationV2Iterations);
336391
simulationIteration(rc, fraction, targetDate, preMigrationIterations + migrationV4Iterations + i, false, emails, emailToUidMapping);
337392
targetDate = targetDate.plusDays(1);
338393
}
@@ -723,6 +778,10 @@ private String toBase64String(byte[] b) {
723778
return Base64.getEncoder().encodeToString(b);
724779
}
725780

781+
private void v2LoadSalts() throws Exception {
782+
HTTP_CLIENT.post(String.format("%s/v2/salts/load", OPERATOR_URL), "", OPERATOR_HEADERS);
783+
}
784+
726785
private JsonNode v3IdentityMap(List<String> emails) throws Exception {
727786
StringBuilder reqBody = new StringBuilder("{ \"email\": [");
728787
for (String email : emails) {

0 commit comments

Comments
 (0)