diff --git a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java index 11fb384ad..a13f3cf81 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java +++ b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java @@ -174,6 +174,7 @@ public class KeysController { if (device.isPresent() && device.get().isActive()) { preKey.setRegistrationId(device.get().getRegistrationId()); + preKey.setIdentityKey(destination.getIdentityKey()); filteredKeys.add(preKey); } } diff --git a/src/main/java/org/whispersystems/textsecuregcm/entities/PreKey.java b/src/main/java/org/whispersystems/textsecuregcm/entities/PreKey.java index 981ba027f..777e23ff9 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/entities/PreKey.java +++ b/src/main/java/org/whispersystems/textsecuregcm/entities/PreKey.java @@ -20,6 +20,7 @@ package org.whispersystems.textsecuregcm.entities; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.annotations.VisibleForTesting; import javax.validation.constraints.NotNull; import javax.xml.bind.annotation.XmlTransient; @@ -58,8 +59,19 @@ public class PreKey { public PreKey() {} public PreKey(long id, String number, long deviceId, long keyId, - String publicKey, String identityKey, - boolean lastResort) + String publicKey, boolean lastResort) + { + this.id = id; + this.number = number; + this.deviceId = deviceId; + this.keyId = keyId; + this.publicKey = publicKey; + this.lastResort = lastResort; + } + + @VisibleForTesting + public PreKey(long id, String number, long deviceId, long keyId, + String publicKey, String identityKey, boolean lastResort) { this.id = id; this.number = number; diff --git a/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java b/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java index 10f68d216..361d750fe 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java +++ b/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java @@ -50,12 +50,12 @@ public abstract class Keys { @SqlUpdate("DELETE FROM keys WHERE id = :id") abstract void removeKey(@Bind("id") long id); - @SqlBatch("INSERT INTO keys (number, device_id, key_id, public_key, identity_key, last_resort) VALUES " + - "(:number, :device_id, :key_id, :public_key, :identity_key, :last_resort)") + @SqlBatch("INSERT INTO keys (number, device_id, key_id, public_key, last_resort) VALUES " + + "(:number, :device_id, :key_id, :public_key, :last_resort)") abstract void append(@PreKeyBinder List preKeys); - @SqlUpdate("INSERT INTO keys (number, device_id, key_id, public_key, identity_key, last_resort) VALUES " + - "(:number, :device_id, :key_id, :public_key, :identity_key, :last_resort)") + @SqlUpdate("INSERT INTO keys (number, device_id, key_id, public_key, last_resort) VALUES " + + "(:number, :device_id, :key_id, :public_key, :last_resort)") abstract void append(@PreKeyBinder PreKey preKey); @SqlQuery("SELECT * FROM keys WHERE number = :number AND device_id = :device_id ORDER BY key_id ASC FOR UPDATE") @@ -129,7 +129,6 @@ public abstract class Keys { sql.bind("device_id", preKey.getDeviceId()); sql.bind("key_id", preKey.getKeyId()); sql.bind("public_key", preKey.getPublicKey()); - sql.bind("identity_key", preKey.getIdentityKey()); sql.bind("last_resort", preKey.isLastResort() ? 1 : 0); } }; @@ -145,7 +144,6 @@ public abstract class Keys { { return new PreKey(resultSet.getLong("id"), resultSet.getString("number"), resultSet.getLong("device_id"), resultSet.getLong("key_id"), resultSet.getString("public_key"), - resultSet.getString("identity_key"), resultSet.getInt("last_resort") == 1); } } diff --git a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java index c3fe4b3b4..c40d5def0 100644 --- a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java +++ b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java @@ -5,6 +5,8 @@ import com.sun.jersey.api.client.ClientResponse; import com.yammer.dropwizard.testing.ResourceTest; import org.junit.Test; import org.mockito.ArgumentCaptor; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; import org.whispersystems.textsecuregcm.controllers.KeysController; import org.whispersystems.textsecuregcm.entities.PreKey; import org.whispersystems.textsecuregcm.entities.PreKeyList; @@ -34,7 +36,7 @@ public class KeyControllerTest extends ResourceTest { private final int SAMPLE_REGISTRATION_ID2 = 1002; private final PreKey SAMPLE_KEY = new PreKey(1, EXISTS_NUMBER, Device.MASTER_ID, 1234, "test1", "test2", false); - private final PreKey SAMPLE_KEY2 = new PreKey(2, EXISTS_NUMBER, 2, 5667, "test3", "test4", false ); + private final PreKey SAMPLE_KEY2 = new PreKey(2, EXISTS_NUMBER, 2, 5667, "test3", "test4,", false ); private final PreKey SAMPLE_KEY3 = new PreKey(3, EXISTS_NUMBER, 3, 334, "test5", "test6", false ); private final Keys keys = mock(Keys.class ); private final AccountsManager accounts = mock(AccountsManager.class); @@ -63,21 +65,34 @@ public class KeyControllerTest extends ResourceTest { when(existsAccount.getDevice(2L)).thenReturn(Optional.of(sampleDevice2)); when(existsAccount.getDevice(3L)).thenReturn(Optional.of(sampleDevice3)); when(existsAccount.isActive()).thenReturn(true); + when(existsAccount.getIdentityKey()).thenReturn("existsidentitykey"); when(accounts.get(EXISTS_NUMBER)).thenReturn(Optional.of(existsAccount)); when(accounts.get(NOT_EXISTS_NUMBER)).thenReturn(Optional.absent()); when(rateLimiters.getPreKeysLimiter()).thenReturn(rateLimiter); - when(keys.get(eq(EXISTS_NUMBER), eq(1L))).thenReturn(Optional.of(new UnstructuredPreKeyList(SAMPLE_KEY))); + when(keys.get(eq(EXISTS_NUMBER), eq(1L))).thenAnswer(new Answer>() { + @Override + public Optional answer(InvocationOnMock invocationOnMock) throws Throwable { + return Optional.of(new UnstructuredPreKeyList(cloneKey(SAMPLE_KEY))); + } + }); + when(keys.get(eq(NOT_EXISTS_NUMBER), eq(1L))).thenReturn(Optional.absent()); - List allKeys = new LinkedList<>(); - allKeys.add(SAMPLE_KEY); - allKeys.add(SAMPLE_KEY2); - allKeys.add(SAMPLE_KEY3); + when(keys.get(EXISTS_NUMBER)).thenAnswer(new Answer>() { + @Override + public Optional answer(InvocationOnMock invocationOnMock) throws Throwable { + List allKeys = new LinkedList<>(); + allKeys.add(cloneKey(SAMPLE_KEY)); + allKeys.add(cloneKey(SAMPLE_KEY2)); + allKeys.add(cloneKey(SAMPLE_KEY3)); + + return Optional.of(new UnstructuredPreKeyList(allKeys)); + } + }); - when(keys.get(EXISTS_NUMBER)).thenReturn(Optional.of(new UnstructuredPreKeyList(allKeys))); when(keys.getCount(eq(AuthHelper.VALID_NUMBER), eq(1L))).thenReturn(5); when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null); @@ -105,7 +120,7 @@ public class KeyControllerTest extends ResourceTest { assertThat(result.getKeyId()).isEqualTo(SAMPLE_KEY.getKeyId()); assertThat(result.getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); - assertThat(result.getIdentityKey()).isEqualTo(SAMPLE_KEY.getIdentityKey()); + assertThat(result.getIdentityKey()).isEqualTo(existsAccount.getIdentityKey()); assertThat(result.getId() == 0); assertThat(result.getNumber() == null); @@ -126,7 +141,7 @@ public class KeyControllerTest extends ResourceTest { assertThat(result.getKeyId()).isEqualTo(SAMPLE_KEY.getKeyId()); assertThat(result.getPublicKey()).isEqualTo(SAMPLE_KEY.getPublicKey()); - assertThat(result.getIdentityKey()).isEqualTo(SAMPLE_KEY.getIdentityKey()); + assertThat(result.getIdentityKey()).isEqualTo(existsAccount.getIdentityKey()); assertThat(result.getRegistrationId()).isEqualTo(SAMPLE_REGISTRATION_ID); assertThat(result.getId() == 0); @@ -135,7 +150,7 @@ public class KeyControllerTest extends ResourceTest { result = results.getKeys().get(1); assertThat(result.getKeyId()).isEqualTo(SAMPLE_KEY2.getKeyId()); assertThat(result.getPublicKey()).isEqualTo(SAMPLE_KEY2.getPublicKey()); - assertThat(result.getIdentityKey()).isEqualTo(SAMPLE_KEY2.getIdentityKey()); + assertThat(result.getIdentityKey()).isEqualTo(existsAccount.getIdentityKey()); assertThat(result.getRegistrationId()).isEqualTo(SAMPLE_REGISTRATION_ID2); assertThat(result.getId() == 0); @@ -209,4 +224,9 @@ public class KeyControllerTest extends ResourceTest { verify(accounts).update(AuthHelper.VALID_ACCOUNT); } + private PreKey cloneKey(PreKey source) { + return new PreKey(source.getId(), source.getNumber(), source.getDeviceId(), source.getKeyId(), + source.getPublicKey(), source.getIdentityKey(), source.isLastResort()); + } + } \ No newline at end of file