Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ian UI d2 1728 identity buckets oom #492

Open
wants to merge 14 commits into
base: main
Choose a base branch
from
Open

Ian UI d2 1728 identity buckets oom #492

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
cb479e9
limit number of entries returned.
Ian-Nara Apr 15, 2024
a4eea10
return buckets in chunks of 30000
Ian-Nara Apr 15, 2024
758f094
Configurable identity buckets chunk size
Ian-Nara Apr 15, 2024
5dd3332
tests, refactoring, encryption
Ian-Nara Apr 16, 2024
d36c702
update v1 buckets
Ian-Nara Apr 16, 2024
8f51843
undo unnecessary whitespace change
Ian-Nara Apr 16, 2024
3962d98
undo unnecessary whitespace change
Ian-Nara Apr 16, 2024
0e69b2f
undo unnecessary whitespace change
Ian-Nara Apr 16, 2024
75fbfcf
undo unnecessary whitespace change
Ian-Nara Apr 16, 2024
8e52064
3 part stream
Ian-Nara Apr 17, 2024
400d7ad
refactoring
Ian-Nara Apr 17, 2024
4f0bbf7
rename interface
Ian-Nara Apr 17, 2024
a8e4506
fixed concurrency bug
Ian-Nara Apr 18, 2024
73b06e7
refactoring, strengthening tests
Ian-Nara Apr 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 104 additions & 18 deletions src/main/java/com/uid2/operator/vertx/UIDOperatorVerticle.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,11 @@
import com.uid2.operator.util.PrivacyBits;
import com.uid2.operator.util.Tuple;
import com.uid2.shared.Const.Data;
import com.uid2.shared.InstantClock;
import com.uid2.shared.Utils;
import com.uid2.shared.auth.*;
import com.uid2.shared.encryption.AesGcm;
import com.uid2.shared.encryption.Random;
import com.uid2.shared.health.HealthComponent;
import com.uid2.shared.health.HealthManager;
import com.uid2.shared.middleware.AuthMiddleware;
Expand Down Expand Up @@ -528,6 +530,12 @@ public void handleKeysRequest(RoutingContext rc) {
private String getSharingTokenExpirySeconds() {
return config.getString(Const.Config.SharingTokenExpiryProp);
}
private int getMaxIdentityBucketsResponseEntries() {
return config.getInteger(Const.Config.MaxIdentityBucketsResponseEntries, 1048576);
}
private int getIdentityBucketsResponseChunkSize() {
return config.getInteger(Const.Config.IdentityBucketsResponseChunkSize, 1048576);
}

public void handleKeysSharing(RoutingContext rc) {
try {
Expand Down Expand Up @@ -1088,23 +1096,31 @@ private void handleBucketsV1(RoutingContext rc) {
return;
}
final List<SaltEntry> modified = this.idService.getModifiedBuckets(sinceTimestamp);
final JsonArray resp = new JsonArray();
if (modified != null) {
for (SaltEntry e : modified) {
final JsonObject o = new JsonObject();
o.put("bucket_id", e.getHashedId());
Instant lastUpdated = Instant.ofEpochMilli(e.getLastUpdated());

o.put("last_updated", APIDateTimeFormatter.format(lastUpdated));
resp.add(o);
if (modified.size() > getMaxIdentityBucketsResponseEntries()) {
ResponseUtil.ClientError(rc, "provided since_timestamp produced large response. please provide a more recent since_timestamp or remap all with /identity/map");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Assuming this text will actually be returned to the caller, can we say something like:
"The time period requested resulted in too many results. Please provide a more recent since_timestamp or re-map all with /identity/map"

return;
}
ResponseUtil.Success(rc, resp);
transmitModifiedBucketsInChunks(rc, modified);
}
} else {
ResponseUtil.ClientError(rc, "missing parameter since_timestamp");
}
}

private String makeSaltEntriesString(List<SaltEntry> entries, int startIndex, int endIndexExclusive) {
StringBuilder s = new StringBuilder();
for(int i = startIndex; i < endIndexExclusive; i++) {
SaltEntry e = entries.get(i);
s.append("{\"bucket_id\":\"")
.append(e.getHashedId())
.append("\",\"last_updated\":\"")
.append(APIDateTimeFormatter.format(Instant.ofEpochMilli(e.getLastUpdated())))
.append("\"},");
}
return s.toString();
}

private void handleBucketsV2(RoutingContext rc) {
final JsonObject req = (JsonObject) rc.data().get("request");
final String qp = req.getString("since_timestamp");
Expand All @@ -1119,23 +1135,93 @@ private void handleBucketsV2(RoutingContext rc) {
return;
}
final List<SaltEntry> modified = this.idService.getModifiedBuckets(sinceTimestamp);
final JsonArray resp = new JsonArray();
if (modified != null) {
for (SaltEntry e : modified) {
final JsonObject o = new JsonObject();
o.put("bucket_id", e.getHashedId());
Instant lastUpdated = Instant.ofEpochMilli(e.getLastUpdated());

o.put("last_updated", APIDateTimeFormatter.format(lastUpdated));
resp.add(o);
if (modified.size() > getMaxIdentityBucketsResponseEntries()) {
ResponseUtil.ClientError(rc, "provided since_timestamp produced large response. please provide a more recent since_timestamp or remap all with /identity/map");
return;
}
try {
transmitModifiedBucketsInChunksEncrypted(rc, modified);
} catch (InvalidKeyException | InvalidAlgorithmParameterException | NoSuchAlgorithmException |
NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException e) {
ResponseUtil.Error(ResponseUtil.ResponseStatus.GenericError, 500, rc, "");
}
ResponseUtil.SuccessV2(rc, resp);
}
} else {
ResponseUtil.ClientError(rc, "missing parameter since_timestamp");
}
}

private void transmitModifiedBucketsInChunks(RoutingContext rc, List<SaltEntry> modified) {

HttpServerResponse response = rc.response();
int chunkSize = getIdentityBucketsResponseChunkSize();

response.setChunked(true).putHeader(HttpHeaders.CONTENT_TYPE, "application/json");
response.write("{\"body\":[");

for(int i =0; i < modified.size(); i+=chunkSize) {
String saltEntries = makeSaltEntriesString(modified, i, Math.min(i + chunkSize, modified.size()));
if(i + chunkSize >= modified.size()) {
saltEntries = saltEntries.substring(0, saltEntries.length() -1);
saltEntries += "], \"status\":\"success\"}";
response.write(saltEntries);
} else {
response.write(saltEntries);
}
}
rc.response().end();
}

private void transmitModifiedBucketsInChunksEncrypted(RoutingContext rc, List<SaltEntry> modified) throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
V2RequestUtil.V2Request request = V2RequestUtil.parseRequest(rc.body().asString(), AuthMiddleware.getAuthClient(ClientKey.class, rc), new InstantClock());
HttpServerResponse response = rc.response();

final String cipherScheme = "AES/GCM/NoPadding";
final int GCM_AUTHTAG_LENGTH = 16;
final int GCM_IV_LENGTH = 12;
final SecretKey k = new SecretKeySpec(request.encryptionKey, "AES");
final Cipher c = Cipher.getInstance(cipherScheme);
final byte[] ivBytes = Random.getBytes(GCM_IV_LENGTH);
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_AUTHTAG_LENGTH * 8, ivBytes);
c.init(Cipher.ENCRYPT_MODE, k, gcmParameterSpec);

int chunkSize = getIdentityBucketsResponseChunkSize();
response.setChunked(true).putHeader(HttpHeaders.CONTENT_TYPE, "application/json");

Buffer b = Buffer.buffer();
b.appendLong(EncodingUtils.NowUTCMillis().toEpochMilli());
b.appendBytes(request.nonce);
b.appendBytes("{\"body\":[".getBytes());
b = Buffer.buffer(ivBytes).appendBytes(c.update(b.getBytes()));

for(int i =0; i < modified.size(); i+=chunkSize) {

String saltEntries = makeSaltEntriesString(modified, i, Math.min(i + chunkSize, modified.size()));

if(i + chunkSize >= modified.size()) {
saltEntries = saltEntries.substring(0, saltEntries.length() -1);
saltEntries += "], \"status\":\"success\"}";
b.appendBytes(c.doFinal(saltEntries.getBytes()));
} else {
b.appendBytes(c.update(saltEntries.getBytes()));
}

if (b.length() % 3 == 0 || b.length() < 3 || (i+chunkSize >= modified.size())) {
response.write(Utils.toBase64String(b.getBytes()));
b = Buffer.buffer();
} else if ((b.length()-1) % 3 == 0) {
response.write(Utils.toBase64String(Arrays.copyOfRange(b.getBytes(), 0, b.length()-1)));
b = Buffer.buffer(Arrays.copyOfRange(b.getBytes(), b.length()-1, b.length()));
} else {
response.write(Utils.toBase64String(Arrays.copyOfRange(b.getBytes(), 0, b.length()-2)));
b = Buffer.buffer(Arrays.copyOfRange(b.getBytes(), b.length()-2, b.length()));
}

}
rc.response().end();
}

