diff --git a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java index 7a976f36..8998a8be 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java +++ b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java @@ -23,6 +23,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.whispersystems.textsecuregcm.entities.PreKey; import org.whispersystems.textsecuregcm.entities.PreKeyList; +import org.whispersystems.textsecuregcm.entities.PreKeyStatus; import org.whispersystems.textsecuregcm.entities.UnstructuredPreKeyList; import org.whispersystems.textsecuregcm.federation.FederatedClientManager; import org.whispersystems.textsecuregcm.federation.NoSuchPeerException; @@ -73,6 +74,20 @@ public class KeysController { keys.store(account.getNumber(), device.getId(), preKeys.getKeys(), preKeys.getLastResortKey()); } + @Timed + @GET + @Path("/") + @Produces(MediaType.APPLICATION_JSON) + public PreKeyStatus getStatus(@Auth Account account) { + int count = keys.getCount(account.getNumber(), account.getAuthenticatedDevice().get().getId()); + + if (count > 0) { + count = count - 1; + } + + return new PreKeyStatus(count); + } + @Timed @GET @Path("/{number}/{device_id}") diff --git a/src/main/java/org/whispersystems/textsecuregcm/entities/PreKeyStatus.java b/src/main/java/org/whispersystems/textsecuregcm/entities/PreKeyStatus.java new file mode 100644 index 00000000..0da6d20a --- /dev/null +++ b/src/main/java/org/whispersystems/textsecuregcm/entities/PreKeyStatus.java @@ -0,0 +1,20 @@ +package org.whispersystems.textsecuregcm.entities; + + +import com.fasterxml.jackson.annotation.JsonProperty; + +public class PreKeyStatus { + + @JsonProperty + private int count; + + public PreKeyStatus(int count) { + this.count = count; + } + + public PreKeyStatus() {} + + public int getCount() { + return count; + } +} diff --git a/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java b/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java index 10be7b3a..10f68d21 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java +++ b/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java @@ -40,7 +40,6 @@ import java.lang.annotation.RetentionPolicy; import java.lang.annotation.Target; import java.sql.ResultSet; import java.sql.SQLException; -import java.util.LinkedList; import java.util.List; public abstract class Keys { @@ -67,6 +66,9 @@ public abstract class Keys { @Mapper(PreKeyMapper.class) abstract List retrieveFirst(@Bind("number") String number); + @SqlQuery("SELECT COUNT(*) FROM keys WHERE number = :number AND device_id = :device_id") + public abstract int getCount(@Bind("number") String number, @Bind("device_id") long deviceId); + @Transaction(TransactionIsolationLevel.SERIALIZABLE) public void store(String number, long deviceId, List keys, PreKey lastResortKey) { for (PreKey key : keys) { 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 3e389f5f..0ec02e55 100644 --- a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java +++ b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java @@ -6,6 +6,7 @@ import com.yammer.dropwizard.testing.ResourceTest; import org.junit.Test; import org.whispersystems.textsecuregcm.controllers.KeysController; import org.whispersystems.textsecuregcm.entities.PreKey; +import org.whispersystems.textsecuregcm.entities.PreKeyStatus; import org.whispersystems.textsecuregcm.entities.UnstructuredPreKeyList; import org.whispersystems.textsecuregcm.limits.RateLimiter; import org.whispersystems.textsecuregcm.limits.RateLimiters; @@ -71,11 +72,25 @@ public class KeyControllerTest extends ResourceTest { allKeys.add(SAMPLE_KEY); allKeys.add(SAMPLE_KEY2); allKeys.add(SAMPLE_KEY3); + when(keys.get(EXISTS_NUMBER)).thenReturn(Optional.of(new UnstructuredPreKeyList(allKeys))); + when(keys.getCount(eq(AuthHelper.VALID_NUMBER), eq(1L))).thenReturn(5); addResource(new KeysController(rateLimiters, keys, accounts, null)); } + @Test + public void validKeyStatusTest() throws Exception { + PreKeyStatus result = client().resource("/v1/keys") + .header("Authorization", + AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD)) + .get(PreKeyStatus.class); + + assertThat(result.getCount() == 4); + + verify(keys).getCount(eq(AuthHelper.VALID_NUMBER), eq(1L)); + } + @Test public void validLegacyRequestTest() throws Exception { PreKey result = client().resource(String.format("/v1/keys/%s", EXISTS_NUMBER))