Require some tree head sizes in key transparency search and monitor r…

…equests
signalapp · Oct 31, 2024 · e627d4e · e627d4e
1 parent 9b5a62e
commit e627d4e
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 62 deletions.
diff --git a/...rc/main/java/org/whispersystems/textsecuregcm/entities/KeyTransparencyMonitorRequest.java b/...rc/main/java/org/whispersystems/textsecuregcm/entities/KeyTransparencyMonitorRequest.java
@@ -36,12 +36,10 @@ public record KeyTransparencyMonitorRequest(
     Optional<@Valid UsernameHashMonitor> usernameHash,
 
     @Schema(description = "The tree head size to prove consistency against.")
-    @NotNull
-    Optional<@Positive Long> lastNonDistinguishedTreeHeadSize,
+    @Positive long lastNonDistinguishedTreeHeadSize,
 
     @Schema(description = "The distinguished tree head size to prove consistency against.")
-    @NotNull
-    Optional<@Positive Long> lastDistinguishedTreeHeadSize
+    @Positive long lastDistinguishedTreeHeadSize
 ) {
 
   public record AciMonitor(

diff --git a/...src/main/java/org/whispersystems/textsecuregcm/entities/KeyTransparencySearchRequest.java b/...src/main/java/org/whispersystems/textsecuregcm/entities/KeyTransparencySearchRequest.java
@@ -51,8 +51,9 @@ public record KeyTransparencySearchRequest(
     @Schema(description = "The non-distinguished tree head size to prove consistency against.")
     Optional<@Positive Long> lastTreeHeadSize,
 
+    @NotNull
     @Schema(description = "The distinguished tree head size to prove consistency against.")
-    Optional<@Positive Long> distinguishedTreeHeadSize
+    @Positive long distinguishedTreeHeadSize
 ) {
     @AssertTrue
     public boolean isUnidentifiedAccessKeyProvidedWithE164() {

diff --git a/...n/java/org/whispersystems/textsecuregcm/keytransparency/KeyTransparencyServiceClient.java b/...n/java/org/whispersystems/textsecuregcm/keytransparency/KeyTransparencyServiceClient.java
@@ -119,7 +119,7 @@ public CompletableFuture<SearchResponse> search(
       final Optional<ByteString> usernameHash,
       final Optional<E164SearchRequest> e164SearchRequest,
       final Optional<Long> lastTreeHeadSize,
-      final Optional<Long> distinguishedTreeHeadSize,
+      final long distinguishedTreeHeadSize,
       final Duration timeout) {
     final SearchRequest.Builder searchKeysRequestBuilder = SearchRequest.newBuilder()
         .setAci(aci)
@@ -128,27 +128,27 @@ public CompletableFuture<SearchResponse> search(
     usernameHash.ifPresent(searchKeysRequestBuilder::setUsernameHash);
     e164SearchRequest.ifPresent(searchKeysRequestBuilder::setE164SearchRequest);
 
-    final ConsistencyParameters.Builder consistency = ConsistencyParameters.newBuilder();
+    final ConsistencyParameters.Builder consistency = ConsistencyParameters.newBuilder()
+        .setDistinguished(distinguishedTreeHeadSize);
     lastTreeHeadSize.ifPresent(consistency::setLast);
-    distinguishedTreeHeadSize.ifPresent(consistency::setDistinguished);
 
     searchKeysRequestBuilder.setConsistency(consistency.build());
 
     return CompletableFutureUtil.toCompletableFuture(stub.withDeadline(toDeadline(timeout))
         .search(searchKeysRequestBuilder.build()), callbackExecutor);
   }
 
-  @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
   public CompletableFuture<byte[]> monitor(final List<MonitorKey> monitorKeys,
-      final Optional<Long> lastTreeHeadSize,
-      final Optional<Long> distinguishedTreeHeadSize,
+      final long lastTreeHeadSize,
+      final long distinguishedTreeHeadSize,
       final Duration timeout) {
     final MonitorRequest.Builder monitorRequestBuilder = MonitorRequest.newBuilder()
         .addAllContactKeys(monitorKeys);
 
-    final ConsistencyParameters.Builder consistency = ConsistencyParameters.newBuilder();
-    lastTreeHeadSize.ifPresent(consistency::setLast);
-    distinguishedTreeHeadSize.ifPresent(consistency::setDistinguished);
+    final ConsistencyParameters consistency = ConsistencyParameters.newBuilder()
+        .setLast(lastTreeHeadSize)
+        .setDistinguished(distinguishedTreeHeadSize)
+        .build();
 
     monitorRequestBuilder.setConsistency(consistency);