private void handleIdentityMapV1(RoutingContext rc) {
final InputUtil.InputVal input = this.phoneSupport ? this.getTokenInputV1(rc) : this.getTokenInput(rc);
if (this.phoneSupport ? !checkTokenInputV1(input, rc) : !checkTokenInput(input, rc)) {
Expand Down
6 changes: 5 additions & 1 deletion src/main/java/com/uid2/operator/vertx/V2PayloadHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,15 +170,19 @@ public void handleTokenRefresh(RoutingContext rc, Handler<RoutingContext> apiHan
private void passThrough(RoutingContext rc, Handler<RoutingContext> apiHandler) {
rc.data().put("request", rc.body().asJsonObject());
apiHandler.handle(rc);
if (rc.response().getStatusCode() != 200) {
if (rc.response().getStatusCode() != 200 || rc.response().ended()) {
return;
}
JsonObject respJson = (JsonObject) rc.data().get("response");
rc.response().putHeader(HttpHeaders.CONTENT_TYPE, "application/json")
.end(respJson.encode());

}

private void writeResponse(RoutingContext rc, byte[] nonce, JsonObject resp, byte[] keyBytes) {
if (rc.response().ended()) {
return;
}
Buffer buffer = Buffer.buffer();
buffer.appendLong(EncodingUtils.NowUTCMillis().toEpochMilli());
buffer.appendBytes(nonce);
Expand Down
63 changes: 63 additions & 0 deletions src/test/java/com/uid2/operator/UIDOperatorVerticleTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.time.temporal.ChronoUnit;
import java.util.*;
import java.util.stream.Collectors;
Expand Down Expand Up @@ -155,6 +157,9 @@ private void setupConfig(JsonObject config) {
config.put("client_side_token_generate_log_invalid_http_origins", true);

config.put(Const.Config.AllowClockSkewSecondsProp, 3600);

config.put(Const.Config.IdentityBucketsResponseChunkSize, 15);
config.put(Const.Config.MaxIdentityBucketsResponseEntries, 50);
}

private static byte[] makeAesKey(String prefix) {
Expand Down Expand Up @@ -212,6 +217,7 @@ private void send(String apiVersion, Vertx vertx, String endpoint, boolean isV1G
byte[] decrypted = AesGcm.decrypt(Utils.decodeBase64String(ar.result().bodyAsString()), 0, ck.getSecretBytes());
assertArrayEquals(Buffer.buffer().appendLong(nonce).getBytes(), Buffer.buffer(decrypted).slice(8, 16).getBytes());

System.out.println(new String(decrypted, 16, decrypted.length - 16, StandardCharsets.UTF_8));
JsonObject respJson = new JsonObject(new String(decrypted, 16, decrypted.length - 16, StandardCharsets.UTF_8));

handler.handle(respJson);
Expand Down Expand Up @@ -802,6 +808,63 @@ RefreshToken decodeRefreshToken(EncryptedTokenEncoder encoder, String refreshTok
return decodeRefreshToken(encoder, refreshTokenString, IdentityType.Email);
}

private void validateIdentityBuckets(JsonArray expectedList, List<SaltEntry> actualList) {
final DateTimeFormatter APIDateTimeFormatter = DateTimeFormatter.ISO_LOCAL_DATE_TIME.withZone(ZoneId.of("UTC"));

assertEquals(expectedList.size(), actualList.size());
for(int i = 0; i < expectedList.size(); i++) {
JsonObject expected = expectedList.getJsonObject(i);
SaltEntry actual = actualList.get(i);
assertAll("Salt Entry Matches",
() -> assertEquals(expected.getString("bucket_id"), actual.getHashedId()),
() -> assertEquals(expected.getString("last_updated"), APIDateTimeFormatter.format(Instant.ofEpochMilli(actual.getLastUpdated()))));
}
}

@ParameterizedTest
@CsvSource({"v1,1", "v2,1", "v1,11", "v2,11", "v1,15", "v2,15", "v1,20","v2,20", "v1,30", "v2,30", "v1,35", "v2,35", "v1,50", "v2,50"})
void identityBucketsChunked(String apiVersion, int numModifiedSalts, Vertx vertx, VertxTestContext testContext) {
final int clientSiteId = 201;
fakeAuth(clientSiteId, newClientCreationDateTime, Role.MAPPER);
List<SaltEntry> modifiedSalts = new ArrayList<>();
for(int i = 0; i < numModifiedSalts; i++) {
modifiedSalts.add(new SaltEntry(i, "someId" + i, i, "someSalt"));
}
when(saltProviderSnapshot.getModifiedSince(any())).thenReturn(modifiedSalts);

JsonObject req = new JsonObject();
req.put("since_timestamp", "2023-04-08T13:00:00");

send(apiVersion, vertx, apiVersion + "/identity/buckets", true, "since_timestamp=" +urlEncode("2023-04-08T13:00:00"), req, 200, respJson -> {
assertTrue(respJson.containsKey("body"));
assertFalse(respJson.containsKey("client_error"));
validateIdentityBuckets(respJson.getJsonArray("body"), modifiedSalts);
testContext.completeNow();
});
}

@ParameterizedTest
@ValueSource(strings = {"v1", "v2"})
void identityBucketsLimit(String apiVersion, Vertx vertx, VertxTestContext testContext) {
final int clientSiteId = 201;
fakeAuth(clientSiteId, newClientCreationDateTime, Role.MAPPER);
List<SaltEntry> modifiedSalts = new ArrayList<>();
for(int i = 0; i < 51; i++) {
modifiedSalts.add(new SaltEntry(i, "someId" + i, i, "someSalt"));
}
when(saltProviderSnapshot.getModifiedSince(any())).thenReturn(modifiedSalts);

JsonObject req = new JsonObject();
req.put("since_timestamp", "2023-04-08T13:00:00");

send(apiVersion, vertx, apiVersion + "/identity/buckets", true, "since_timestamp=" + urlEncode("2023-04-08T13:00:00"), req, 400, respJson -> {
assertFalse(respJson.containsKey("body"));
assertEquals("client_error", respJson.getString("status"));
assertEquals("provided since_timestamp produced large response. please provide a more recent since_timestamp or remap all with /identity/map", respJson.getString("message"));
testContext.completeNow();
});
}

@ParameterizedTest
@ValueSource(strings = {"v1", "v2"})
void identityMapNewClientNoPolicySpecified(String apiVersion, Vertx vertx, VertxTestContext testContext) {
Expand Down
Loading