Bump version to 0.13

Step 3 in migration.

.gitignore

@@ -7,3 +7,5 @@ run.sh
 local.yml
 config/production.yml
 config/federated.yml
+config/staging.yml
+.opsmanage

pom.xml

@@ -9,7 +9,7 @@
 
     <groupId>org.whispersystems.textsecure</groupId>
     <artifactId>TextSecureServer</artifactId>
-    <version>0.4</version>
+    <version>0.13</version>
 
     <dependencies>
         <dependency>

src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java

@@ -20,11 +20,11 @@ import com.fasterxml.jackson.annotation.JsonProperty;
 import com.yammer.dropwizard.config.Configuration;
 import com.yammer.dropwizard.db.DatabaseConfiguration;
 import org.whispersystems.textsecuregcm.configuration.ApnConfiguration;
-import org.whispersystems.textsecuregcm.configuration.DataDogConfiguration;
 import org.whispersystems.textsecuregcm.configuration.FederationConfiguration;
 import org.whispersystems.textsecuregcm.configuration.GcmConfiguration;
 import org.whispersystems.textsecuregcm.configuration.GraphiteConfiguration;
 import org.whispersystems.textsecuregcm.configuration.MemcacheConfiguration;
+import org.whispersystems.textsecuregcm.configuration.MetricsConfiguration;
 import org.whispersystems.textsecuregcm.configuration.NexmoConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RedisConfiguration;
