From 33f506a43192764ae6217c6690004c93ba34e0eb Mon Sep 17 00:00:00 2001 From: Katherine Yen Date: Tue, 9 Dec 2025 18:06:01 -0500 Subject: [PATCH] Add registration ID to `GetPreKeysResponse` --- .../textsecuregcm/grpc/KeysGrpcHelper.java | 48 ++++++++++--------- .../src/main/proto/org/signal/chat/keys.proto | 5 ++ .../grpc/KeysGrpcServiceTest.java | 22 +++++---- 3 files changed, 45 insertions(+), 30 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java index 423c4a99d..71cf58ca7 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcHelper.java @@ -35,28 +35,32 @@ class KeysGrpcHelper { final String userAgent = RequestAttributesUtil.getUserAgent().orElse(null); return devices - .flatMap(device -> Mono - .fromFuture(keysManager.takeDevicePreKeys(device.getId(), targetServiceIdentifier, userAgent)) - .flatMap(Mono::justOrEmpty) - .map(devicePreKeys -> { - final GetPreKeysResponse.PreKeyBundle.Builder builder = GetPreKeysResponse.PreKeyBundle.newBuilder() - .setEcSignedPreKey(EcSignedPreKey.newBuilder() - .setKeyId(devicePreKeys.ecSignedPreKey().keyId()) - .setPublicKey(ByteString.copyFrom(devicePreKeys.ecSignedPreKey().serializedPublicKey())) - .setSignature(ByteString.copyFrom(devicePreKeys.ecSignedPreKey().signature())) - .build()) - .setKemOneTimePreKey(KemSignedPreKey.newBuilder() - .setKeyId(devicePreKeys.kemSignedPreKey().keyId()) - .setPublicKey(ByteString.copyFrom(devicePreKeys.kemSignedPreKey().serializedPublicKey())) - .setSignature(ByteString.copyFrom(devicePreKeys.kemSignedPreKey().signature())) - .build()); - devicePreKeys.ecPreKey().ifPresent(ecPreKey -> builder.setEcOneTimePreKey(EcPreKey.newBuilder() - .setKeyId(ecPreKey.keyId()) - .setPublicKey(ByteString.copyFrom(ecPreKey.serializedPublicKey())) - .build())); - // Cast device IDs to `int` to match data types in the response object’s protobuf definition - return Tuples.of((int) device.getId(), builder.build()); - })) + .flatMap(device -> { + final int registrationId = device.getRegistrationId(targetServiceIdentifier.identityType()); + return Mono + .fromFuture(keysManager.takeDevicePreKeys(device.getId(), targetServiceIdentifier, userAgent)) + .flatMap(Mono::justOrEmpty) + .map(devicePreKeys -> { + final GetPreKeysResponse.PreKeyBundle.Builder builder = GetPreKeysResponse.PreKeyBundle.newBuilder() + .setEcSignedPreKey(EcSignedPreKey.newBuilder() + .setKeyId(devicePreKeys.ecSignedPreKey().keyId()) + .setPublicKey(ByteString.copyFrom(devicePreKeys.ecSignedPreKey().serializedPublicKey())) + .setSignature(ByteString.copyFrom(devicePreKeys.ecSignedPreKey().signature())) + .build()) + .setKemOneTimePreKey(KemSignedPreKey.newBuilder() + .setKeyId(devicePreKeys.kemSignedPreKey().keyId()) + .setPublicKey(ByteString.copyFrom(devicePreKeys.kemSignedPreKey().serializedPublicKey())) + .setSignature(ByteString.copyFrom(devicePreKeys.kemSignedPreKey().signature())) + .build()) + .setRegistrationId(registrationId); + devicePreKeys.ecPreKey().ifPresent(ecPreKey -> builder.setEcOneTimePreKey(EcPreKey.newBuilder() + .setKeyId(ecPreKey.keyId()) + .setPublicKey(ByteString.copyFrom(ecPreKey.serializedPublicKey())) + .build())); + // Cast device IDs to `int` to match data types in the response object’s protobuf definition + return Tuples.of((int) device.getId(), builder.build()); + }); + }) // If there were no devices with valid prekey bundles in the account, the account is gone .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException())) .collectMap(Tuple2::getT1, Tuple2::getT2) diff --git a/service/src/main/proto/org/signal/chat/keys.proto b/service/src/main/proto/org/signal/chat/keys.proto index 368bd432e..80c5e46f6 100644 --- a/service/src/main/proto/org/signal/chat/keys.proto +++ b/service/src/main/proto/org/signal/chat/keys.proto @@ -198,6 +198,11 @@ message GetPreKeysResponse { * account/device/identity. */ common.KemSignedPreKey kem_one_time_pre_key = 3; + + /** + * The registration ID for the targeted account/device/identity. + */ + uint32 registration_id = 4; } /** diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java index bf1bf0d54..c358dd845 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java @@ -483,16 +483,20 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest deviceRegistrations = Map.of( + deviceId1, 123, + deviceId2, 456 + ); - for (final byte deviceId : List.of(deviceId1, deviceId2)) { + for (Map.Entry entry : deviceRegistrations.entrySet()) { final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(3, identityKeyPair); final Optional maybeEcPreKey = Optional .of(new ECPreKey(1, ECKeyPair.generate().getPublicKey())) - .filter(_ -> deviceId == deviceId1); + .filter(_ -> entry.getKey() == deviceId1); final KEMSignedPreKey kemSignedPreKey = KeysHelper.signedKEMPreKey(2, identityKeyPair); - devicePreKeysMap.put(deviceId, new KeysManager.DevicePreKeys(ecSignedPreKey, maybeEcPreKey, kemSignedPreKey)); + devicePreKeysMap.put(entry.getKey(), new KeysManager.DevicePreKeys(ecSignedPreKey, maybeEcPreKey, kemSignedPreKey)); final GetPreKeysResponse.PreKeyBundle.Builder builder = GetPreKeysResponse.PreKeyBundle.newBuilder() .setEcSignedPreKey(EcSignedPreKey.newBuilder() @@ -504,19 +508,21 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest builder .setEcOneTimePreKey(EcPreKey.newBuilder() .setKeyId(ecPreKey.keyId()) .setPublicKey(ByteString.copyFrom(ecPreKey.serializedPublicKey())) .build())); - expectedPreKeyBundles.put(deviceId, builder.build()); + expectedPreKeyBundles.put(entry.getKey(), builder.build()); final Device device = mock(Device.class); - when(device.getId()).thenReturn(deviceId); + when(device.getId()).thenReturn(entry.getKey()); + when(device.getRegistrationId(any())).thenReturn(entry.getValue()); - devices.put(deviceId, device); - when(targetAccount.getDevice(deviceId)).thenReturn(Optional.of(device)); + devices.put(entry.getKey(), device); + when(targetAccount.getDevice(entry.getKey())).thenReturn(Optional.of(device)); } when(targetAccount.getDevices()).thenReturn(new ArrayList<>(devices.values()));