Support KeyTransparencyClient with chat connection

RELEASE_NOTES.md

@@ -1,3 +1,4 @@
 v0.65.4
 
 - backup: Remove DirectStoryReplyMessage.storySentTimestamp
+- net: Enable using ChatConnection for key transparency operations (still Java only)

java/client/src/main/java/org/signal/libsignal/net/KeyTransparencyClient.java

@@ -18,8 +18,8 @@
  * Typed API to access the key transparency subsystem using an existing unauthenticated chat
  * connection.
  *
- * <p>Unlike {@link ChatService}, key transparency client does not export "raw" send/receive APIs,
- * and instead uses them internally to implement high-level operations.
+ * <p>Unlike {@link ChatService} and {@link ChatConnection}, key transparency client does not export
+ * "raw" send/receive APIs, and instead uses them internally to implement high-level operations.
  *
  * <p>Note: {@code Store} APIs may be invoked concurrently. Here are possible strategies to make
  * sure there are no thread safety violations:
@@ -31,11 +31,25 @@
  */
 public class KeyTransparencyClient {
   private final TokioAsyncContext tokioAsyncContext;
-  private final UnauthenticatedChatService chat;
+  private final UnauthenticatedChatService chatService;
+  private final UnauthenticatedChatConnection chatConnection;
+  private final Network.Environment environment;
 
   KeyTransparencyClient(UnauthenticatedChatService chat, TokioAsyncContext tokioAsyncContext) {
-    this.chat = chat;
+    this.chatService = chat;
+    this.chatConnection = null;
     this.tokioAsyncContext = tokioAsyncContext;
+    this.environment = chat.environment;
+  }
+
+  KeyTransparencyClient(
+      UnauthenticatedChatConnection chat,
+      TokioAsyncContext tokioAsyncContext,
+      Network.Environment environment) {
+    this.chatConnection = chat;
+    this.chatService = null;
+    this.tokioAsyncContext = tokioAsyncContext;
+    this.environment = environment;
   }
 
   /**
@@ -96,12 +110,14 @@ public CompletableFuture<SearchResult> search(
     // requests.
     // It may result in an IllegalArgumentException.
     try (NativeHandleGuard tokioContextGuard = this.tokioAsyncContext.guard();
-        NativeHandleGuard chatGuard = chat.guard();
         NativeHandleGuard identityKeyGuard = aciIdentityKey.getPublicKey().guard()) {
+      NativeHandleGuard chatServiceGuard = new NativeHandleGuard(chatService);
+      NativeHandleGuard chatConnectionGuard = new NativeHandleGuard(chatConnection);
       return Native.KeyTransparency_Search(
               tokioContextGuard.nativeHandle(),
-              chat.environment.value,
-              chatGuard.nativeHandle(),
+              this.environment.value,
+              chatServiceGuard.nativeHandle(),
+              chatConnectionGuard.nativeHandle(),
               aci.toServiceIdFixedWidthBinary(),
               identityKeyGuard.nativeHandle(),
               e164,
@@ -147,11 +163,13 @@ public CompletableFuture<SearchResult> search(
   public CompletableFuture<Void> updateDistinguished(final Store store) {
     byte[] lastDistinguished = store.getLastDistinguishedTreeHead().orElse(null);
     try (NativeHandleGuard tokioContextGuard = this.tokioAsyncContext.guard();
-        NativeHandleGuard chatGuard = chat.guard()) {
+        NativeHandleGuard chatServiceGuard = new NativeHandleGuard(chatService);
+        NativeHandleGuard chatConnectionGuard = new NativeHandleGuard(chatConnection)) {
       return Native.KeyTransparency_Distinguished(
               tokioContextGuard.nativeHandle(),
-              chat.environment.value,
-              chatGuard.nativeHandle(),
+              this.environment.value,
+              chatServiceGuard.nativeHandle(),
+              chatConnectionGuard.nativeHandle(),
               lastDistinguished)
           .thenApply(
               bytes -> {
@@ -205,11 +223,13 @@ public CompletableFuture<Void> monitor(
           .thenCompose((ignored) -> this.monitor(aci, e164, usernameHash, store));
     }
     try (NativeHandleGuard tokioContextGuard = this.tokioAsyncContext.guard();
-        NativeHandleGuard chatGuard = chat.guard()) {
+        NativeHandleGuard chatServiceGuard = new NativeHandleGuard(chatService);
+        NativeHandleGuard chatConnectionGuard = new NativeHandleGuard(chatConnection)) {
       return Native.KeyTransparency_Monitor(
               tokioContextGuard.nativeHandle(),
-              chat.environment.value,
-              chatGuard.nativeHandle(),
+              this.environment.value,
+              chatServiceGuard.nativeHandle(),
+              chatConnectionGuard.nativeHandle(),
               aci.toServiceIdFixedWidthBinary(),
               e164,
               usernameHash,

java/client/src/main/java/org/signal/libsignal/net/UnauthenticatedChatConnection.java

@@ -23,10 +23,14 @@ public class UnauthenticatedChatConnection extends ChatConnection {
   private UnauthenticatedChatConnection(
       final TokioAsyncContext tokioAsyncContext,
       long nativeHandle,
-      ChatConnectionListener listener) {
+      ChatConnectionListener listener,
+      Network.Environment ktEnvironment) {
     super(tokioAsyncContext, nativeHandle, listener);
+    this.keyTransparencyClient = new KeyTransparencyClient(this, tokioAsyncContext, ktEnvironment);
   }
 
+  private KeyTransparencyClient keyTransparencyClient;
+
   static CompletableFuture<UnauthenticatedChatConnection> connect(
       final TokioAsyncContext tokioAsyncContext,
       final Network.ConnectionManager connectionManager,
@@ -40,7 +44,20 @@ static CompletableFuture<UnauthenticatedChatConnection> connect(
                         .thenApply(
                             nativeHandle ->
                                 new UnauthenticatedChatConnection(
-                                    tokioAsyncContext, nativeHandle, chatListener))));
+                                    tokioAsyncContext,
+                                    nativeHandle,
+                                    chatListener,
+                                    connectionManager.environment()))));
+  }
+
+  /**
+   * High-level key transparency subsystem client on top using {@code this} to communicate with the
+   * chat server.
+   *
+   * @return an instance of {@link KeyTransparencyClient}
+   */
+  public KeyTransparencyClient keyTransparencyClient() {
+    return this.keyTransparencyClient;
   }
 
   // Implementing these abstract methods from ChatConnection allows UnauthenticatedChatConnection

java/client/src/test/java/org/signal/libsignal/net/KeyTransparencyClientTest.java

@@ -12,31 +12,66 @@
 import static org.signal.libsignal.net.KeyTransparencyTest.TEST_UNIDENTIFIED_ACCESS_KEY;
 import static org.signal.libsignal.net.KeyTransparencyTest.TEST_USERNAME_HASH;
 
+import java.util.ArrayList;
 import java.util.Deque;
+import java.util.function.Function;
 import org.junit.Assume;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+import org.signal.libsignal.internal.CompletableFuture;
 import org.signal.libsignal.keytrans.SearchResult;
 import org.signal.libsignal.keytrans.TestStore;
 import org.signal.libsignal.protocol.util.Hex;
 import org.signal.libsignal.util.TestEnvironment;
 
+@RunWith(Parameterized.class)
 public class KeyTransparencyClientTest {
   private static final String USER_AGENT = "test";
   private static final boolean INTEGRATION_TESTS_ENABLED =
       TestEnvironment.get("LIBSIGNAL_TESTING_RUN_NONHERMETIC_TESTS") != null;
 
+  @Parameters
+  public static Iterable<Function<Network, CompletableFuture<KeyTransparencyClient>>>
+      connectUnauthChatAndGetKtClientFns() {
+    ArrayList<Function<Network, CompletableFuture<KeyTransparencyClient>>> fns = new ArrayList<>(2);
+    fns.add(
+        (Network net) -> {
+          // Chat service
+          final UnauthenticatedChatService chat = net.createUnauthChatService(null);
+          return chat.connect()
+              .thenApply((ChatService.DebugInfo debugInfo) -> chat.keyTransparencyClient());
+        });
+    fns.add(
+        (Network net) -> {
+          // Chat connection
+          return net.connectUnauthChat(null)
+              .thenApply(
+                  (UnauthenticatedChatConnection chat) -> {
+                    chat.start();
+                    return chat.keyTransparencyClient();
+                  });
+        });
+    return fns;
+  }
+
+  @Parameter
+  public Function<Network, CompletableFuture<KeyTransparencyClient>>
+      connectUnauthChatAndGetKtClient;
+
   @Test
   public void searchInStagingIntegration() throws Exception {
     Assume.assumeTrue(INTEGRATION_TESTS_ENABLED);
 
     final Network net = new Network(Network.Environment.STAGING, USER_AGENT);
-    final UnauthenticatedChatService chat = net.createUnauthChatService(null);
-    chat.connect().get();
+    final KeyTransparencyClient ktClient = connectUnauthChatAndGetKtClient.apply(net).get();
 
     TestStore store = new TestStore();
 
     SearchResult result =
-        chat.keyTransparencyClient()
+        ktClient
             .search(
                 TEST_ACI,
                 TEST_ACI_IDENTITY_KEY,
@@ -62,11 +97,10 @@ public void updateDistinguishedStagingIntegration() throws Exception {
     Assume.assumeTrue(INTEGRATION_TESTS_ENABLED);
 
     final Network net = new Network(Network.Environment.STAGING, USER_AGENT);
-    final UnauthenticatedChatService chat = net.createUnauthChatService(null);
-    chat.connect().get();
+    final KeyTransparencyClient ktClient = connectUnauthChatAndGetKtClient.apply(net).get();
 
     TestStore store = new TestStore();
-    chat.keyTransparencyClient().updateDistinguished(store).get();
+    ktClient.updateDistinguished(store).get();
 
     assertTrue(store.getLastDistinguishedTreeHead().isPresent());
   }
@@ -76,13 +110,12 @@ public void monitorInStagingIntegration() throws Exception {
     Assume.assumeTrue(INTEGRATION_TESTS_ENABLED);
 
     final Network net = new Network(Network.Environment.STAGING, USER_AGENT);
-    final UnauthenticatedChatService chat = net.createUnauthChatService(null);
-    chat.connect().get();
+    final KeyTransparencyClient ktClient = connectUnauthChatAndGetKtClient.apply(net).get();
 
     TestStore store = new TestStore();
 
     SearchResult ignoredSearchResult =
-        chat.keyTransparencyClient()
+        ktClient
             .search(
                 TEST_ACI,
                 TEST_ACI_IDENTITY_KEY,
@@ -97,7 +130,7 @@ public void monitorInStagingIntegration() throws Exception {
     // Following search there should be a single entry in the account history
     assertEquals(1, accountDataHistory.size());
 
-    chat.keyTransparencyClient().monitor(TEST_ACI, TEST_E164, TEST_USERNAME_HASH, store).get();
+    ktClient.monitor(TEST_ACI, TEST_E164, TEST_USERNAME_HASH, store).get();
     // Another entry in the account history after a successful monitor request
     assertEquals(2, accountDataHistory.size());
   }

java/shared/java/org/signal/libsignal/internal/Native.java

@@ -376,10 +376,10 @@ private Native() {}
   public static native byte[] IncrementalMac_Update(long mac, byte[] bytes, int offset, int length);
 
   public static native byte[] KeyTransparency_AciSearchKey(byte[] aci);
-  public static native CompletableFuture<byte[]> KeyTransparency_Distinguished(long asyncRuntime, int environment, long chat, byte[] lastDistinguishedTreeHead);
+  public static native CompletableFuture<byte[]> KeyTransparency_Distinguished(long asyncRuntime, int environment, long chatService, long chatConnection, byte[] lastDistinguishedTreeHead);
   public static native byte[] KeyTransparency_E164SearchKey(String e164);
-  public static native CompletableFuture<byte[]> KeyTransparency_Monitor(long asyncRuntime, int environment, long chat, byte[] aci, String e164, byte[] usernameHash, byte[] accountData, byte[] lastDistinguishedTreeHead);
-  public static native CompletableFuture<Long> KeyTransparency_Search(long asyncRuntime, int environment, long chat, byte[] aci, long aciIdentityKey, String e164, byte[] unidentifiedAccessKey, byte[] usernameHash, byte[] accountData, byte[] lastDistinguishedTreeHead);
+  public static native CompletableFuture<byte[]> KeyTransparency_Monitor(long asyncRuntime, int environment, long chatService, long chatConnection, byte[] aci, String e164, byte[] usernameHash, byte[] accountData, byte[] lastDistinguishedTreeHead);
+  public static native CompletableFuture<Long> KeyTransparency_Search(long asyncRuntime, int environment, long chatService, long chatConnection, byte[] aci, long aciIdentityKey, String e164, byte[] unidentifiedAccessKey, byte[] usernameHash, byte[] accountData, byte[] lastDistinguishedTreeHead);
   public static native byte[] KeyTransparency_UsernameHashSearchKey(byte[] hash);
 
   public static native void KyberKeyPair_Destroy(long handle);

rust/bridge/shared/src/net/keytrans.rs

@@ -6,14 +6,16 @@
 use std::time::SystemTime;
 
 use libsignal_bridge_macros::{bridge_fn, bridge_io};
-use libsignal_bridge_types::net::chat::UnauthChat;
+use libsignal_bridge_types::net::chat::{UnauthChat, UnauthenticatedChatConnection};
 pub use libsignal_bridge_types::net::{Environment, TokioAsyncContext};
 use libsignal_bridge_types::support::AsType;
 use libsignal_core::{Aci, E164};
 use libsignal_keytrans::{
     AccountData, KeyTransparency, LocalStateUpdate, StoredAccountData, StoredTreeHead,
 };
-use libsignal_net::keytrans::{Error, Kt, SearchKey, SearchResult, UsernameHash};
+use libsignal_net::keytrans::{
+    Error, Kt, SearchKey, SearchResult, UnauthenticatedChat, UsernameHash,
+};
 use libsignal_protocol::PublicKey;
 use prost::{DecodeError, Message};
 
@@ -76,12 +78,27 @@ where
     T::decode(bytes.as_ref())
 }
 
+#[cfg(feature = "jni")]
+fn pick_chat_or_panic<'a>(
+    chat_service: Option<&'a UnauthChat>,
+    chat_connection: Option<&'a UnauthenticatedChatConnection>,
+) -> &'a (dyn UnauthenticatedChat + Sync) {
+    match (chat_service, chat_connection) {
+        (None, None) => panic!("no chat impl was provided"),
+        (None, Some(connection)) => connection,
+        (Some(service), None) => service,
+        (Some(_), Some(_)) => panic!("two chat impls were provided"),
+    }
+}
+
 #[bridge_io(TokioAsyncContext, node = false, ffi = false)]
 #[allow(clippy::too_many_arguments)]
 async fn KeyTransparency_Search(
     // TODO: it is currently possible to pass an env that does not match chat
     environment: AsType<Environment, u8>,
-    chat: &UnauthChat,
+    // TODO remove chatService when the switch to chat connection is complete.
+    chatService: Option<&UnauthChat>,
+    chatConnection: Option<&UnauthenticatedChatConnection>,
     aci: Aci,
     aci_identity_key: &PublicKey,
     e164: Option<E164>,
@@ -90,6 +107,7 @@ async fn KeyTransparency_Search(
     account_data: Option<Box<[u8]>>,
     last_distinguished_tree_head: Box<[u8]>,
 ) -> Result<SearchResult, Error> {
+    let chat = pick_chat_or_panic(chatService, chatConnection);
     let username_hash = username_hash.map(UsernameHash::from);
     let config = environment
         .into_inner()
@@ -99,7 +117,7 @@ async fn KeyTransparency_Search(
         .into();
     let kt = Kt {
         inner: KeyTransparency { config },
-        chat: &chat.service.0,
+        chat,
         config: Default::default(),
     };
 
@@ -149,13 +167,16 @@ async fn KeyTransparency_Search(
 async fn KeyTransparency_Monitor(
     // TODO: it is currently possible to pass an env that does not match chat
     environment: AsType<Environment, u8>,
-    chat: &UnauthChat,
+    // TODO remove chatService when the switch to chat connection is complete.
+    chatService: Option<&UnauthChat>,
+    chatConnection: Option<&UnauthenticatedChatConnection>,
     aci: Aci,
     e164: Option<E164>,
     username_hash: Option<Box<[u8]>>,
     account_data: Box<[u8]>,
     last_distinguished_tree_head: Box<[u8]>,
 ) -> Result<Vec<u8>, Error> {
+    let chat = pick_chat_or_panic(chatService, chatConnection);
     let username_hash = username_hash.map(UsernameHash::from);
 
     let account_data = {
@@ -177,7 +198,7 @@ async fn KeyTransparency_Monitor(
         .into();
     let kt = Kt {
         inner: KeyTransparency { config },
-        chat: &chat.service.0,
+        chat,
         config: Default::default(),
     };
     let updated_account_data = kt
@@ -196,9 +217,12 @@ async fn KeyTransparency_Monitor(
 async fn KeyTransparency_Distinguished(
     // TODO: it is currently possible to pass an env that does not match chat
     environment: AsType<Environment, u8>,
-    chat: &UnauthChat,
+    // TODO remove chatService when the switch to chat connection is complete.
+    chatService: Option<&UnauthChat>,
+    chatConnection: Option<&UnauthenticatedChatConnection>,
     last_distinguished_tree_head: Option<Box<[u8]>>,
 ) -> Result<Vec<u8>, Error> {
+    let chat = pick_chat_or_panic(chatService, chatConnection);
     let config = environment
         .into_inner()
         .env()
@@ -207,7 +231,7 @@ async fn KeyTransparency_Distinguished(
         .into();
     let kt = Kt {
         inner: KeyTransparency { config },
-        chat: &chat.service.0,
+        chat,
         config: Default::default(),
     };