@@ -87,7 +87,7 @@ public class WhisperServerConfiguration extends Configuration {
 
   @Valid
   @JsonProperty
-  private DataDogConfiguration datadog = new DataDogConfiguration();
+  private MetricsConfiguration metrics = new MetricsConfiguration();
 
   @Valid
   @JsonProperty
@@ -141,7 +141,7 @@ public class WhisperServerConfiguration extends Configuration {
     return graphite;
   }
 
-  public DataDogConfiguration getDataDogConfiguration() {
-    return datadog;
+  public MetricsConfiguration getMetricsConfiguration() {
+    return metrics;
   }
 }

src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -24,9 +24,7 @@ import com.yammer.dropwizard.config.HttpConfiguration;
 import com.yammer.dropwizard.db.DatabaseConfiguration;
 import com.yammer.dropwizard.jdbi.DBIFactory;
 import com.yammer.dropwizard.migrations.MigrationsBundle;
-import com.yammer.metrics.core.Clock;
-import com.yammer.metrics.core.MetricPredicate;
-import com.yammer.metrics.reporting.DatadogReporter;
+import com.yammer.metrics.Metrics;
 import com.yammer.metrics.reporting.GraphiteReporter;
 import net.spy.memcached.MemcachedClient;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
@@ -48,6 +46,11 @@ import org.whispersystems.textsecuregcm.federation.FederatedPeer;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.IOExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
+import org.whispersystems.textsecuregcm.metrics.CpuUsageGauge;
+import org.whispersystems.textsecuregcm.metrics.FreeMemoryGauge;
+import org.whispersystems.textsecuregcm.metrics.JsonMetricsReporter;
+import org.whispersystems.textsecuregcm.metrics.NetworkReceivedGauge;
+import org.whispersystems.textsecuregcm.metrics.NetworkSentGauge;
 import org.whispersystems.textsecuregcm.providers.MemcacheHealthCheck;
 import org.whispersystems.textsecuregcm.providers.MemcachedClientFactory;
 import org.whispersystems.textsecuregcm.providers.RedisClientFactory;
@@ -162,17 +165,22 @@ public class WhisperServerService extends Service<WhisperServerConfiguration> {
     environment.addProvider(new IOExceptionMapper());
     environment.addProvider(new RateLimitExceededExceptionMapper());
 
+    Metrics.newGauge(CpuUsageGauge.class, "cpu", new CpuUsageGauge());
+    Metrics.newGauge(FreeMemoryGauge.class, "free_memory", new FreeMemoryGauge());
+    Metrics.newGauge(NetworkSentGauge.class, "bytes_sent", new NetworkSentGauge());
+    Metrics.newGauge(NetworkReceivedGauge.class, "bytes_received", new NetworkReceivedGauge());
+
     if (config.getGraphiteConfiguration().isEnabled()) {
       GraphiteReporter.enable(15, TimeUnit.SECONDS,
                               config.getGraphiteConfiguration().getHost(),
                               config.getGraphiteConfiguration().getPort());
     }
 
-    if (config.getDataDogConfiguration().isEnabled()) {
-      new DatadogReporter.Builder().withApiKey(config.getDataDogConfiguration().getApiKey())
-                                   .withVmMetricsEnabled(true)
-                                   .build()
-                                   .start(15, TimeUnit.SECONDS);
+    if (config.getMetricsConfiguration().isEnabled()) {
+      new JsonMetricsReporter("textsecure", Metrics.defaultRegistry(),
+                              config.getMetricsConfiguration().getToken(),
+                              config.getMetricsConfiguration().getHost())
+          .start(60, TimeUnit.SECONDS);
     }
   }
 

src/main/java/org/whispersystems/textsecuregcm/configuration/MetricsConfiguration.java

@@ -0,0 +1,27 @@
+package org.whispersystems.textsecuregcm.configuration;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class MetricsConfiguration {
+
+  @JsonProperty
+  private String token;
+
+  @JsonProperty
+  private String host;
+
+  @JsonProperty
+  private boolean enabled = false;
+
+  public String getHost() {
+    return host;
+  }
+
+  public String getToken() {
+    return token;
+  }
+
+  public boolean isEnabled() {
+    return enabled && token != null && host != null;
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java

@@ -135,6 +135,10 @@ public class AccountController {
         throw new WebApplicationException(Response.status(403).build());
       }
 
+      if (accounts.isRelayListed(number)) {
+        throw new WebApplicationException(Response.status(417).build());
+      }
+
       Device device = new Device();
       device.setId(Device.MASTER_ID);
       device.setAuthenticationCredentials(new AuthenticationCredentials(password));

src/main/java/org/whispersystems/textsecuregcm/controllers/DirectoryController.java

@@ -18,7 +18,9 @@ package org.whispersystems.textsecuregcm.controllers;
 
 import com.google.common.base.Optional;
 import com.yammer.dropwizard.auth.Auth;
+import com.yammer.metrics.Metrics;
 import com.yammer.metrics.annotation.Timed;
+import com.yammer.metrics.core.Histogram;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.ClientContact;
@@ -46,7 +48,8 @@ import java.util.List;
 @Path("/v1/directory")
 public class DirectoryController {
 
-  private final Logger logger = LoggerFactory.getLogger(DirectoryController.class);
+  private final Logger    logger            = LoggerFactory.getLogger(DirectoryController.class);
+  private final Histogram contactsHistogram = Metrics.newHistogram(DirectoryController.class, "contacts");
 
   private final RateLimiters     rateLimiters;
   private final DirectoryManager directory;
@@ -56,7 +59,7 @@ public class DirectoryController {
     this.rateLimiters = rateLimiters;
   }
 
-  @Timed()
+  @Timed
   @GET
   @Path("/{token}")
   @Produces(MediaType.APPLICATION_JSON)
@@ -77,7 +80,7 @@ public class DirectoryController {
     }
   }
 
-  @Timed()
+  @Timed
   @PUT
   @Path("/tokens")
   @Produces(MediaType.APPLICATION_JSON)
@@ -86,6 +89,7 @@ public class DirectoryController {
       throws RateLimitExceededException
   {
     rateLimiters.getContactsLimiter().validate(account.getNumber(), contacts.getContacts().size());
+    contactsHistogram.update(contacts.getContacts().size());
 
     try {
       List<byte[]> tokens = new LinkedList<>();

src/main/java/org/whispersystems/textsecuregcm/controllers/FederationController.java

@@ -76,7 +76,7 @@ public class FederationController {
                                               @PathParam("attachmentId") long attachmentId)
       throws IOException
   {
-    return attachmentController.redirectToAttachment(new NonLimitedAccount("Unknown", peer.getName()),
+    return attachmentController.redirectToAttachment(new NonLimitedAccount("Unknown", -1, peer.getName()),
                                                      attachmentId, Optional.<String>absent());
   }
 
@@ -89,7 +89,7 @@ public class FederationController {
       throws IOException
   {
     try {
-      return keysController.get(new NonLimitedAccount("Unknown", peer.getName()), number, Optional.<String>absent());
+      return keysController.get(new NonLimitedAccount("Unknown", -1, peer.getName()), number, Optional.<String>absent());
     } catch (RateLimitExceededException e) {
       logger.warn("Rate limiting on federated channel", e);
       throw new IOException(e);
@@ -106,7 +106,7 @@ public class FederationController {
       throws IOException
   {
     try {
-      return keysController.getDeviceKey(new NonLimitedAccount("Unknown", peer.getName()),
+      return keysController.getDeviceKey(new NonLimitedAccount("Unknown", -1, peer.getName()),
                                          number, device, Optional.<String>absent());
     } catch (RateLimitExceededException e) {
       logger.warn("Rate limiting on federated channel", e);
@@ -116,16 +116,17 @@ public class FederationController {
 
   @Timed
   @PUT
-  @Path("/messages/{source}/{destination}")
-  public void sendMessages(@Auth                     FederatedPeer peer,
-                           @PathParam("source")      String source,
-                           @PathParam("destination") String destination,
-                           @Valid                    IncomingMessageList messages)
+  @Path("/messages/{source}/{sourceDeviceId}/{destination}")
+  public void sendMessages(@Auth                        FederatedPeer peer,
+                           @PathParam("source")         String source,
+                           @PathParam("sourceDeviceId") long sourceDeviceId,
+                           @PathParam("destination")    String destination,
+                           @Valid                       IncomingMessageList messages)
       throws IOException
   {
     try {
       messages.setRelay(null);
-      messageController.sendMessage(new NonLimitedAccount(source, peer.getName()), destination, messages);
+      messageController.sendMessage(new NonLimitedAccount(source, sourceDeviceId, peer.getName()), destination, messages);
     } catch (RateLimitExceededException e) {
       logger.warn("Rate limiting on federated channel", e);
       throw new IOException(e);

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;
@@ -69,10 +70,31 @@ public class KeysController {
   @PUT
   @Consumes(MediaType.APPLICATION_JSON)
   public void setKeys(@Auth Account account, @Valid PreKeyList preKeys)  {
-    Device device = account.getAuthenticatedDevice().get();
+    Device device      = account.getAuthenticatedDevice().get();
+    String identityKey = preKeys.getLastResortKey().getIdentityKey();
+
+    if (!identityKey.equals(account.getIdentityKey())) {
+      account.setIdentityKey(identityKey);
+      accounts.update(account);
+    }
+
     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}")
@@ -152,6 +174,7 @@ public class KeysController {
 
       if (device.isPresent() && device.get().isActive()) {
         preKey.setRegistrationId(device.get().getRegistrationId());
+        preKey.setIdentityKey(destination.getIdentityKey());
         filteredKeys.add(preKey);
       }
     }

src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java

@@ -186,7 +186,8 @@ public class MessageController {
   {
     try {
       FederatedClient client = federatedClientManager.getClient(messages.getRelay());
-      client.sendMessages(source.getNumber(), destinationName, messages);
+      client.sendMessages(source.getNumber(), source.getAuthenticatedDevice().get().getId(),
+                          destinationName, messages);
     } catch (NoSuchPeerException e) {
       throw new NoSuchUserException(e);
     }

src/main/java/org/whispersystems/textsecuregcm/entities/ClientContactTokens.java

@@ -30,4 +30,11 @@ public class ClientContactTokens {
   public List<String> getContacts() {
     return contacts;
   }
+
+  public ClientContactTokens() {}
+
+  public ClientContactTokens(List<String> contacts) {
+    this.contacts = contacts;
+  }
+
 }

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;

src/main/java/org/whispersystems/textsecuregcm/entities/PreKeyList.java

@@ -17,6 +17,7 @@
 package org.whispersystems.textsecuregcm.entities;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.annotations.VisibleForTesting;
 import org.hibernate.validator.constraints.NotEmpty;
 
 import javax.validation.Valid;
@@ -27,6 +28,7 @@ public class PreKeyList {
 
   @JsonProperty
   @NotNull
+  @Valid
   private PreKey lastResortKey;
 
   @JsonProperty
@@ -38,7 +40,17 @@ public class PreKeyList {
     return keys;
   }
 
+  @VisibleForTesting
+  public void setKeys(List<PreKey> keys) {
+    this.keys = keys;
+  }
+
   public PreKey getLastResortKey() {
     return lastResortKey;
   }
+
+  @VisibleForTesting
+  public void setLastResortKey(PreKey lastResortKey) {
+    this.lastResortKey = lastResortKey;
+  }
 }

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;
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/federation/FederatedClient.java

@@ -64,7 +64,7 @@ public class FederatedClient {
 
   private static final String USER_COUNT_PATH     = "/v1/federation/user_count";
   private static final String USER_TOKENS_PATH    = "/v1/federation/user_tokens/%d";
-  private static final String RELAY_MESSAGE_PATH  = "/v1/federation/messages/%s/%s";
+  private static final String RELAY_MESSAGE_PATH  = "/v1/federation/messages/%s/%d/%s";
   private static final String PREKEY_PATH_DEVICE  = "/v1/federation/key/%s/%s";
   private static final String ATTACHMENT_URI_PATH = "/v1/federation/attachment/%d";
 
@@ -155,11 +155,11 @@ public class FederatedClient {
     }
   }
 
-  public void sendMessages(String source, String destination, IncomingMessageList messages)
+  public void sendMessages(String source, long sourceDeviceId, String destination, IncomingMessageList messages)
       throws IOException
   {
     try {
-      WebResource    resource = client.resource(peer.getUrl()).path(String.format(RELAY_MESSAGE_PATH, source, destination));
+      WebResource    resource = client.resource(peer.getUrl()).path(String.format(RELAY_MESSAGE_PATH, source, sourceDeviceId, destination));
       ClientResponse response = resource.type(MediaType.APPLICATION_JSON)
                                         .header("Authorization", authorizationHeader)
                                         .entity(messages)

src/main/java/org/whispersystems/textsecuregcm/federation/NonLimitedAccount.java

@@ -4,6 +4,7 @@ package org.whispersystems.textsecuregcm.federation;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.google.common.base.Optional;
 import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Device;
 
 public class NonLimitedAccount extends Account {
 
@@ -13,20 +14,32 @@ public class NonLimitedAccount extends Account {
   @JsonIgnore
   private final String relay;
 
-  public NonLimitedAccount(String number, String relay) {
-    this.number = number;
-    this.relay  = relay;
+  @JsonIgnore
+  private final long deviceId;
+
+  public NonLimitedAccount(String number, long deviceId, String relay) {
+    this.number   = number;
+    this.deviceId = deviceId;
+    this.relay    = relay;
   }
 
+  @Override
   public String getNumber() {
     return number;
   }
 
+  @Override
   public boolean isRateLimited() {
     return false;
   }
 
+  @Override
   public Optional<String> getRelay() {
     return Optional.of(relay);
   }
+
+  @Override
+  public Optional<Device> getAuthenticatedDevice() {
+    return Optional.of(new Device(deviceId, null, null, null, null, null, false, 0));
+  }
 }

src/main/java/org/whispersystems/textsecuregcm/metrics/CpuUsageGauge.java

@@ -0,0 +1,16 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+import com.sun.management.OperatingSystemMXBean;
+import com.yammer.metrics.core.Gauge;
+
+import java.lang.management.ManagementFactory;
+
+public class CpuUsageGauge extends Gauge<Integer> {
+  @Override
+  public Integer value() {
+    OperatingSystemMXBean mbean = (com.sun.management.OperatingSystemMXBean)
+        ManagementFactory.getOperatingSystemMXBean();
+
+    return (int) Math.ceil(mbean.getSystemCpuLoad() * 100);
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/metrics/FreeMemoryGauge.java

@@ -0,0 +1,20 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+import com.sun.management.OperatingSystemMXBean;
+import com.yammer.metrics.core.Gauge;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.management.ManagementFactory;
+
+public class FreeMemoryGauge extends Gauge<Long> {
+
+  @Override
+  public Long value() {
+    OperatingSystemMXBean mbean = (com.sun.management.OperatingSystemMXBean)
+        ManagementFactory.getOperatingSystemMXBean();
+
+    return mbean.getFreePhysicalMemorySize();
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/metrics/JsonMetricsReporter.java

@@ -0,0 +1,324 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+import com.fasterxml.jackson.core.JsonEncoding;
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.yammer.metrics.core.Clock;
+import com.yammer.metrics.core.Counter;
+import com.yammer.metrics.core.Gauge;
+import com.yammer.metrics.core.Histogram;
+import com.yammer.metrics.core.Metered;
+import com.yammer.metrics.core.Metric;
+import com.yammer.metrics.core.MetricName;
+import com.yammer.metrics.core.MetricProcessor;
+import com.yammer.metrics.core.MetricsRegistry;
+import com.yammer.metrics.core.Sampling;
+import com.yammer.metrics.core.Summarizable;
+import com.yammer.metrics.core.Timer;
+import com.yammer.metrics.core.VirtualMachineMetrics;
+import com.yammer.metrics.reporting.AbstractPollingReporter;
+import com.yammer.metrics.stats.Snapshot;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.net.HttpURLConnection;
+import java.net.InetAddress;
+import java.net.URL;
+import java.net.UnknownHostException;
+import java.util.Map;
+import java.util.SortedMap;
+import java.util.concurrent.TimeUnit;
+import java.util.regex.Pattern;
+
+/**
+ * Adapted from MetricsServlet.
+ */
+public class JsonMetricsReporter extends AbstractPollingReporter implements MetricProcessor<JsonMetricsReporter.Context> {
+  private final Clock clock = Clock.defaultClock();
+  private final VirtualMachineMetrics vm = VirtualMachineMetrics.getInstance();
+  private final String service;
+  private final MetricsRegistry registry;
+  private final JsonFactory factory = new JsonFactory();
+
+  private final String table;
+  private final String sunnylabsHost;
+  private final String host;
+
+  private final boolean includeVMMetrics;
+
+  public JsonMetricsReporter(String service, MetricsRegistry registry, String token, String sunnylabsHost) throws UnknownHostException {
+    this(service, registry, token, sunnylabsHost, true);
+  }
+
+  public JsonMetricsReporter(String service, MetricsRegistry registry, String token, String sunnylabsHost, boolean includeVMMetrics) throws UnknownHostException {
+    super(registry, "jsonmetrics-reporter");
+    this.service = service;
+    this.registry = registry;
+    this.table = token;
+    this.sunnylabsHost = sunnylabsHost;
+    this.host = InetAddress.getLocalHost().getHostName();
+    this.includeVMMetrics = includeVMMetrics;
+  }
+
+  @Override
+  public void run() {
+    try {
+      URL http = new URL("https", sunnylabsHost, 443, "/report/metrics?t=" + table + "&h=" + host);
+      System.out.println("Reporting started to: " + http);
+      HttpURLConnection urlc = (HttpURLConnection) http.openConnection();
+      urlc.setDoOutput(true);
+      urlc.addRequestProperty("Content-Type", "application/json");
+      OutputStream outputStream = urlc.getOutputStream();
+      writeJson(outputStream);
+      outputStream.close();
+      System.out.println("Reporting complete: " + urlc.getResponseCode());
+    } catch (IOException e) {
+      e.printStackTrace();
+    }
+  }
+
+  static final class Context {
+    final boolean showFullSamples;
+    final JsonGenerator json;
+
+    Context(JsonGenerator json, boolean showFullSamples) {
+      this.json = json;
+      this.showFullSamples = showFullSamples;
+    }
+  }
+
+  public void writeJson(OutputStream out) throws IOException {
+    final JsonGenerator json = factory.createGenerator(out, JsonEncoding.UTF8);
+    json.writeStartObject();
+    if (includeVMMetrics) {
+      writeVmMetrics(json);
+    }
+    writeRegularMetrics(json, false);
+    json.writeEndObject();
+    json.close();
+  }
+
+  private void writeVmMetrics(JsonGenerator json) throws IOException {
+    json.writeFieldName(service);
+    json.writeStartObject();
+    json.writeFieldName("jvm");
+    json.writeStartObject();
+    {
+      json.writeFieldName("vm");
+      json.writeStartObject();
+      {
+        json.writeStringField("name", vm.name());
+        json.writeStringField("version", vm.version());
+      }
+      json.writeEndObject();
+
+      json.writeFieldName("memory");
+      json.writeStartObject();
+      {
+        json.writeNumberField("totalInit", vm.totalInit());
+        json.writeNumberField("totalUsed", vm.totalUsed());
+        json.writeNumberField("totalMax", vm.totalMax());
+        json.writeNumberField("totalCommitted", vm.totalCommitted());
+
+        json.writeNumberField("heapInit", vm.heapInit());
+        json.writeNumberField("heapUsed", vm.heapUsed());
+        json.writeNumberField("heapMax", vm.heapMax());
+        json.writeNumberField("heapCommitted", vm.heapCommitted());
+
+        json.writeNumberField("heap_usage", vm.heapUsage());
+        json.writeNumberField("non_heap_usage", vm.nonHeapUsage());
+        json.writeFieldName("memory_pool_usages");
+        json.writeStartObject();
+        {
+          for (Map.Entry<String, Double> pool : vm.memoryPoolUsage().entrySet()) {
+            json.writeNumberField(pool.getKey(), pool.getValue());
+          }
+        }
+        json.writeEndObject();
+      }
+      json.writeEndObject();
+
+      final Map<String, VirtualMachineMetrics.BufferPoolStats> bufferPoolStats = vm.getBufferPoolStats();
+      if (!bufferPoolStats.isEmpty()) {
+        json.writeFieldName("buffers");
+        json.writeStartObject();
+        {
+          json.writeFieldName("direct");
+          json.writeStartObject();
+          {
+            json.writeNumberField("count", bufferPoolStats.get("direct").getCount());
+            json.writeNumberField("memoryUsed", bufferPoolStats.get("direct").getMemoryUsed());
+            json.writeNumberField("totalCapacity", bufferPoolStats.get("direct").getTotalCapacity());
+          }
+          json.writeEndObject();
+
+          json.writeFieldName("mapped");
+          json.writeStartObject();
+          {
+            json.writeNumberField("count", bufferPoolStats.get("mapped").getCount());
+            json.writeNumberField("memoryUsed", bufferPoolStats.get("mapped").getMemoryUsed());
+            json.writeNumberField("totalCapacity", bufferPoolStats.get("mapped").getTotalCapacity());
+          }
+          json.writeEndObject();
+        }
+        json.writeEndObject();
+      }
+
+
+      json.writeNumberField("daemon_thread_count", vm.daemonThreadCount());
+      json.writeNumberField("thread_count", vm.threadCount());
+      json.writeNumberField("current_time", clock.time());
+      json.writeNumberField("uptime", vm.uptime());
+      json.writeNumberField("fd_usage", vm.fileDescriptorUsage());
+
+      json.writeFieldName("thread-states");
+      json.writeStartObject();
+      {
+        for (Map.Entry<Thread.State, Double> entry : vm.threadStatePercentages()
+                                                       .entrySet()) {
+          json.writeNumberField(entry.getKey().toString().toLowerCase(),
+                                entry.getValue());
+        }
+      }
+      json.writeEndObject();
+
+      json.writeFieldName("garbage-collectors");
+      json.writeStartObject();
+      {
+        for (Map.Entry<String, VirtualMachineMetrics.GarbageCollectorStats> entry : vm.garbageCollectors()
+                                                                                      .entrySet()) {
+          json.writeFieldName(entry.getKey());
+          json.writeStartObject();
+          {
+            final VirtualMachineMetrics.GarbageCollectorStats gc = entry.getValue();
+            json.writeNumberField("runs", gc.getRuns());
+            json.writeNumberField("time", gc.getTime(TimeUnit.MILLISECONDS));
+          }
+          json.writeEndObject();
+        }
+      }
+      json.writeEndObject();
+    }
+    json.writeEndObject();
+    json.writeEndObject();
+  }
+
+  public void writeRegularMetrics(JsonGenerator json, boolean showFullSamples) throws IOException {
+    for (Map.Entry<String, SortedMap<MetricName, Metric>> entry : registry.groupedMetrics().entrySet()) {
+      for (Map.Entry<MetricName, Metric> subEntry : entry.getValue().entrySet()) {
+        json.writeFieldName(sanitize(subEntry.getKey()));
+        try {
+          subEntry.getValue()
+                  .processWith(this,
+                               subEntry.getKey(),
+                               new Context(json, showFullSamples));
+        } catch (Exception e) {
+          e.printStackTrace();
+        }
+      }
+    }
+  }
+
+  @Override
+  public void processHistogram(MetricName name, Histogram histogram, Context context) throws Exception {
+    final JsonGenerator json = context.json;
+    json.writeStartObject();
+    {
+      json.writeNumberField("count", histogram.count());
+      writeSummarizable(histogram, json);
+      writeSampling(histogram, json);
+      if (context.showFullSamples) {
+        json.writeObjectField("values", histogram.getSnapshot().getValues());
+      }
+      histogram.clear();
+    }
+    json.writeEndObject();
+  }
+
+  @Override
+  public void processCounter(MetricName name, Counter counter, Context context) throws Exception {
+    final JsonGenerator json = context.json;
+    json.writeNumber(counter.count());
+  }
+
+  @Override
+  public void processGauge(MetricName name, Gauge<?> gauge, Context context) throws Exception {
+    final JsonGenerator json = context.json;
+    json.writeObject(evaluateGauge(gauge));
+  }
+
+  @Override
+  public void processMeter(MetricName name, Metered meter, Context context) throws Exception {
+    final JsonGenerator json = context.json;
+    json.writeStartObject();
+    {
+      writeMeteredFields(meter, json);
+    }
+    json.writeEndObject();
+  }
+
+  @Override
+  public void processTimer(MetricName name, Timer timer, Context context) throws Exception {
+    final JsonGenerator json = context.json;
+    json.writeStartObject();
+    {
+      json.writeFieldName("duration");
+      json.writeStartObject();
+      {
+        json.writeStringField("unit", timer.durationUnit().toString().toLowerCase());
+        writeSummarizable(timer, json);
+        writeSampling(timer, json);
+        if (context.showFullSamples) {
+          json.writeObjectField("values", timer.getSnapshot().getValues());
+        }
+      }
+      json.writeEndObject();
+
+      json.writeFieldName("rate");
+      json.writeStartObject();
+      {
+        writeMeteredFields(timer, json);
+      }
+      json.writeEndObject();
+    }
+    json.writeEndObject();
+  }
+
+  private static Object evaluateGauge(Gauge<?> gauge) {
+    try {
+      return gauge.value();
+    } catch (RuntimeException e) {
+      return "error reading gauge: " + e.getMessage();
+    }
+  }
+
+  private static void writeSummarizable(Summarizable metric, JsonGenerator json) throws IOException {
+    json.writeNumberField("min", metric.min());
+    json.writeNumberField("max", metric.max());
+    json.writeNumberField("mean", metric.mean());
+  }
+
+  private static void writeSampling(Sampling metric, JsonGenerator json) throws IOException {
+    final Snapshot snapshot = metric.getSnapshot();
+    json.writeNumberField("median", snapshot.getMedian());
+    json.writeNumberField("p75", snapshot.get75thPercentile());
+    json.writeNumberField("p95", snapshot.get95thPercentile());
+    json.writeNumberField("p99", snapshot.get99thPercentile());
+    json.writeNumberField("p999", snapshot.get999thPercentile());
+  }
+
+  private static void writeMeteredFields(Metered metered, JsonGenerator json) throws IOException {
+    json.writeNumberField("count", metered.count());
+    json.writeNumberField("mean", metered.meanRate());
+    json.writeNumberField("m1", metered.oneMinuteRate());
+    json.writeNumberField("m5", metered.fiveMinuteRate());
+    json.writeNumberField("m15", metered.fifteenMinuteRate());
+  }
+
+  private static final Pattern SIMPLE_NAMES = Pattern.compile("[^a-zA-Z0-9_.\\-~]");
+
+  private String sanitize(MetricName metricName) {
+    return SIMPLE_NAMES.matcher(metricName.getGroup() + "." + metricName.getName()).replaceAll("_");
+  }
+
+}

src/main/java/org/whispersystems/textsecuregcm/metrics/NetworkGauge.java

@@ -0,0 +1,37 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+
+import com.yammer.metrics.core.Gauge;
+import org.whispersystems.textsecuregcm.util.Pair;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.FileReader;
+import java.io.IOException;
+
+public abstract class NetworkGauge extends Gauge<Long> {
+
+  protected Pair<Long, Long> getSentReceived() throws IOException {
+    File           proc          = new File("/proc/net/dev");
+    BufferedReader reader        = new BufferedReader(new FileReader(proc));
+    String         header        = reader.readLine();
+    String         header2       = reader.readLine();
+
+    long           bytesSent     = 0;
+    long           bytesReceived = 0;
+
+    String interfaceStats;
+
+      while ((interfaceStats = reader.readLine()) != null) {
+        String[] stats = interfaceStats.split("\\s+");
+
+        if (!stats[1].equals("lo:")) {
+          bytesReceived += Long.parseLong(stats[2]);
+          bytesSent     += Long.parseLong(stats[10]);
+        }
+      }
+
+    return new Pair<>(bytesSent, bytesReceived);
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/metrics/NetworkReceivedGauge.java

@@ -0,0 +1,36 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.util.Pair;
+
+import java.io.IOException;
+
+public class NetworkReceivedGauge extends NetworkGauge {
+
+  private final Logger logger = LoggerFactory.getLogger(NetworkSentGauge.class);
+
+  private long lastTimestamp;
+  private long lastReceived;
+
+  @Override
+  public Long value() {
+    try {
+      long             timestamp       = System.currentTimeMillis();
+      Pair<Long, Long> sentAndReceived = getSentReceived();
+      long             result          = 0;
+
+      if (lastTimestamp != 0) {
+        result       = sentAndReceived.second() - lastReceived;
+        lastReceived = sentAndReceived.second();
+      }
+
+      lastTimestamp = timestamp;
+      return result;
+    } catch (IOException e) {
+      logger.warn("NetworkReceivedGauge", e);
+      return -1L;
+    }
+  }
+
+}

src/main/java/org/whispersystems/textsecuregcm/metrics/NetworkSentGauge.java

@@ -0,0 +1,35 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.util.Pair;
+
+import java.io.IOException;
+
+public class NetworkSentGauge extends NetworkGauge {
+
+  private final Logger logger = LoggerFactory.getLogger(NetworkSentGauge.class);
+
+  private long lastTimestamp;
+  private long lastSent;
+
+  @Override
+  public Long value() {
+    try {
+      long             timestamp       = System.currentTimeMillis();
+      Pair<Long, Long> sentAndReceived = getSentReceived();
+      long             result          = 0;
+
+      if (lastTimestamp != 0) {
+        result        = sentAndReceived.first() - lastSent;
+        lastSent      = sentAndReceived.first();
+      }
+
+      lastTimestamp = timestamp;
+      return result;
+    } catch (IOException e) {
+      logger.warn("NetworkSentGauge", e);
+      return -1L;
+    }
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/storage/Account.java

@@ -28,7 +28,7 @@ import java.util.List;
 
 public class Account implements Serializable {
 
-  public static final int MEMCACHE_VERION = 2;
+  public static final int MEMCACHE_VERION = 3;
 
   @JsonIgnore
   private long id;
@@ -42,16 +42,14 @@ public class Account implements Serializable {
   @JsonProperty
   private List<Device> devices = new LinkedList<>();
 
+  @JsonProperty
+  private String identityKey;
+
   @JsonIgnore
   private Optional<Device> authenticatedDevice;
 
   public Account() {}
 
-  public Account(String number, boolean supportsSms) {
-    this.number      = number;
-    this.supportsSms = supportsSms;
-  }
-
   @VisibleForTesting
   public Account(String number, boolean supportsSms, List<Device> devices) {
     this.number      = number;
@@ -142,4 +140,12 @@ public class Account implements Serializable {
   public Optional<String> getRelay() {
     return Optional.absent();
   }
+
+  public void setIdentityKey(String identityKey) {
+    this.identityKey = identityKey;
+  }
+
+  public String getIdentityKey() {
+    return identityKey;
+  }
 }

src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java

@@ -90,6 +90,13 @@ public class AccountsManager {
     else                 return Optional.absent();
   }
 
+  public boolean isRelayListed(String number) {
+    byte[]                  token   = Util.getContactToken(number);
+    Optional<ClientContact> contact = directory.get(token);
+
+    return contact.isPresent() && !Util.isEmpty(contact.get().getRelay());
+  }
+
   private void updateDirectory(Account account) {
     if (account.isActive()) {
       byte[]        token         = Util.getContactToken(account.getNumber());

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 {
@@ -51,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<PreKey> 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")
@@ -67,6 +66,9 @@ public abstract class Keys {
   @Mapper(PreKeyMapper.class)
   abstract List<PreKey> 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<PreKey> keys, PreKey lastResortKey) {
     for (PreKey key : keys) {
@@ -127,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);
           }
         };
@@ -143,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);
     }
   }

src/main/java/org/whispersystems/textsecuregcm/workers/DirectoryCommand.java

@@ -26,6 +26,8 @@ import com.yammer.dropwizard.jdbi.args.OptionalArgumentFactory;
 import net.sourceforge.argparse4j.inf.Namespace;
 import net.spy.memcached.MemcachedClient;
 import org.skife.jdbi.v2.DBI;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.federation.FederatedClientManager;
 import org.whispersystems.textsecuregcm.providers.MemcachedClientFactory;
@@ -38,6 +40,8 @@ import redis.clients.jedis.JedisPool;
 
 public class DirectoryCommand extends ConfiguredCommand<WhisperServerConfiguration> {
 
+  private final Logger logger = LoggerFactory.getLogger(DirectoryCommand.class);
+
   public DirectoryCommand() {
     super("directory", "Update directory from DB and peers.");
   }
@@ -68,6 +72,9 @@ public class DirectoryCommand extends ConfiguredCommand<WhisperServerConfigurati
 
       update.updateFromLocalDatabase();
       update.updateFromPeers();
+    } catch (Exception ex) {
+      logger.warn("Directory Exception", ex);
+      throw new RuntimeException(ex);
     } finally {
       Thread.sleep(3000);
       System.exit(0);

src/main/java/org/whispersystems/textsecuregcm/workers/DirectoryUpdater.java

@@ -27,6 +27,7 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.DirectoryManager;
 import org.whispersystems.textsecuregcm.storage.DirectoryManager.BatchOperationHandle;
 import org.whispersystems.textsecuregcm.util.Base64;
+import org.whispersystems.textsecuregcm.util.Hex;
 import org.whispersystems.textsecuregcm.util.Util;
 
 import java.util.Iterator;
@@ -85,7 +86,7 @@ public class DirectoryUpdater {
 
     for (FederatedClient client : clients) {
       logger.info("Updating directory from peer: " + client.getPeerName());
-      BatchOperationHandle handle = directory.startBatchOperation();
+//      BatchOperationHandle handle = directory.startBatchOperation();
 
       try {
         int userCount = client.getUserCount();
@@ -94,31 +95,38 @@ public class DirectoryUpdater {
         logger.info("Remote peer user count: " + userCount);
 
         while (retrieved < userCount) {
+          logger.info("Retrieving remote tokens...");
           List<ClientContact> clientContacts = client.getUserTokens(retrieved);
 
-          if (clientContacts == null)
+          if (clientContacts == null) {
+            logger.info("Remote tokens empty, ending...");
             break;
+          } else {
+            logger.info("Retrieved " + clientContacts.size() + " remote tokens...");
+          }
 
           for (ClientContact clientContact : clientContacts) {
             clientContact.setRelay(client.getPeerName());
 
             Optional<ClientContact> existing = directory.get(clientContact.getToken());
 
-            if (!clientContact.isInactive() && (!existing.isPresent() || existing.get().getRelay().equals(client.getPeerName()))) {
-              directory.add(handle, clientContact);
+            if (!clientContact.isInactive() && (!existing.isPresent() || client.getPeerName().equals(existing.get().getRelay()))) {
+//              directory.add(handle, clientContact);
+              directory.add(clientContact);
             } else {
-              if (existing != null && client.getPeerName().equals(existing.get().getRelay())) {
+              if (existing.isPresent() && client.getPeerName().equals(existing.get().getRelay())) {
                 directory.remove(clientContact.getToken());
               }
             }
           }
 
           retrieved += clientContacts.size();
+          logger.info("Processed: " + retrieved + " remote tokens.");
         }
 
         logger.info("Update from peer complete.");
       } finally {
-        directory.stopBatchOperation(handle);
+//        directory.stopBatchOperation(handle);
       }
     }
 

src/main/resources/migrations.xml

@@ -142,4 +142,33 @@
         </createIndex>
 
     </changeSet>
+
+    <changeSet id="3" author="moxie">
+        <sql>CREATE OR REPLACE FUNCTION "custom_json_object_set_key"(
+            "json"          json,
+            "key_to_set"    TEXT,
+            "value_to_set"  anyelement
+            )
+            RETURNS json
+            LANGUAGE sql
+            IMMUTABLE
+            STRICT
+            AS $function$
+            SELECT COALESCE(
+            (SELECT ('{' || string_agg(to_json("key") || ':' || "value", ',') || '}')
+            FROM (SELECT *
+            FROM json_each("json")
+            WHERE "key" &lt;&gt; "key_to_set"
+            UNION ALL
+            SELECT "key_to_set", to_json("value_to_set")) AS "fields"),
+            '{}'
+            )::json
+            $function$;</sql>
+        <sql>UPDATE accounts SET data = custom_json_object_set_key(data, 'identityKey', k.identity_key) FROM keys k WHERE (data->>'identityKey')::text is null AND k.number = data->>'number' AND k.last_resort = 1;</sql>
+        <sql>UPDATE accounts SET data = custom_json_object_set_key(data, 'identityKey', k.identity_key) FROM keys k WHERE (data->>'identityKey')::text is null AND k.number = data->>'number';</sql>
+    </changeSet>
+
+    <changeSet id="4" author="moxie">
+        <dropColumn tableName="keys" columnName="identity_key"/>
+    </changeSet>
 </databaseChangeLog>

src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DirectoryControllerTest.java

@@ -0,0 +1,74 @@
+package org.whispersystems.textsecuregcm.tests.controllers;
+
+import com.sun.jersey.api.client.ClientResponse;
+import com.yammer.dropwizard.testing.ResourceTest;
+import org.junit.Test;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+import org.whispersystems.textsecuregcm.controllers.DirectoryController;
+import org.whispersystems.textsecuregcm.entities.ClientContactTokens;
+import org.whispersystems.textsecuregcm.limits.RateLimiter;
+import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.storage.DirectoryManager;
+import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
+import org.whispersystems.textsecuregcm.util.Base64;
+
+import javax.ws.rs.core.MediaType;
+import java.util.LinkedList;
+import java.util.List;
+
+import static org.fest.assertions.api.Assertions.assertThat;
+import static org.mockito.Matchers.anyList;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class DirectoryControllerTest extends ResourceTest {
+
+  private RateLimiters     rateLimiters     = mock(RateLimiters.class    );
+  private RateLimiter      rateLimiter      = mock(RateLimiter.class     );
+  private DirectoryManager directoryManager = mock(DirectoryManager.class);
+
+  @Override
+  protected void setUpResources() throws Exception {
+    addProvider(AuthHelper.getAuthenticator());
+
+    when(rateLimiters.getContactsLimiter()).thenReturn(rateLimiter);
+    when(directoryManager.get(anyList())).thenAnswer(new Answer<List<byte[]>>() {
+      @Override
+      public List<byte[]> answer(InvocationOnMock invocationOnMock) throws Throwable {
+        List<byte[]> query    = (List<byte[]>) invocationOnMock.getArguments()[0];
+        List<byte[]> response = new LinkedList<>(query);
+        response.remove(0);
+        return response;
+      }
+    });
+
+    addResource(new DirectoryController(rateLimiters, directoryManager));
+  }
+
+  @Test
+  public void testContactIntersection() throws Exception {
+    List<String> tokens = new LinkedList<String>() {{
+      add(Base64.encodeBytes("foo".getBytes()));
+      add(Base64.encodeBytes("bar".getBytes()));
+      add(Base64.encodeBytes("baz".getBytes()));
+    }};
+
+    List<String> expectedResponse = new LinkedList<>(tokens);
+    expectedResponse.remove(0);
+
+    ClientResponse response =
+        client().resource("/v1/directory/tokens/")
+            .entity(new ClientContactTokens(tokens))
+            .type(MediaType.APPLICATION_JSON_TYPE)
+            .header("Authorization",
+                    AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER,
+                                             AuthHelper.VALID_PASSWORD))
+            .put(ClientResponse.class);
+
+
+
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getEntity(ClientContactTokens.class).getContacts()).isEqualTo(expectedResponse);
+  }
+}

src/test/java/org/whispersystems/textsecuregcm/tests/controllers/FederatedControllerTest.java

@@ -0,0 +1,89 @@
+package org.whispersystems.textsecuregcm.tests.controllers;
+
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Optional;
+import com.sun.jersey.api.client.ClientResponse;
+import com.yammer.dropwizard.testing.ResourceTest;
+import org.junit.Test;
+import org.whispersystems.textsecuregcm.controllers.FederationController;
+import org.whispersystems.textsecuregcm.controllers.MessageController;
+import org.whispersystems.textsecuregcm.entities.IncomingMessageList;
+import org.whispersystems.textsecuregcm.entities.MessageProtos;
+import org.whispersystems.textsecuregcm.federation.FederatedClientManager;
+import org.whispersystems.textsecuregcm.limits.RateLimiter;
+import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.push.PushSender;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
+
+import javax.ws.rs.core.MediaType;
+import java.util.LinkedList;
+import java.util.List;
+
+import static com.yammer.dropwizard.testing.JsonHelpers.jsonFixture;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class FederatedControllerTest extends ResourceTest {
+
+  private static final String SINGLE_DEVICE_RECIPIENT = "+14151111111";
+  private static final String MULTI_DEVICE_RECIPIENT  = "+14152222222";
+
+  private PushSender             pushSender             = mock(PushSender.class            );
+  private FederatedClientManager federatedClientManager = mock(FederatedClientManager.class);
+  private AccountsManager        accountsManager        = mock(AccountsManager.class       );
+  private RateLimiters           rateLimiters           = mock(RateLimiters.class          );
+  private RateLimiter            rateLimiter            = mock(RateLimiter.class           );
+
+  private final ObjectMapper mapper = new ObjectMapper();
+
+  @Override
+  protected void setUpResources() throws Exception {
+    addProvider(AuthHelper.getAuthenticator());
+
+    List<Device> singleDeviceList = new LinkedList<Device>() {{
+      add(new Device(1, "foo", "bar", "baz", "isgcm", null, false, 111));
+    }};
+
+    List<Device> multiDeviceList = new LinkedList<Device>() {{
+      add(new Device(1, "foo", "bar", "baz", "isgcm", null, false, 222));
+      add(new Device(2, "foo", "bar", "baz", "isgcm", null, false, 333));
+    }};
+
+    Account singleDeviceAccount = new Account(SINGLE_DEVICE_RECIPIENT, false, singleDeviceList);
+    Account multiDeviceAccount  = new Account(MULTI_DEVICE_RECIPIENT, false, multiDeviceList);
+
+    when(accountsManager.get(eq(SINGLE_DEVICE_RECIPIENT))).thenReturn(Optional.of(singleDeviceAccount));
+    when(accountsManager.get(eq(MULTI_DEVICE_RECIPIENT))).thenReturn(Optional.of(multiDeviceAccount));
+
+    when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter);
+
+    MessageController messageController = new MessageController(rateLimiters, pushSender, accountsManager, federatedClientManager);
+    addResource(new FederationController(accountsManager, null, null, messageController));
+  }
+
+  @Test
+  public void testSingleDeviceCurrent() throws Exception {
+    ClientResponse response =
+        client().resource(String.format("/v1/federation/messages/+14152223333/1/%s", SINGLE_DEVICE_RECIPIENT))
+            .header("Authorization", AuthHelper.getAuthHeader("cyanogen", "foofoo"))
+            .entity(mapper.readValue(jsonFixture("fixtures/current_message_single_device.json"), IncomingMessageList.class))
+            .type(MediaType.APPLICATION_JSON_TYPE)
+            .put(ClientResponse.class);
+
+    assertThat("Good Response", response.getStatus(), is(equalTo(204)));
+
+    verify(pushSender).sendMessage(any(Account.class), any(Device.class), any(MessageProtos.OutgoingMessageSignal.class));
+  }
+
+
+}

src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java

@@ -4,8 +4,13 @@ import com.google.common.base.Optional;
 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;
+import org.whispersystems.textsecuregcm.entities.PreKeyStatus;
 import org.whispersystems.textsecuregcm.entities.UnstructuredPreKeyList;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -15,6 +20,7 @@ import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.Keys;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 
+import javax.ws.rs.core.MediaType;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -29,11 +35,13 @@ public class KeyControllerTest extends ResourceTest {
   private final int SAMPLE_REGISTRATION_ID  =  999;
   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_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);
+  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_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);
+  private final Account         existsAccount = mock(Account.class        );
+
 
   @Override
   protected void setUpResources() {
@@ -45,7 +53,6 @@ public class KeyControllerTest extends ResourceTest {
     Device  sampleDevice  = mock(Device.class );
     Device  sampleDevice2 = mock(Device.class);
     Device  sampleDevice3 = mock(Device.class);
-    Account existsAccount = mock(Account.class);
 
     when(sampleDevice.getRegistrationId()).thenReturn(SAMPLE_REGISTRATION_ID);
     when(sampleDevice2.getRegistrationId()).thenReturn(SAMPLE_REGISTRATION_ID2);
@@ -58,24 +65,53 @@ 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.<Account>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<Optional<UnstructuredPreKeyList>>() {
+      @Override
+      public Optional<UnstructuredPreKeyList> answer(InvocationOnMock invocationOnMock) throws Throwable {
+        return Optional.of(new UnstructuredPreKeyList(cloneKey(SAMPLE_KEY)));
+      }
+    });
+
     when(keys.get(eq(NOT_EXISTS_NUMBER), eq(1L))).thenReturn(Optional.<UnstructuredPreKeyList>absent());
 
-    List<PreKey> allKeys = new LinkedList<>();
-    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.get(EXISTS_NUMBER)).thenAnswer(new Answer<Optional<UnstructuredPreKeyList>>() {
+      @Override
+      public Optional<UnstructuredPreKeyList> answer(InvocationOnMock invocationOnMock) throws Throwable {
+        List<PreKey> 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.getCount(eq(AuthHelper.VALID_NUMBER), eq(1L))).thenReturn(5);
+
+    when(AuthHelper.VALID_ACCOUNT.getIdentityKey()).thenReturn(null);
 
     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))
@@ -84,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);
@@ -105,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);
@@ -114,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);
@@ -150,4 +186,47 @@ public class KeyControllerTest extends ResourceTest {
     assertThat(response.getClientResponseStatus().getStatusCode()).isEqualTo(401);
   }
 
+  @Test
+  public void putKeysTest() throws Exception {
+    final PreKey newKey        = new PreKey(0, null, 1L, 31337, "foobar", "foobarbaz", false);
+    final PreKey lastResortKey = new PreKey(0, null, 1L, 0xFFFFFF, "fooz", "foobarbaz", false);
+
+    List<PreKey> preKeys = new LinkedList<PreKey>() {{
+      add(newKey);
+    }};
+
+    PreKeyList preKeyList = new PreKeyList();
+    preKeyList.setKeys(preKeys);
+    preKeyList.setLastResortKey(lastResortKey);
+
+    ClientResponse response =
+        client().resource("/v1/keys")
+            .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
+            .type(MediaType.APPLICATION_JSON_TYPE)
+            .put(ClientResponse.class, preKeyList);
+
+    assertThat(response.getClientResponseStatus().getStatusCode()).isEqualTo(204);
+
+    ArgumentCaptor<List>   listCaptor       = ArgumentCaptor.forClass(List.class  );
+    ArgumentCaptor<PreKey> lastResortCaptor = ArgumentCaptor.forClass(PreKey.class);
+    verify(keys).store(eq(AuthHelper.VALID_NUMBER), eq(1L), listCaptor.capture(), lastResortCaptor.capture());
+
+    List<PreKey> capturedList = listCaptor.getValue();
+    assertThat(capturedList.size() == 1);
+    assertThat(capturedList.get(0).getIdentityKey().equals("foobarbaz"));
+    assertThat(capturedList.get(0).getKeyId() == 31337);
+    assertThat(capturedList.get(0).getPublicKey().equals("foobar"));
+
+    assertThat(lastResortCaptor.getValue().getPublicKey().equals("fooz"));
+    assertThat(lastResortCaptor.getValue().getIdentityKey().equals("foobarbaz"));
+
+    verify(AuthHelper.VALID_ACCOUNT).setIdentityKey(eq("foobarbaz"));
+    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());
+  }
+
 }

src/test/java/org/whispersystems/textsecuregcm/tests/util/AuthHelper.java

@@ -13,6 +13,8 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.util.Base64;
 
 import java.util.Arrays;
+import java.util.LinkedList;
+import java.util.List;
 
 import static org.mockito.Matchers.anyLong;
 import static org.mockito.Mockito.mock;
@@ -25,24 +27,31 @@ public class AuthHelper {
   public static final String INVVALID_NUMBER  = "+14151111111";
   public static final String INVALID_PASSWORD = "bar";
 
+  public static AccountsManager           ACCOUNTS_MANAGER  = mock(AccountsManager.class          );
+  public static Account                   VALID_ACCOUNT     = mock(Account.class                  );
+  public static Device                    VALID_DEVICE      = mock(Device.class                   );
+  public static AuthenticationCredentials VALID_CREDENTIALS = mock(AuthenticationCredentials.class);
+
   public static MultiBasicAuthProvider<FederatedPeer, Account> getAuthenticator() {
-    AccountsManager           accounts    = mock(AccountsManager.class          );
-    Account                   account     = mock(Account.class                  );
-    Device                    device      = mock(Device.class                   );
-    AuthenticationCredentials credentials = mock(AuthenticationCredentials.class);
+    when(VALID_CREDENTIALS.verify("foo")).thenReturn(true);
+    when(VALID_DEVICE.getAuthenticationCredentials()).thenReturn(VALID_CREDENTIALS);
+    when(VALID_DEVICE.getId()).thenReturn(1L);
+    when(VALID_ACCOUNT.getDevice(anyLong())).thenReturn(Optional.of(VALID_DEVICE));
+    when(VALID_ACCOUNT.getNumber()).thenReturn(VALID_NUMBER);
+    when(VALID_ACCOUNT.getAuthenticatedDevice()).thenReturn(Optional.of(VALID_DEVICE));
+    when(VALID_ACCOUNT.getRelay()).thenReturn(Optional.<String>absent());
+    when(ACCOUNTS_MANAGER.get(VALID_NUMBER)).thenReturn(Optional.of(VALID_ACCOUNT));
 
-    when(credentials.verify("foo")).thenReturn(true);
-    when(device.getAuthenticationCredentials()).thenReturn(credentials);
-    when(device.getId()).thenReturn(1L);
-    when(account.getDevice(anyLong())).thenReturn(Optional.of(device));
-    when(account.getNumber()).thenReturn(VALID_NUMBER);
-    when(account.getAuthenticatedDevice()).thenReturn(Optional.of(device));
-    when(account.getRelay()).thenReturn(Optional.<String>absent());
-    when(accounts.get(VALID_NUMBER)).thenReturn(Optional.of(account));
+    List<FederatedPeer> peer = new LinkedList<FederatedPeer>() {{
+      add(new FederatedPeer("cyanogen", "https://foo", "foofoo", "bazzzzz"));
+    }};
 
-    return new MultiBasicAuthProvider<>(new FederatedPeerAuthenticator(new FederationConfiguration()),
+    FederationConfiguration federationConfiguration = mock(FederationConfiguration.class);
+    when(federationConfiguration.getPeers()).thenReturn(peer);
+
+    return new MultiBasicAuthProvider<>(new FederatedPeerAuthenticator(federationConfiguration),
                                         FederatedPeer.class,
-                                        new AccountAuthenticator(accounts),
+                                        new AccountAuthenticator(ACCOUNTS_MANAGER),
                                         Account.class, "WhisperServer");
   }