Add registration ID to GetPreKeysResponse

This commit is contained in:
Katherine Yen
2025-12-09 18:06:01 -05:00
committed by ravi-signal
parent 7fd0cb32f6
commit 33f506a431
3 changed files with 45 additions and 30 deletions

View File

@@ -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 objects 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 objects 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)

View File

@@ -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;
}
/**

View File

@@ -483,16 +483,20 @@ class KeysGrpcServiceTest extends SimpleBaseGrpcTest<KeysGrpcService, KeysGrpc.K
final byte deviceId1 = 1;
final byte deviceId2 = 2;
final Map<Byte, Integer> deviceRegistrations = Map.of(
deviceId1, 123,
deviceId2, 456
);
for (final byte deviceId : List.of(deviceId1, deviceId2)) {
for (Map.Entry<Byte, Integer> entry : deviceRegistrations.entrySet()) {
final ECSignedPreKey ecSignedPreKey = KeysHelper.signedECPreKey(3, identityKeyPair);
final Optional<ECPreKey> 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<KeysGrpcService, KeysGrpc.K
.setKeyId(kemSignedPreKey.keyId())
.setPublicKey(ByteString.copyFrom(kemSignedPreKey.serializedPublicKey()))
.setSignature(ByteString.copyFrom(kemSignedPreKey.signature()))
.build());
.build())
.setRegistrationId(entry.getValue());
maybeEcPreKey.ifPresent(ecPreKey -> 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()));