Bump version to 0.16

1) Use WebSockets for delivery if a client is connected.

2) If a client isn't connected, write to a redis queue and send
   an APN push.

config/sample.yml

@@ -56,9 +56,6 @@ graphite:
   host:
   port:
 
-http:
-  shutdownGracePeriod: 0s
-
 database:
   # the name of your JDBC driver
   driverClass: org.postgresql.Driver
@@ -75,24 +72,3 @@ database:
   # any properties specific to your JDBC driver:
   properties:
     charSet: UTF-8
-
-  # the maximum amount of time to wait on an empty pool before throwing an exception
-  maxWaitForConnection: 1s
-
-  # the SQL query to run when validating a connection's liveness
-  validationQuery: "/* MyService Health Check */ SELECT 1"
-
-  # the minimum number of connections to keep open
-  minSize: 8
-
-  # the maximum number of connections to keep open
-  maxSize: 32
-
-  # whether or not idle connections should be validated
-  checkConnectionWhileIdle: false
-
-  # how long a connection must be held before it can be validated
-  checkConnectionHealthWhenIdleFor: 10s
-
-  # the maximum lifetime of an idle connection
-  closeConnectionIfIdleFor: 1 minute

pom.xml

@@ -9,24 +9,72 @@
 
     <groupId>org.whispersystems.textsecure</groupId>
     <artifactId>TextSecureServer</artifactId>
-    <version>0.10</version>
+    <version>0.16</version>
+
+    <properties>
+        <dropwizard.version>0.7.0</dropwizard.version>
+        <jackson.api.version>2.3.3</jackson.api.version>
+        <commons-codec.version>1.6</commons-codec.version>
+    </properties>
 
     <dependencies>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-core</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-jdbi</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-auth</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-client</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-migrations</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-testing</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.dropwizard</groupId>
+            <artifactId>dropwizard-metrics-graphite</artifactId>
+            <version>${dropwizard.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.sun.jersey</groupId>
+            <artifactId>jersey-json</artifactId>
+            <version>1.18.1</version>
+        </dependency>
+        <dependency>
+            <groupId>com.codahale.metrics</groupId>
+            <artifactId>metrics-graphite</artifactId>
+            <version>3.0.2</version>
+        </dependency>
+        <dependency>
+            <groupId>org.eclipse.jetty.websocket</groupId>
+            <artifactId>websocket-server</artifactId>
+            <version>9.0.7.v20131107</version>
+        </dependency>
+
         <dependency>
             <groupId>bouncycastle</groupId>
             <artifactId>bcprov-jdk16</artifactId>
             <version>140</version>
         </dependency>
-        <dependency>
-            <groupId>com.yammer.dropwizard</groupId>
-            <artifactId>dropwizard-core</artifactId>
-            <version>0.6.2</version>
-        </dependency>
-        <dependency>
-            <groupId>com.yammer.metrics</groupId>
-            <artifactId>metrics-graphite</artifactId>
-            <version>2.2.0</version>
-        </dependency>
         <dependency>
             <groupId>com.google.android.gcm</groupId>
             <artifactId>gcm-server</artifactId>
@@ -66,35 +114,10 @@
             <type>jar</type>
             <scope>compile</scope>
         </dependency>
-        <dependency>
-            <groupId>com.yammer.dropwizard</groupId>
-            <artifactId>dropwizard-jdbi</artifactId>
-            <version>0.6.2</version>
-        </dependency>
-        <dependency>
-            <groupId>com.yammer.dropwizard</groupId>
-            <artifactId>dropwizard-auth</artifactId>
-            <version>0.6.2</version>
-        </dependency>
-        <dependency>
-            <groupId>com.yammer.dropwizard</groupId>
-            <artifactId>dropwizard-client</artifactId>
-            <version>0.6.2</version>
-        </dependency>
-        <dependency>
-            <groupId>com.yammer.dropwizard</groupId>
-            <artifactId>dropwizard-migrations</artifactId>
-            <version>0.6.2</version>
-        </dependency>
-        <dependency>
-            <groupId>com.yammer.dropwizard</groupId>
-            <artifactId>dropwizard-testing</artifactId>
-            <version>0.6.2</version>
-        </dependency>
         <dependency>
             <groupId>com.twilio.sdk</groupId>
             <artifactId>twilio-java-sdk</artifactId>
-            <version>3.4.1</version>
+            <version>3.4.5</version>
         </dependency>
 
         <dependency>
@@ -103,24 +126,23 @@
             <version>9.1-901.jdbc4</version>
         </dependency>
 
-        <dependency>
-            <groupId>com.sun.jersey</groupId>
-            <artifactId>jersey-json</artifactId>
-            <version>1.17.1</version>
-        </dependency>
 
-        <dependency>
+    </dependencies>
-            <groupId>org.eclipse.jetty</groupId>
-            <artifactId>jetty-websocket</artifactId>
-            <version>8.1.14.v20131031</version>
-        </dependency>
 
+    <dependencyManagement>
+        <dependencies>
             <dependency>
-            <groupId>org.coursera</groupId>
+                <groupId>com.fasterxml.jackson.core</groupId>
-            <artifactId>metrics-datadog</artifactId>
+                <artifactId>jackson-databind</artifactId>
-            <version>0.1.5</version>
+                <version>${jackson.api.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>commons-codec</groupId>
+                <artifactId>commons-codec</artifactId>
+                <version>${commons-codec.version}</version>
             </dependency>
         </dependencies>
+    </dependencyManagement>
 
     <build>
         <plugins>

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

@@ -17,8 +17,6 @@
 package org.whispersystems.textsecuregcm;
 
 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.FederationConfiguration;
 import org.whispersystems.textsecuregcm.configuration.GcmConfiguration;
@@ -35,6 +33,9 @@ import org.whispersystems.textsecuregcm.configuration.WebsocketConfiguration;
 import javax.validation.Valid;
 import javax.validation.constraints.NotNull;
 
+import io.dropwizard.Configuration;
+import io.dropwizard.db.DataSourceFactory;
+
 public class WhisperServerConfiguration extends Configuration {
 
   @NotNull
@@ -74,7 +75,7 @@ public class WhisperServerConfiguration extends Configuration {
   @Valid
   @NotNull
   @JsonProperty
-  private DatabaseConfiguration database = new DatabaseConfiguration();
+  private DataSourceFactory database = new DataSourceFactory();
 
   @Valid
   @NotNull
@@ -87,7 +88,7 @@ public class WhisperServerConfiguration extends Configuration {
 
   @Valid
   @JsonProperty
-  private MetricsConfiguration metrics = new MetricsConfiguration();
+  private MetricsConfiguration viz = new MetricsConfiguration();
 
   @Valid
   @JsonProperty
@@ -125,7 +126,7 @@ public class WhisperServerConfiguration extends Configuration {
     return redis;
   }
 
-  public DatabaseConfiguration getDatabaseConfiguration() {
+  public DataSourceFactory getDataSourceFactory() {
     return database;
   }
 
@@ -142,6 +143,6 @@ public class WhisperServerConfiguration extends Configuration {
   }
 
   public MetricsConfiguration getMetricsConfiguration() {
-    return metrics;
+    return viz;
   }
 }

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

@@ -16,18 +16,13 @@
  */
 package org.whispersystems.textsecuregcm;
 
+import com.codahale.metrics.SharedMetricRegistries;
+import com.codahale.metrics.graphite.GraphiteReporter;
+import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.Service;
-import com.yammer.dropwizard.config.Bootstrap;
-import com.yammer.dropwizard.config.Environment;
-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.Metrics;
-import com.yammer.metrics.reporting.GraphiteReporter;
 import net.spy.memcached.MemcachedClient;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.eclipse.jetty.servlets.CrossOriginFilter;
 import org.skife.jdbi.v2.DBI;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.FederatedPeerAuthenticator;
@@ -40,7 +35,6 @@ import org.whispersystems.textsecuregcm.controllers.DirectoryController;
 import org.whispersystems.textsecuregcm.controllers.FederationController;
 import org.whispersystems.textsecuregcm.controllers.KeysController;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
-import org.whispersystems.textsecuregcm.controllers.WebsocketControllerFactory;
 import org.whispersystems.textsecuregcm.federation.FederatedClientManager;
 import org.whispersystems.textsecuregcm.federation.FederatedPeer;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -69,18 +63,30 @@ import org.whispersystems.textsecuregcm.storage.PendingAccountsManager;
 import org.whispersystems.textsecuregcm.storage.PendingDevices;
 import org.whispersystems.textsecuregcm.storage.PendingDevicesManager;
 import org.whispersystems.textsecuregcm.storage.PubSubManager;
-import org.whispersystems.textsecuregcm.storage.StoredMessageManager;
 import org.whispersystems.textsecuregcm.storage.StoredMessages;
-import org.whispersystems.textsecuregcm.util.CORSHeaderFilter;
+import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.UrlSigner;
+import org.whispersystems.textsecuregcm.websocket.WebsocketControllerFactory;
 import org.whispersystems.textsecuregcm.workers.DirectoryCommand;
 
+import javax.servlet.DispatcherType;
+import javax.servlet.FilterRegistration;
+import javax.servlet.ServletRegistration;
 import java.security.Security;
+import java.util.EnumSet;
 import java.util.concurrent.TimeUnit;
 
+import static com.codahale.metrics.MetricRegistry.name;
+import io.dropwizard.Application;
+import io.dropwizard.db.DataSourceFactory;
+import io.dropwizard.jdbi.DBIFactory;
+import io.dropwizard.metrics.graphite.GraphiteReporterFactory;
+import io.dropwizard.migrations.MigrationsBundle;
+import io.dropwizard.setup.Bootstrap;
+import io.dropwizard.setup.Environment;
 import redis.clients.jedis.JedisPool;
 
-public class WhisperServerService extends Service<WhisperServerConfiguration> {
+public class WhisperServerService extends Application<WhisperServerConfiguration> {
 
   static {
     Security.addProvider(new BouncyCastleProvider());
@@ -88,30 +94,34 @@ public class WhisperServerService extends Service<WhisperServerConfiguration> {
 
   @Override
   public void initialize(Bootstrap<WhisperServerConfiguration> bootstrap) {
-    bootstrap.setName("whisper-server");
     bootstrap.addCommand(new DirectoryCommand());
     bootstrap.addBundle(new MigrationsBundle<WhisperServerConfiguration>() {
       @Override
-      public DatabaseConfiguration getDatabaseConfiguration(WhisperServerConfiguration configuration) {
+      public DataSourceFactory getDataSourceFactory(WhisperServerConfiguration configuration) {
-        return configuration.getDatabaseConfiguration();
+        return configuration.getDataSourceFactory();
       }
     });
   }
 
+  @Override
+  public String getName() {
+    return "whisper-server";
+  }
+
   @Override
   public void run(WhisperServerConfiguration config, Environment environment)
       throws Exception
   {
-    config.getHttpConfiguration().setConnectorType(HttpConfiguration.ConnectorType.NONBLOCKING);
+    SharedMetricRegistries.add(Constants.METRICS_NAME, environment.metrics());
+    environment.getObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
 
     DBIFactory dbiFactory = new DBIFactory();
-    DBI        jdbi       = dbiFactory.build(environment, config.getDatabaseConfiguration(), "postgresql");
+    DBI        jdbi       = dbiFactory.build(environment, config.getDataSourceFactory(), "postgresql");
 
     Accounts        accounts        = jdbi.onDemand(Accounts.class);
     PendingAccounts pendingAccounts = jdbi.onDemand(PendingAccounts.class);
     PendingDevices  pendingDevices  = jdbi.onDemand(PendingDevices.class);
     Keys            keys            = jdbi.onDemand(Keys.class);
-    StoredMessages  storedMessages  = jdbi.onDemand(StoredMessages.class );
 
     MemcachedClient memcachedClient = new MemcachedClientFactory(config.getMemcacheConfiguration()).getClient();
     JedisPool       redisClient     = new RedisClientFactory(config.getRedisConfiguration()).getRedisClientPool();
@@ -121,8 +131,8 @@ public class WhisperServerService extends Service<WhisperServerConfiguration> {
     PendingDevicesManager  pendingDevicesManager  = new PendingDevicesManager (pendingDevices, memcachedClient );
     AccountsManager        accountsManager        = new AccountsManager(accounts, directory, memcachedClient);
     FederatedClientManager federatedClientManager = new FederatedClientManager(config.getFederationConfiguration());
+    StoredMessages         storedMessages         = new StoredMessages(redisClient);
     PubSubManager          pubSubManager          = new PubSubManager(redisClient);
-    StoredMessageManager     storedMessageManager   = new StoredMessageManager(storedMessages, pubSubManager);
 
     AccountAuthenticator     deviceAuthenticator    = new AccountAuthenticator(accountsManager);
     RateLimiters             rateLimiters           = new RateLimiters(config.getLimitsConfiguration(), memcachedClient);
@@ -133,51 +143,67 @@ public class WhisperServerService extends Service<WhisperServerConfiguration> {
     UrlSigner                urlSigner              = new UrlSigner(config.getS3Configuration());
     PushSender               pushSender             = new PushSender(config.getGcmConfiguration(),
                                                                      config.getApnConfiguration(),
-                                                                     storedMessageManager,
+                                                                     storedMessages, pubSubManager,
                                                                      accountsManager);
 
     AttachmentController attachmentController = new AttachmentController(rateLimiters, federatedClientManager, urlSigner);
     KeysController       keysController       = new KeysController(rateLimiters, keys, accountsManager, federatedClientManager);
     MessageController    messageController    = new MessageController(rateLimiters, pushSender, accountsManager, federatedClientManager);
 
-    environment.addProvider(new MultiBasicAuthProvider<>(new FederatedPeerAuthenticator(config.getFederationConfiguration()),
+    environment.jersey().register(new MultiBasicAuthProvider<>(new FederatedPeerAuthenticator(config.getFederationConfiguration()),
                                                                FederatedPeer.class,
                                                                deviceAuthenticator,
                                                                Device.class, "WhisperServer"));
 
-    environment.addResource(new AccountController(pendingAccountsManager, accountsManager, rateLimiters, smsSender));
+    environment.jersey().register(new AccountController(pendingAccountsManager, accountsManager, rateLimiters, smsSender));
-    environment.addResource(new DeviceController(pendingDevicesManager, accountsManager, rateLimiters));
+    environment.jersey().register(new DeviceController(pendingDevicesManager, accountsManager, rateLimiters));
-    environment.addResource(new DirectoryController(rateLimiters, directory));
+    environment.jersey().register(new DirectoryController(rateLimiters, directory));
-    environment.addResource(new FederationController(accountsManager, attachmentController, keysController, messageController));
+    environment.jersey().register(new FederationController(accountsManager, attachmentController, keysController, messageController));
-    environment.addResource(attachmentController);
+    environment.jersey().register(attachmentController);
-    environment.addResource(keysController);
+    environment.jersey().register(keysController);
-    environment.addResource(messageController);
+    environment.jersey().register(messageController);
 
     if (config.getWebsocketConfiguration().isEnabled()) {
-      environment.addServlet(new WebsocketControllerFactory(deviceAuthenticator, storedMessageManager, pubSubManager),
+      WebsocketControllerFactory servlet = new WebsocketControllerFactory(deviceAuthenticator,
-                             "/v1/websocket/");
+                                                                          pushSender,
-      environment.addFilter(new CORSHeaderFilter(), "/*");
+                                                                          storedMessages,
+                                                                          pubSubManager);
+
+      ServletRegistration.Dynamic websocket = environment.servlets().addServlet("WebSocket", servlet);
+      websocket.addMapping("/v1/websocket/*");
+      websocket.setAsyncSupported(true);
+
+      FilterRegistration.Dynamic filter = environment.servlets().addFilter("CORS", CrossOriginFilter.class);
+      filter.addMappingForUrlPatterns(EnumSet.allOf(DispatcherType.class), true, "/*");
+      filter.setInitParameter("allowedOrigins", "*");
+      filter.setInitParameter("allowedHeaders", "Content-Type,Authorization,X-Requested-With,Content-Length,Accept,Origin");
+      filter.setInitParameter("allowedMethods", "GET,PUT,POST,DELETE,OPTIONS");
+      filter.setInitParameter("preflightMaxAge", "5184000");
+      filter.setInitParameter("allowCredentials", "true");
     }
 
-    environment.addHealthCheck(new RedisHealthCheck(redisClient));
+    environment.healthChecks().register("redis", new RedisHealthCheck(redisClient));
-    environment.addHealthCheck(new MemcacheHealthCheck(memcachedClient));
+    environment.healthChecks().register("memcache", new MemcacheHealthCheck(memcachedClient));
 
-    environment.addProvider(new IOExceptionMapper());
+    environment.jersey().register(new IOExceptionMapper());
-    environment.addProvider(new RateLimitExceededExceptionMapper());
+    environment.jersey().register(new RateLimitExceededExceptionMapper());
 
-    Metrics.newGauge(CpuUsageGauge.class, "cpu", new CpuUsageGauge());
+    environment.metrics().register(name(CpuUsageGauge.class, "cpu"), new CpuUsageGauge());
-    Metrics.newGauge(FreeMemoryGauge.class, "free_memory", new FreeMemoryGauge());
+    environment.metrics().register(name(FreeMemoryGauge.class, "free_memory"), new FreeMemoryGauge());
-    Metrics.newGauge(NetworkSentGauge.class, "bytes_sent", new NetworkSentGauge());
+    environment.metrics().register(name(NetworkSentGauge.class, "bytes_sent"), new NetworkSentGauge());
-    Metrics.newGauge(NetworkReceivedGauge.class, "bytes_received", new NetworkReceivedGauge());
+    environment.metrics().register(name(NetworkReceivedGauge.class, "bytes_received"), new NetworkReceivedGauge());
 
     if (config.getGraphiteConfiguration().isEnabled()) {
-      GraphiteReporter.enable(15, TimeUnit.SECONDS,
+      GraphiteReporterFactory graphiteReporterFactory = new GraphiteReporterFactory();
-                              config.getGraphiteConfiguration().getHost(),
+      graphiteReporterFactory.setHost(config.getGraphiteConfiguration().getHost());
-                              config.getGraphiteConfiguration().getPort());
+      graphiteReporterFactory.setPort(config.getGraphiteConfiguration().getPort());
+
+      GraphiteReporter graphiteReporter = (GraphiteReporter) graphiteReporterFactory.build(environment.metrics());
+      graphiteReporter.start(15, TimeUnit.SECONDS);
     }
 
     if (config.getMetricsConfiguration().isEnabled()) {
-      new JsonMetricsReporter("textsecure", Metrics.defaultRegistry(),
+      new JsonMetricsReporter(environment.metrics(),
                               config.getMetricsConfiguration().getToken(),
                               config.getMetricsConfiguration().getHost())
           .start(60, TimeUnit.SECONDS);

src/main/java/org/whispersystems/textsecuregcm/auth/AccountAuthenticator.java

@@ -16,29 +16,27 @@
  */
 package org.whispersystems.textsecuregcm.auth;
 
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.AuthenticationException;
-import com.yammer.dropwizard.auth.Authenticator;
-import com.yammer.dropwizard.auth.basic.BasicCredentials;
-import com.yammer.metrics.Metrics;
-import com.yammer.metrics.core.Meter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.storage.Account;
-import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.util.Constants;
 
-import java.util.concurrent.TimeUnit;
+import static com.codahale.metrics.MetricRegistry.name;
+import io.dropwizard.auth.AuthenticationException;
+import io.dropwizard.auth.Authenticator;
+import io.dropwizard.auth.basic.BasicCredentials;
 
 public class AccountAuthenticator implements Authenticator<BasicCredentials, Account> {
 
-  private final Meter authenticationFailedMeter    = Metrics.newMeter(AccountAuthenticator.class,
+  private final MetricRegistry metricRegistry               = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-                                                                      "authentication", "failed",
+  private final Meter          authenticationFailedMeter    = metricRegistry.meter(name(getClass(), "authentication", "failed"   ));
-                                                                      TimeUnit.MINUTES);
+  private final Meter          authenticationSucceededMeter = metricRegistry.meter(name(getClass(), "authentication", "succeeded"));
-
-  private final Meter authenticationSucceededMeter = Metrics.newMeter(AccountAuthenticator.class,
-                                                                      "authentication", "succeeded",
-                                                                      TimeUnit.MINUTES);
 
   private final Logger logger = LoggerFactory.getLogger(AccountAuthenticator.class);
 

src/main/java/org/whispersystems/textsecuregcm/auth/FederatedPeerAuthenticator.java

@@ -16,30 +16,35 @@
  */
 package org.whispersystems.textsecuregcm.auth;
 
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.AuthenticationException;
-import com.yammer.dropwizard.auth.Authenticator;
-import com.yammer.dropwizard.auth.basic.BasicCredentials;
-import com.yammer.metrics.Metrics;
-import com.yammer.metrics.core.Meter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.configuration.FederationConfiguration;
 import org.whispersystems.textsecuregcm.federation.FederatedPeer;
+import org.whispersystems.textsecuregcm.util.Constants;
 
 import java.util.List;
-import java.util.concurrent.TimeUnit;
+
+import static com.codahale.metrics.MetricRegistry.name;
+import io.dropwizard.auth.AuthenticationException;
+import io.dropwizard.auth.Authenticator;
+import io.dropwizard.auth.basic.BasicCredentials;
 
 
 public class FederatedPeerAuthenticator implements Authenticator<BasicCredentials, FederatedPeer> {
 
-  private final Meter authenticationFailedMeter    = Metrics.newMeter(FederatedPeerAuthenticator.class,
+  private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-                                                                      "authentication", "failed",
-                                                                      TimeUnit.MINUTES);
 
-  private final Meter authenticationSucceededMeter = Metrics.newMeter(FederatedPeerAuthenticator.class,
+  private final Meter authenticationFailedMeter    = metricRegistry.meter(name(getClass(),
-                                                                      "authentication", "succeeded",
+                                                                               "authentication",
-                                                                      TimeUnit.MINUTES);
+                                                                               "failed"));
+
+  private final Meter authenticationSucceededMeter = metricRegistry.meter(name(getClass(),
+                                                                               "authentication",
+                                                                               "succeeded"));
 
   private final Logger logger = LoggerFactory.getLogger(FederatedPeerAuthenticator.class);
 

src/main/java/org/whispersystems/textsecuregcm/auth/MultiBasicAuthProvider.java

@@ -21,17 +21,14 @@ import com.sun.jersey.core.spi.component.ComponentContext;
 import com.sun.jersey.core.spi.component.ComponentScope;
 import com.sun.jersey.spi.inject.Injectable;
 import com.sun.jersey.spi.inject.InjectableProvider;
-import com.yammer.dropwizard.auth.Auth;
+
-import com.yammer.dropwizard.auth.Authenticator;
+import io.dropwizard.auth.Auth;
-import com.yammer.dropwizard.auth.basic.BasicAuthProvider;
+import io.dropwizard.auth.Authenticator;
-import com.yammer.dropwizard.auth.basic.BasicCredentials;
+import io.dropwizard.auth.basic.BasicAuthProvider;
-import org.slf4j.Logger;
+import io.dropwizard.auth.basic.BasicCredentials;
-import org.slf4j.LoggerFactory;
 
 public class MultiBasicAuthProvider<T1,T2> implements InjectableProvider<Auth, Parameter> {
 
-  private final Logger logger = LoggerFactory.getLogger(MultiBasicAuthProvider.class);
-
   private final BasicAuthProvider<T1> provider1;
   private final BasicAuthProvider<T2> provider2;
 
@@ -44,8 +41,8 @@ public class MultiBasicAuthProvider<T1,T2> implements InjectableProvider<Auth, P
                                 Class<?> clazz2,
                                 String realm)
   {
-    this.provider1 = new BasicAuthProvider<T1>(authenticator1, realm);
+    this.provider1 = new BasicAuthProvider<>(authenticator1, realm);
-    this.provider2 = new BasicAuthProvider<T2>(authenticator2, realm);
+    this.provider2 = new BasicAuthProvider<>(authenticator2, realm);
     this.clazz1    = clazz1;
     this.clazz2    = clazz2;
   }

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

@@ -1,20 +0,0 @@
-package org.whispersystems.textsecuregcm.configuration;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-
-public class DataDogConfiguration {
-
-  @JsonProperty
-  private String apiKey;
-
-  @JsonProperty
-  private boolean enabled = false;
-
-  public String getApiKey() {
-    return apiKey;
-  }
-
-  public boolean isEnabled() {
-    return enabled && apiKey != null;
-  }
-}

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

@@ -16,10 +16,9 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.codahale.metrics.annotation.Timed;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.Auth;
-import com.yammer.metrics.annotation.Timed;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
@@ -55,6 +54,8 @@ import java.io.IOException;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
 
+import io.dropwizard.auth.Auth;
+
 @Path("/v1/accounts")
 public class AccountController {
 

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

@@ -17,9 +17,8 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 import com.amazonaws.HttpMethod;
+import com.codahale.metrics.annotation.Timed;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.Auth;
-import com.yammer.metrics.annotation.Timed;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.AttachmentDescriptor;
@@ -44,6 +43,8 @@ import java.net.URL;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
 
+import io.dropwizard.auth.Auth;
+
 
 @Path("/v1/attachments")
 public class AttachmentController {

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

@@ -16,10 +16,9 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.codahale.metrics.annotation.Timed;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.Auth;
-import com.yammer.metrics.annotation.Timed;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
@@ -29,8 +28,8 @@ import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.entities.DeviceResponse;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.Account;
-import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.PendingDevicesManager;
 import org.whispersystems.textsecuregcm.util.VerificationCode;
 
@@ -48,6 +47,8 @@ import javax.ws.rs.core.Response;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
 
+import io.dropwizard.auth.Auth;
+
 @Path("/v1/devices")
 public class DeviceController {
 

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

@@ -16,11 +16,11 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.codahale.metrics.Histogram;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
+import com.codahale.metrics.annotation.Timed;
 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;
@@ -30,6 +30,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.DirectoryManager;
 import org.whispersystems.textsecuregcm.util.Base64;
+import org.whispersystems.textsecuregcm.util.Constants;
 
 import javax.validation.Valid;
 import javax.ws.rs.Consumes;
@@ -45,11 +46,15 @@ import java.io.IOException;
 import java.util.LinkedList;
 import java.util.List;
 
+import static com.codahale.metrics.MetricRegistry.name;
+import io.dropwizard.auth.Auth;
+
 @Path("/v1/directory")
 public class DirectoryController {
 
   private final Logger         logger            = LoggerFactory.getLogger(DirectoryController.class);
-  private final Histogram contactsHistogram = Metrics.newHistogram(DirectoryController.class, "contacts");
+  private final MetricRegistry metricRegistry    = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
+  private final Histogram      contactsHistogram = metricRegistry.histogram(name(getClass(), "contacts"));
 
   private final RateLimiters     rateLimiters;
   private final DirectoryManager directory;

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

@@ -16,9 +16,8 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.codahale.metrics.annotation.Timed;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.Auth;
-import com.yammer.metrics.annotation.Timed;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.AccountCount;
@@ -45,6 +44,8 @@ import java.io.IOException;
 import java.util.LinkedList;
 import java.util.List;
 
+import io.dropwizard.auth.Auth;
+
 @Path("/v1/federation")
 public class FederationController {
 

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

@@ -16,9 +16,8 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.codahale.metrics.annotation.Timed;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.Auth;
-import com.yammer.metrics.annotation.Timed;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.PreKey;
@@ -47,6 +46,8 @@ import javax.ws.rs.core.Response;
 import java.util.LinkedList;
 import java.util.List;
 
+import io.dropwizard.auth.Auth;
+
 @Path("/v1/keys")
 public class KeysController {
 
@@ -83,7 +84,6 @@ public class KeysController {
 
   @Timed
   @GET
-  @Path("/")
   @Produces(MediaType.APPLICATION_JSON)
   public PreKeyStatus getStatus(@Auth Account account) {
     int count = keys.getCount(account.getNumber(), account.getAuthenticatedDevice().get().getId());
@@ -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);
       }
     }

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

@@ -16,10 +16,9 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import com.codahale.metrics.annotation.Timed;
 import com.google.common.base.Optional;
 import com.google.protobuf.ByteString;
-import com.yammer.dropwizard.auth.Auth;
-import com.yammer.metrics.annotation.Timed;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.entities.IncomingMessage;
@@ -56,6 +55,8 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
 
+import io.dropwizard.auth.Auth;
+
 @Path("/v1/messages")
 public class MessageController {
 

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

@@ -1,7 +1,6 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 import java.util.List;
-import java.util.Set;
 
 public class MismatchedDevicesException extends Exception {
 

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

@@ -16,8 +16,6 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
-import org.whispersystems.textsecuregcm.federation.NoSuchPeerException;
-
 import java.util.LinkedList;
 import java.util.List;
 
@@ -27,7 +25,7 @@ public class NoSuchUserException extends Exception {
 
   public NoSuchUserException(String user) {
     super(user);
-    missing = new LinkedList<String>();
+    missing = new LinkedList<>();
     missing.add(user);
   }
 

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

@@ -1,17 +1,26 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import org.eclipse.jetty.websocket.WebSocket;
+import com.google.common.base.Optional;
+import org.eclipse.jetty.websocket.api.CloseStatus;
+import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.api.UpgradeRequest;
+import org.eclipse.jetty.websocket.api.WebSocketListener;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.entities.AcknowledgeWebsocketMessage;
+import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
 import org.whispersystems.textsecuregcm.entities.IncomingWebsocketMessage;
+import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
+import org.whispersystems.textsecuregcm.push.PushSender;
+import org.whispersystems.textsecuregcm.push.TransientPushFailureException;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.PubSubListener;
 import org.whispersystems.textsecuregcm.storage.PubSubManager;
 import org.whispersystems.textsecuregcm.storage.PubSubMessage;
-import org.whispersystems.textsecuregcm.storage.StoredMessageManager;
+import org.whispersystems.textsecuregcm.storage.StoredMessages;
 import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
 import org.whispersystems.textsecuregcm.websocket.WebsocketMessage;
 
@@ -22,65 +31,132 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 
-public class WebsocketController implements WebSocket.OnTextMessage, PubSubListener {
+import io.dropwizard.auth.AuthenticationException;
+import io.dropwizard.auth.basic.BasicCredentials;
+
+public class WebsocketController implements WebSocketListener, PubSubListener {
 
   private static final Logger            logger          = LoggerFactory.getLogger(WebsocketController.class);
   private static final ObjectMapper      mapper          = new ObjectMapper();
   private static final Map<Long, String> pendingMessages = new HashMap<>();
 
-  private final StoredMessageManager storedMessageManager;
+  private final AccountAuthenticator accountAuthenticator;
   private final PubSubManager        pubSubManager;
+  private final StoredMessages       storedMessages;
+  private final PushSender           pushSender;
 
-  private final Account account;
+  private WebsocketAddress address;
-  private final Device  device;
+  private Account          account;
+  private Device           device;
+  private Session          session;
 
-  private Connection connection;
   private long pendingMessageSequence;
 
-  public WebsocketController(StoredMessageManager storedMessageManager,
+  public WebsocketController(AccountAuthenticator accountAuthenticator,
+                             PushSender           pushSender,
                              PubSubManager        pubSubManager,
-                             Account account)
+                             StoredMessages       storedMessages)
   {
-    this.storedMessageManager = storedMessageManager;
+    this.accountAuthenticator = accountAuthenticator;
+    this.pushSender           = pushSender;
     this.pubSubManager        = pubSubManager;
-    this.account              = account;
+    this.storedMessages       = storedMessages;
-    this.device               = account.getAuthenticatedDevice().get();
   }
 
-
   @Override
-  public void onOpen(Connection connection) {
+  public void onWebSocketConnect(Session session) {
-    this.connection = connection;
+    try {
-    pubSubManager.subscribe(new WebsocketAddress(this.account.getId(), this.device.getId()), this);
+      UpgradeRequest        request    = session.getUpgradeRequest();
+      Map<String, String[]> parameters = request.getParameterMap();
+      String[]              usernames  = parameters.get("login"   );
+      String[]              passwords  = parameters.get("password");
+
+      if (usernames == null || usernames.length == 0 ||
+          passwords == null || passwords.length == 0)
+      {
+        session.close(new CloseStatus(4001, "Unauthorized"));
+        return;
+      }
+
+      BasicCredentials  credentials = new BasicCredentials(usernames[0], passwords[0]);
+      Optional<Account> account     = accountAuthenticator.authenticate(credentials);
+
+      if (!account.isPresent()) {
+        session.close(new CloseStatus(4001, "Unauthorized"));
+        return;
+      }
+
+      this.account = account.get();
+      this.device  = account.get().getAuthenticatedDevice().get();
+      this.address = new WebsocketAddress(this.account.getId(), this.device.getId());
+      this.session = session;
+
+      this.session.setIdleTimeout(10 * 60 * 1000);
+      this.pubSubManager.subscribe(this.address, this);
+
       handleQueryDatabase();
+    } catch (AuthenticationException e) {
+      try { session.close(1011, "Server Error");} catch (IOException e1) {}
+    } catch (IOException ioe) {
+      logger.info("Abrupt session close.");
+    }
   }
 
   @Override
-  public void onClose(int i, String s) {
+  public void onWebSocketText(String body) {
-    handleClose();
-  }
-
-  @Override
-  public void onMessage(String body) {
     try {
       IncomingWebsocketMessage incomingMessage = mapper.readValue(body, IncomingWebsocketMessage.class);
 
       switch (incomingMessage.getType()) {
-        case IncomingWebsocketMessage.TYPE_ACKNOWLEDGE_MESSAGE: handleMessageAck(body); break;
+        case IncomingWebsocketMessage.TYPE_ACKNOWLEDGE_MESSAGE:
-        case IncomingWebsocketMessage.TYPE_PING_MESSAGE:        handlePing();           break;
+          handleMessageAck(body);
-        default:                                                handleClose();          break;
+          break;
+        default:
+          close(new CloseStatus(1008, "Unknown Type"));
       }
     } catch (IOException e) {
       logger.debug("Parse", e);
-      handleClose();
+      close(new CloseStatus(1008, "Badly Formatted"));
     }
   }
 
+  @Override
+  public void onWebSocketClose(int i, String s) {
+    pubSubManager.unsubscribe(this.address, this);
+
+    List<String> remainingMessages = new LinkedList<>();
+
+    synchronized (pendingMessages) {
+      Long[] pendingKeys = pendingMessages.keySet().toArray(new Long[0]);
+      Arrays.sort(pendingKeys);
+
+      for (long pendingKey : pendingKeys) {
+        remainingMessages.add(pendingMessages.get(pendingKey));
+      }
+
+      pendingMessages.clear();
+    }
+
+    for (String remainingMessage : remainingMessages) {
+      try {
+        pushSender.sendMessage(account, device, new EncryptedOutgoingMessage(remainingMessage));
+      } catch (NotPushRegisteredException | TransientPushFailureException e) {
+        logger.warn("onWebSocketClose", e);
+        storedMessages.insert(account.getId(), device.getId(), remainingMessage);
+      }
+    }
+  }
+
+
   @Override
   public void onPubSubMessage(PubSubMessage outgoingMessage) {
     switch (outgoingMessage.getType()) {
-      case PubSubMessage.TYPE_DELIVER:  handleDeliverOutgoingMessage(outgoingMessage.getContents()); break;
+      case PubSubMessage.TYPE_DELIVER:
-      case PubSubMessage.TYPE_QUERY_DB: handleQueryDatabase();                                       break;
+        handleDeliverOutgoingMessage(outgoingMessage.getContents());
+        break;
+      case PubSubMessage.TYPE_QUERY_DB:
+        handleQueryDatabase();
+        break;
       default:
         logger.warn("Unknown pubsub message: " + outgoingMessage.getType());
     }
@@ -95,10 +171,11 @@ public class WebsocketController implements WebSocket.OnTextMessage, PubSubListe
         pendingMessages.put(messageSequence, message);
       }
 
-      connection.sendMessage(mapper.writeValueAsString(new WebsocketMessage(messageSequence, message)));
+      WebsocketMessage websocketMessage = new WebsocketMessage(messageSequence, message);
+      session.getRemote().sendStringByFuture(mapper.writeValueAsString(websocketMessage));
     } catch (IOException e) {
       logger.debug("Response failed", e);
-      handleClose();
+      close(null);
     }
   }
 
@@ -114,42 +191,33 @@ public class WebsocketController implements WebSocket.OnTextMessage, PubSubListe
     }
   }
 
-  private void handlePing() {
-    try {
-      IncomingWebsocketMessage pongMessage = new IncomingWebsocketMessage(IncomingWebsocketMessage.TYPE_PONG_MESSAGE);
-      connection.sendMessage(mapper.writeValueAsString(pongMessage));
-    } catch (IOException e) {
-      logger.warn("Pong failed", e);
-      handleClose();
-    }
-  }
-
-  private void handleClose() {
-    pubSubManager.unsubscribe(new WebsocketAddress(account.getId(), device.getId()), this);
-    connection.close();
-
-    List<String> remainingMessages = new LinkedList<>();
-
-    synchronized (pendingMessages) {
-      Long[] pendingKeys = pendingMessages.keySet().toArray(new Long[0]);
-      Arrays.sort(pendingKeys);
-
-      for (long pendingKey : pendingKeys) {
-        remainingMessages.add(pendingMessages.get(pendingKey));
-      }
-
-      pendingMessages.clear();
-    }
-
-    storedMessageManager.storeMessages(account, device, remainingMessages);
-  }
-
   private void handleQueryDatabase() {
-    List<String> messages = storedMessageManager.getOutgoingMessages(account, device);
+    List<String> messages = storedMessages.getMessagesForDevice(account.getId(), device.getId());
 
     for (String message : messages) {
       handleDeliverOutgoingMessage(message);
     }
   }
 
+  @Override
+  public void onWebSocketBinary(byte[] bytes, int i, int i2) {
+    logger.info("Received binary message!");
+  }
+
+  @Override
+  public void onWebSocketError(Throwable throwable) {
+    logger.info("onWebSocketError", throwable);
+  }
+
+
+  private void close(CloseStatus closeStatus) {
+    try {
+      if (this.session != null) {
+        if (closeStatus != null) this.session.close(closeStatus);
+        else                     this.session.close();
+      }
+    } catch (IOException e) {
+      logger.info("close()", e);
+    }
+  }
 }

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

@@ -1,98 +0,0 @@
-package org.whispersystems.textsecuregcm.controllers;
-
-import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.AuthenticationException;
-import com.yammer.dropwizard.auth.basic.BasicCredentials;
-import org.eclipse.jetty.websocket.WebSocket;
-import org.eclipse.jetty.websocket.WebSocketServlet;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
-import org.whispersystems.textsecuregcm.storage.Account;
-import org.whispersystems.textsecuregcm.storage.PubSubManager;
-import org.whispersystems.textsecuregcm.storage.StoredMessageManager;
-
-import javax.servlet.http.HttpServletRequest;
-import java.util.LinkedHashMap;
-import java.util.Map;
-
-
-public class WebsocketControllerFactory extends WebSocketServlet {
-
-  private final Logger logger = LoggerFactory.getLogger(WebsocketControllerFactory.class);
-
-  private final StoredMessageManager storedMessageManager;
-  private final PubSubManager        pubSubManager;
-  private final AccountAuthenticator accountAuthenticator;
-
-  private final LinkedHashMap<BasicCredentials, Optional<Account>> cache =
-      new LinkedHashMap<BasicCredentials, Optional<Account>>() {
-    @Override
-    protected boolean removeEldestEntry(Map.Entry<BasicCredentials, Optional<Account>> eldest) {
-      return size() > 10;
-    }
-  };
-
-  public WebsocketControllerFactory(AccountAuthenticator accountAuthenticator,
-                                    StoredMessageManager storedMessageManager,
-                                    PubSubManager pubSubManager)
-  {
-    this.accountAuthenticator = accountAuthenticator;
-    this.storedMessageManager = storedMessageManager;
-    this.pubSubManager        = pubSubManager;
-  }
-
-  @Override
-  public WebSocket doWebSocketConnect(HttpServletRequest request, String s) {
-    try {
-      String username = request.getParameter("user");
-      String password = request.getParameter("password");
-
-      if (username == null || password == null) {
-        return null;
-      }
-
-      BasicCredentials credentials = new BasicCredentials(username, password);
-
-      Optional<Account> account = cache.remove(credentials);
-
-      if (account == null) {
-        account = accountAuthenticator.authenticate(new BasicCredentials(username, password));
-      }
-
-      if (!account.isPresent()) {
-        return null;
-      }
-
-      return new WebsocketController(storedMessageManager, pubSubManager, account.get());
-    } catch (AuthenticationException e) {
-      throw new AssertionError(e);
-    }
-  }
-
-  @Override
-  public boolean checkOrigin(HttpServletRequest request, String origin) {
-    try {
-      String username = request.getParameter("user");
-      String password = request.getParameter("password");
-
-      if (username == null || password == null) {
-        return false;
-      }
-
-      BasicCredentials  credentials = new BasicCredentials(username, password);
-      Optional<Account> account     = accountAuthenticator.authenticate(credentials);
-
-      if (!account.isPresent()) {
-        return false;
-      }
-
-      cache.put(credentials, account);
-
-      return true;
-    } catch (AuthenticationException e) {
-      logger.warn("Auth Failure", e);
-      return false;
-    }
-  }
-}

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

@@ -41,23 +41,26 @@ public class EncryptedOutgoingMessage {
   private static final int    MAC_KEY_SIZE    = 20;
   private static final int    MAC_SIZE        = 10;
 
-  private final OutgoingMessageSignal outgoingMessage;
+  private final String serialized;
-  private final String signalingKey;
 
   public EncryptedOutgoingMessage(OutgoingMessageSignal outgoingMessage,
                                   String signalingKey)
+      throws CryptoEncodingException
   {
-    this.outgoingMessage = outgoingMessage;
-    this.signalingKey    = signalingKey;
-  }
-
-  public String serialize() throws CryptoEncodingException {
     byte[]        plaintext  = outgoingMessage.toByteArray();
     SecretKeySpec cipherKey  = getCipherKey (signalingKey);
     SecretKeySpec macKey     = getMacKey(signalingKey);
     byte[]        ciphertext = getCiphertext(plaintext, cipherKey, macKey);
 
-    return Base64.encodeBytes(ciphertext);
+    this.serialized = Base64.encodeBytes(ciphertext);
+  }
+
+  public EncryptedOutgoingMessage(String serialized) {
+    this.serialized = serialized;
+  }
+
+  public String serialize() {
+    return serialized;
   }
 
   private byte[] getCiphertext(byte[] plaintext, SecretKeySpec cipherKey, SecretKeySpec macKey)

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,
+                String publicKey, boolean lastResort)
-                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/limits/RateLimiter.java

@@ -16,13 +16,14 @@
  */
 package org.whispersystems.textsecuregcm.limits;
 
-import com.yammer.metrics.Metrics;
+import com.codahale.metrics.Meter;
-import com.yammer.metrics.core.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import net.spy.memcached.MemcachedClient;
-
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import org.whispersystems.textsecuregcm.util.Constants;
 
-import java.util.concurrent.TimeUnit;
+import static com.codahale.metrics.MetricRegistry.name;
 
 public class RateLimiter {
 
@@ -35,7 +36,9 @@ public class RateLimiter {
   public RateLimiter(MemcachedClient memcachedClient, String name,
                      int bucketSize, double leakRatePerMinute)
   {
-    this.meter             = Metrics.newMeter(RateLimiter.class, name, "exceeded", TimeUnit.MINUTES);
+    MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
+
+    this.meter             = metricRegistry.meter(name(getClass(), name, "exceeded"));
     this.memcachedClient   = memcachedClient;
     this.name              = name;
     this.bucketSize        = bucketSize;

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

@@ -1,13 +1,13 @@
 package org.whispersystems.textsecuregcm.metrics;
 
+import com.codahale.metrics.Gauge;
 import com.sun.management.OperatingSystemMXBean;
-import com.yammer.metrics.core.Gauge;
 
 import java.lang.management.ManagementFactory;
 
-public class CpuUsageGauge extends Gauge<Integer> {
+public class CpuUsageGauge implements Gauge<Integer> {
   @Override
-  public Integer value() {
+  public Integer getValue() {
     OperatingSystemMXBean mbean = (com.sun.management.OperatingSystemMXBean)
         ManagementFactory.getOperatingSystemMXBean();
 

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

@@ -1,17 +1,14 @@
 package org.whispersystems.textsecuregcm.metrics;
 
+import com.codahale.metrics.Gauge;
 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> {
+public class FreeMemoryGauge implements Gauge<Long> {
 
   @Override
-  public Long value() {
+  public Long getValue() {
     OperatingSystemMXBean mbean = (com.sun.management.OperatingSystemMXBean)
         ManagementFactory.getOperatingSystemMXBean();
 

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

@@ -1,23 +1,20 @@
 package org.whispersystems.textsecuregcm.metrics;
 
+import com.codahale.metrics.Counter;
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Histogram;
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.Metered;
+import com.codahale.metrics.MetricFilter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.ScheduledReporter;
+import com.codahale.metrics.Snapshot;
+import com.codahale.metrics.Timer;
 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 org.slf4j.Logger;
-import com.yammer.metrics.core.Counter;
+import org.slf4j.LoggerFactory;
-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;
@@ -33,292 +30,155 @@ import java.util.regex.Pattern;
 /**
  * Adapted from MetricsServlet.
  */
-public class JsonMetricsReporter extends AbstractPollingReporter implements MetricProcessor<JsonMetricsReporter.Context> {
+public class JsonMetricsReporter extends ScheduledReporter {
-  private final Clock clock = Clock.defaultClock();
+
-  private final VirtualMachineMetrics vm = VirtualMachineMetrics.getInstance();
+  private final Logger      logger  = LoggerFactory.getLogger(JsonMetricsReporter.class);
-  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(MetricRegistry registry, String token, String sunnylabsHost)
-
+      throws UnknownHostException
-  public JsonMetricsReporter(String service, MetricsRegistry registry, String token, String sunnylabsHost) throws UnknownHostException {
+  {
-    this(service, registry, token, sunnylabsHost, true);
+    super(registry, "jsonmetrics-reporter", MetricFilter.ALL, TimeUnit.SECONDS, TimeUnit.MILLISECONDS);
-  }
-
-  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() {
+  public void report(SortedMap<String, Gauge>     stringGaugeSortedMap,
+                     SortedMap<String, Counter>   stringCounterSortedMap,
+                     SortedMap<String, Histogram> stringHistogramSortedMap,
+                     SortedMap<String, Meter>     stringMeterSortedMap,
+                     SortedMap<String, Timer>     stringTimerSortedMap)
+  {
     try {
-      URL http = new URL("https", sunnylabsHost, 443, "/report/metrics?t=" + table + "&h=" + host);
+      logger.info("Reporting metrics...");
-      System.out.println("Reporting started to: " + http);
+      URL url = new URL("https", sunnylabsHost, 443, "/report/metrics?t=" + table + "&h=" + host);
-      HttpURLConnection urlc = (HttpURLConnection) http.openConnection();
+      HttpURLConnection connection = (HttpURLConnection) url.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 {
+      connection.setDoOutput(true);
-    final boolean showFullSamples;
+      connection.addRequestProperty("Content-Type", "application/json");
-    final JsonGenerator json;
 
-    Context(JsonGenerator json, boolean showFullSamples) {
+      OutputStream  outputStream = connection.getOutputStream();
-      this.json = json;
+      JsonGenerator json         = factory.createGenerator(outputStream, JsonEncoding.UTF8);
-      this.showFullSamples = showFullSamples;
-    }
-  }
 
-  public void writeJson(OutputStream out) throws IOException {
-    final JsonGenerator json = factory.createGenerator(out, JsonEncoding.UTF8);
       json.writeStartObject();
-    if (includeVMMetrics) {
+
-      writeVmMetrics(json);
+      for (Map.Entry<String, Gauge> gauge : stringGaugeSortedMap.entrySet()) {
+        reportGauge(json, gauge.getKey(), gauge.getValue());
       }
-    writeRegularMetrics(json, false);
+
+      for (Map.Entry<String, Counter> counter : stringCounterSortedMap.entrySet()) {
+        reportCounter(json, counter.getKey(), counter.getValue());
+      }
+
+      for (Map.Entry<String, Histogram> histogram : stringHistogramSortedMap.entrySet()) {
+        reportHistogram(json, histogram.getKey(), histogram.getValue());
+      }
+
+      for (Map.Entry<String, Meter> meter : stringMeterSortedMap.entrySet()) {
+        reportMeter(json, meter.getKey(), meter.getValue());
+      }
+
+      for (Map.Entry<String, Timer> timer : stringTimerSortedMap.entrySet()) {
+        reportTimer(json, timer.getKey(), timer.getValue());
+      }
+
       json.writeEndObject();
       json.close();
-  }
 
-  private void writeVmMetrics(JsonGenerator json) throws IOException {
+      outputStream.close();
-    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");
+      logger.info("Metrics server response: " + connection.getResponseCode());
-      json.writeStartObject();
+    } catch (IOException e) {
-      {
+      logger.warn("Error sending metrics", e);
-        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();
+      logger.warn("error", e);
-        }
-      }
     }
   }
 
-  @Override
+  private void reportGauge(JsonGenerator json, String name, Gauge gauge) throws IOException {
-  public void processHistogram(MetricName name, Histogram histogram, Context context) throws Exception {
+    Object gaugeValue = evaluateGauge(gauge);
-    final JsonGenerator json = context.json;
+
+    if (gaugeValue instanceof Number) {
+      json.writeFieldName(sanitize(name));
+      json.writeObject(gaugeValue);
+    }
+  }
+
+  private void reportCounter(JsonGenerator json, String name, Counter counter) throws IOException {
+    json.writeFieldName(sanitize(name));
+    json.writeNumber(counter.getCount());
+  }
+
+  private void reportHistogram(JsonGenerator json, String name, Histogram histogram) throws IOException {
+    Snapshot snapshot = histogram.getSnapshot();
+    json.writeFieldName(sanitize(name));
     json.writeStartObject();
-    {
+    json.writeNumberField("count", histogram.getCount());
-      json.writeNumberField("count", histogram.count());
+    writeSnapshot(json, snapshot);
-      writeSummarizable(histogram, json);
-      writeSampling(histogram, json);
-      if (context.showFullSamples) {
-        json.writeObjectField("values", histogram.getSnapshot().getValues());
-      }
-      histogram.clear();
-    }
     json.writeEndObject();
   }
 
-  @Override
+  private void reportMeter(JsonGenerator json, String name, Meter meter) throws IOException {
-  public void processCounter(MetricName name, Counter counter, Context context) throws Exception {
+    json.writeFieldName(sanitize(name));
-    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();
-    {
+    writeMetered(json, meter);
-      writeMeteredFields(meter, json);
-    }
     json.writeEndObject();
   }
 
-  @Override
+  private void reportTimer(JsonGenerator json, String name, Timer timer) throws IOException {
-  public void processTimer(MetricName name, Timer timer, Context context) throws Exception {
+    json.writeFieldName(sanitize(name));
-    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();
-      {
+    writeMetered(json, timer);
-        writeMeteredFields(timer, json);
+    json.writeEndObject();
-      }
+    json.writeFieldName("duration");
+    json.writeStartObject();
+    writeSnapshot(json, timer.getSnapshot());
     json.writeEndObject();
-    }
     json.writeEndObject();
   }
 
-  private static Object evaluateGauge(Gauge<?> gauge) {
+  private Object evaluateGauge(Gauge gauge) {
     try {
-      return gauge.value();
+      return gauge.getValue();
     } catch (RuntimeException e) {
-      return "error reading gauge: " + e.getMessage();
+      logger.warn("Error reading gauge", e);
+      return "error reading gauge";
     }
   }
 
-  private static void writeSummarizable(Summarizable metric, JsonGenerator json) throws IOException {
+  private void writeSnapshot(JsonGenerator json, Snapshot snapshot) throws IOException {
-    json.writeNumberField("min", metric.min());
+    json.writeNumberField("max", convertDuration(snapshot.getMax()));
-    json.writeNumberField("max", metric.max());
+    json.writeNumberField("mean", convertDuration(snapshot.getMean()));
-    json.writeNumberField("mean", metric.mean());
+    json.writeNumberField("min", convertDuration(snapshot.getMin()));
+    json.writeNumberField("stddev", convertDuration(snapshot.getStdDev()));
+    json.writeNumberField("median", convertDuration(snapshot.getMedian()));
+    json.writeNumberField("p75", convertDuration(snapshot.get75thPercentile()));
+    json.writeNumberField("p95", convertDuration(snapshot.get95thPercentile()));
+    json.writeNumberField("p98", convertDuration(snapshot.get98thPercentile()));
+    json.writeNumberField("p99", convertDuration(snapshot.get99thPercentile()));
+    json.writeNumberField("p999", convertDuration(snapshot.get999thPercentile()));
   }
 
-  private static void writeSampling(Sampling metric, JsonGenerator json) throws IOException {
+  private void writeMetered(JsonGenerator json, Metered meter) throws IOException {
-    final Snapshot snapshot = metric.getSnapshot();
+    json.writeNumberField("count", convertRate(meter.getCount()));
-    json.writeNumberField("median", snapshot.getMedian());
+    json.writeNumberField("mean", convertRate(meter.getMeanRate()));
-    json.writeNumberField("p75", snapshot.get75thPercentile());
+    json.writeNumberField("m1", convertRate(meter.getOneMinuteRate()));
-    json.writeNumberField("p95", snapshot.get95thPercentile());
+    json.writeNumberField("m5", convertRate(meter.getFiveMinuteRate()));
-    json.writeNumberField("p99", snapshot.get99thPercentile());
+    json.writeNumberField("m15", convertRate(meter.getFifteenMinuteRate()));
-    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) {
+  private String sanitize(String metricName) {
-    return SIMPLE_NAMES.matcher(metricName.getGroup() + "." + metricName.getName()).replaceAll("_");
+    return SIMPLE_NAMES.matcher(metricName).replaceAll("_");
   }
 
 }

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

@@ -1,16 +1,15 @@
 package org.whispersystems.textsecuregcm.metrics;
 
 
-import com.yammer.metrics.core.Gauge;
+import com.codahale.metrics.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> {
+public abstract class NetworkGauge implements Gauge<Long> {
 
   protected Pair<Long, Long> getSentReceived() throws IOException {
     File           proc          = new File("/proc/net/dev");

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

@@ -14,7 +14,7 @@ public class NetworkReceivedGauge extends NetworkGauge {
   private long lastReceived;
 
   @Override
-  public Long value() {
+  public Long getValue() {
     try {
       long             timestamp       = System.currentTimeMillis();
       Pair<Long, Long> sentAndReceived = getSentReceived();

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

@@ -14,7 +14,7 @@ public class NetworkSentGauge extends NetworkGauge {
   private long lastSent;
 
   @Override
-  public Long value() {
+  public Long getValue() {
     try {
       long             timestamp       = System.currentTimeMillis();
       Pair<Long, Long> sentAndReceived = getSentReceived();

src/main/java/org/whispersystems/textsecuregcm/providers/MemcacheHealthCheck.java

@@ -16,8 +16,7 @@
  */
 package org.whispersystems.textsecuregcm.providers;
 
-import com.yammer.metrics.core.HealthCheck;
+import com.codahale.metrics.health.HealthCheck;
-import com.yammer.metrics.core.HealthCheck.Result;
 import net.spy.memcached.MemcachedClient;
 
 import java.security.SecureRandom;
@@ -27,7 +26,6 @@ public class MemcacheHealthCheck extends HealthCheck {
   private final MemcachedClient client;
 
   public MemcacheHealthCheck(MemcachedClient client) {
-    super("memcached");
     this.client = client;
   }
 

src/main/java/org/whispersystems/textsecuregcm/providers/RedisHealthCheck.java

@@ -16,7 +16,7 @@
  */
 package org.whispersystems.textsecuregcm.providers;
 
-import com.yammer.metrics.core.HealthCheck;
+import com.codahale.metrics.health.HealthCheck;
 
 import redis.clients.jedis.Jedis;
 import redis.clients.jedis.JedisPool;
@@ -26,7 +26,6 @@ public class RedisHealthCheck extends HealthCheck {
   private final JedisPool clientPool;
 
   public RedisHealthCheck(JedisPool clientPool) {
-    super("redis");
     this.clientPool = clientPool;
   }
 

src/main/java/org/whispersystems/textsecuregcm/push/APNSender.java

@@ -16,18 +16,25 @@
  */
 package org.whispersystems.textsecuregcm.push;
 
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import com.google.common.base.Optional;
 import com.notnoop.apns.APNS;
 import com.notnoop.apns.ApnsService;
 import com.notnoop.exceptions.NetworkIOException;
-import com.yammer.metrics.Metrics;
-import com.yammer.metrics.core.Meter;
 import org.bouncycastle.openssl.PEMReader;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.entities.CryptoEncodingException;
 import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.PubSubManager;
+import org.whispersystems.textsecuregcm.storage.PubSubMessage;
+import org.whispersystems.textsecuregcm.storage.StoredMessages;
+import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.Util;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -40,21 +47,31 @@ import java.security.NoSuchAlgorithmException;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
-import java.util.concurrent.TimeUnit;
+
+import static com.codahale.metrics.MetricRegistry.name;
 
 public class APNSender {
 
-  private final Meter  success = Metrics.newMeter(APNSender.class, "sent", "success", TimeUnit.MINUTES);
+  private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-  private final Meter  failure = Metrics.newMeter(APNSender.class, "sent", "failure", TimeUnit.MINUTES);
+  private final Meter          websocketMeter = metricRegistry.meter(name(getClass(), "websocket"));
+  private final Meter          pushMeter      = metricRegistry.meter(name(getClass(), "push"));
+  private final Meter          failureMeter   = metricRegistry.meter(name(getClass(), "failure"));
   private final Logger         logger         = LoggerFactory.getLogger(APNSender.class);
 
   private static final String MESSAGE_BODY = "m";
 
   private final Optional<ApnsService> apnService;
+  private final PubSubManager         pubSubManager;
+  private final StoredMessages        storedMessages;
 
-  public APNSender(String apnCertificate, String apnKey)
+  public APNSender(PubSubManager pubSubManager,
+                   StoredMessages storedMessages,
+                   String apnCertificate, String apnKey)
       throws CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException
   {
+    this.pubSubManager  = pubSubManager;
+    this.storedMessages = storedMessages;
+
     if (!Util.isEmpty(apnCertificate) && !Util.isEmpty(apnKey)) {
       byte[] keyStore = initializeKeyStore(apnCertificate, apnKey);
       this.apnService = Optional.of(APNS.newService()
@@ -65,31 +82,44 @@ public class APNSender {
     }
   }
 
-  public void sendMessage(String registrationId, EncryptedOutgoingMessage message)
+  public void sendMessage(Account account, Device device,
+                          String registrationId, EncryptedOutgoingMessage message)
       throws TransientPushFailureException, NotPushRegisteredException
+  {
+    if (pubSubManager.publish(new WebsocketAddress(account.getId(), device.getId()),
+                              new PubSubMessage(PubSubMessage.TYPE_DELIVER, message.serialize())))
+    {
+      websocketMeter.mark();
+    } else {
+      storedMessages.insert(account.getId(), device.getId(), message.serialize());
+      sendPush(registrationId, message.serialize());
+    }
+  }
+
+  private void sendPush(String registrationId, String message)
+      throws TransientPushFailureException
   {
     try {
       if (!apnService.isPresent()) {
-        failure.mark();
+        failureMeter.mark();
         throw new TransientPushFailureException("APN access not configured!");
       }
 
       String payload = APNS.newPayload()
                            .alertBody("Message!")
-                           .customField(MESSAGE_BODY, message.serialize())
+                           .customField(MESSAGE_BODY, message)
                            .build();
 
       logger.debug("APN Payload: " + payload);
 
       apnService.get().push(registrationId, payload);
-      success.mark();
+      pushMeter.mark();
     } catch (NetworkIOException nioe) {
       logger.warn("Network Error", nioe);
-      failure.mark();
+      failureMeter.mark();
       throw new TransientPushFailureException(nioe);
-    } catch (CryptoEncodingException e) {
-      throw new NotPushRegisteredException(e);
     }
+
   }
 
   private static byte[] initializeKeyStore(String pemCertificate, String pemKey)

src/main/java/org/whispersystems/textsecuregcm/push/GCMSender.java

@@ -16,22 +16,24 @@
  */
 package org.whispersystems.textsecuregcm.push;
 
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import com.google.android.gcm.server.Constants;
 import com.google.android.gcm.server.Message;
 import com.google.android.gcm.server.Result;
 import com.google.android.gcm.server.Sender;
-import com.yammer.metrics.Metrics;
-import com.yammer.metrics.core.Meter;
-import org.whispersystems.textsecuregcm.entities.CryptoEncodingException;
 import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
 
 import java.io.IOException;
-import java.util.concurrent.TimeUnit;
+
+import static com.codahale.metrics.MetricRegistry.name;
 
 public class GCMSender {
 
-  private final Meter success = Metrics.newMeter(GCMSender.class, "sent", "success", TimeUnit.MINUTES);
+  private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(org.whispersystems.textsecuregcm.util.Constants.METRICS_NAME);
-  private final Meter failure = Metrics.newMeter(GCMSender.class, "sent", "failure", TimeUnit.MINUTES);
+  private final Meter          success        = metricRegistry.meter(name(getClass(), "sent", "success"));
+  private final Meter          failure        = metricRegistry.meter(name(getClass(), "sent", "failure"));
 
   private final Sender sender;
 
@@ -62,8 +64,6 @@ public class GCMSender {
       }
     } catch (IOException e) {
       throw new TransientPushFailureException(e);
-    } catch (CryptoEncodingException e) {
-      throw new NotPushRegisteredException(e);
     }
   }
 }

src/main/java/org/whispersystems/textsecuregcm/push/PushSender.java

@@ -26,7 +26,8 @@ import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
-import org.whispersystems.textsecuregcm.storage.StoredMessageManager;
+import org.whispersystems.textsecuregcm.storage.PubSubManager;
+import org.whispersystems.textsecuregcm.storage.StoredMessages;
 
 import java.io.IOException;
 import java.security.KeyStoreException;
@@ -40,29 +41,42 @@ public class PushSender {
   private final AccountsManager accounts;
   private final GCMSender       gcmSender;
   private final APNSender       apnSender;
-  private final StoredMessageManager storedMessageManager;
+  private final WebsocketSender webSocketSender;
 
   public PushSender(GcmConfiguration gcmConfiguration,
                     ApnConfiguration apnConfiguration,
-                    StoredMessageManager storedMessageManager,
+                    StoredMessages   storedMessages,
+                    PubSubManager    pubSubManager,
                     AccountsManager  accounts)
       throws CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException
   {
     this.accounts        = accounts;
-    this.storedMessageManager = storedMessageManager;
+    this.webSocketSender = new WebsocketSender(storedMessages, pubSubManager);
     this.gcmSender       = new GCMSender(gcmConfiguration.getApiKey());
-    this.apnSender            = new APNSender(apnConfiguration.getCertificate(), apnConfiguration.getKey());
+    this.apnSender       = new APNSender(pubSubManager, storedMessages,
+                                         apnConfiguration.getCertificate(),
+                                         apnConfiguration.getKey());
   }
 
-  public void sendMessage(Account account, Device device, MessageProtos.OutgoingMessageSignal outgoingMessage)
+  public void sendMessage(Account account, Device device, MessageProtos.OutgoingMessageSignal message)
       throws NotPushRegisteredException, TransientPushFailureException
   {
+    try {
       String                   signalingKey     = device.getSignalingKey();
-    EncryptedOutgoingMessage message      = new EncryptedOutgoingMessage(outgoingMessage, signalingKey);
+      EncryptedOutgoingMessage encryptedMessage = new EncryptedOutgoingMessage(message, signalingKey);
 
+      sendMessage(account, device, encryptedMessage);
+    } catch (CryptoEncodingException e) {
+      throw new NotPushRegisteredException(e);
+    }
+  }
+
+  public void sendMessage(Account account, Device device, EncryptedOutgoingMessage message)
+      throws NotPushRegisteredException, TransientPushFailureException
+  {
     if      (device.getGcmId() != null)   sendGcmMessage(account, device, message);
     else if (device.getApnId() != null)   sendApnMessage(account, device, message);
-    else if (device.getFetchesMessages()) storeFetchedMessage(account, device, message);
+    else if (device.getFetchesMessages()) sendWebSocketMessage(account, device, message);
     else                                  throw new NotPushRegisteredException("No delivery possible!");
   }
 
@@ -89,7 +103,7 @@ public class PushSender {
       throws TransientPushFailureException, NotPushRegisteredException
   {
     try {
-      apnSender.sendMessage(device.getApnId(), outgoingMessage);
+      apnSender.sendMessage(account, device, device.getApnId(), outgoingMessage);
     } catch (NotPushRegisteredException e) {
       device.setApnId(null);
       accounts.update(account);
@@ -97,11 +111,11 @@ public class PushSender {
     }
   }
 
-  private void storeFetchedMessage(Account account, Device device, EncryptedOutgoingMessage outgoingMessage)
+  private void sendWebSocketMessage(Account account, Device device, EncryptedOutgoingMessage outgoingMessage)
       throws NotPushRegisteredException
   {
     try {
-      storedMessageManager.storeMessage(account, device, outgoingMessage);
+      webSocketSender.sendMessage(account, device, outgoingMessage);
     } catch (CryptoEncodingException e) {
       throw new NotPushRegisteredException(e);
     }

src/main/java/org/whispersystems/textsecuregcm/push/WebsocketSender.java

@@ -0,0 +1,68 @@
+/**
+ * Copyright (C) 2014 Open WhisperSystems
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.whispersystems.textsecuregcm.push;
+
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
+import org.whispersystems.textsecuregcm.entities.CryptoEncodingException;
+import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.PubSubManager;
+import org.whispersystems.textsecuregcm.storage.PubSubMessage;
+import org.whispersystems.textsecuregcm.storage.StoredMessages;
+import org.whispersystems.textsecuregcm.util.Constants;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
+
+import java.util.List;
+
+import static com.codahale.metrics.MetricRegistry.name;
+
+public class WebsocketSender {
+
+  private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
+  private final Meter          onlineMeter    = metricRegistry.meter(name(getClass(), "online"));
+  private final Meter          offlineMeter   = metricRegistry.meter(name(getClass(), "offline"));
+
+  private final StoredMessages storedMessages;
+  private final PubSubManager  pubSubManager;
+
+  public WebsocketSender(StoredMessages storedMessages, PubSubManager pubSubManager) {
+    this.storedMessages = storedMessages;
+    this.pubSubManager  = pubSubManager;
+  }
+
+  public void sendMessage(Account account, Device device, EncryptedOutgoingMessage outgoingMessage)
+      throws CryptoEncodingException
+  {
+    sendMessage(account, device, outgoingMessage.serialize());
+  }
+
+  private void sendMessage(Account account, Device device, String serializedMessage) {
+    WebsocketAddress address       = new WebsocketAddress(account.getId(), device.getId());
+    PubSubMessage    pubSubMessage = new PubSubMessage(PubSubMessage.TYPE_DELIVER, serializedMessage);
+
+    if (pubSubManager.publish(address, pubSubMessage)) {
+      onlineMeter.mark();
+    } else {
+      offlineMeter.mark();
+      storedMessages.insert(account.getId(), device.getId(), serializedMessage);
+      pubSubManager.publish(address, new PubSubMessage(PubSubMessage.TYPE_QUERY_DB, null));
+    }
+  }
+}

src/main/java/org/whispersystems/textsecuregcm/sms/NexmoSmsSender.java

@@ -16,11 +16,13 @@
  */
 package org.whispersystems.textsecuregcm.sms;
 
-import com.yammer.metrics.Metrics;
+import com.codahale.metrics.Meter;
-import com.yammer.metrics.core.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.configuration.NexmoConfiguration;
+import org.whispersystems.textsecuregcm.util.Constants;
 
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -28,12 +30,14 @@ import java.io.InputStreamReader;
 import java.net.URL;
 import java.net.URLConnection;
 import java.net.URLEncoder;
-import java.util.concurrent.TimeUnit;
+
+import static com.codahale.metrics.MetricRegistry.name;
 
 public class NexmoSmsSender {
 
-  private final Meter  smsMeter = Metrics.newMeter(NexmoSmsSender.class, "sms", "delivered", TimeUnit.MINUTES);
+  private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-  private final Meter  voxMeter = Metrics.newMeter(NexmoSmsSender.class, "vox", "delivered", TimeUnit.MINUTES);
+  private final Meter          smsMeter       = metricRegistry.meter(name(getClass(), "sms", "delivered"));
+  private final Meter          voxMeter       = metricRegistry.meter(name(getClass(), "vox", "delivered"));
   private final Logger         logger         = LoggerFactory.getLogger(NexmoSmsSender.class);
 
   private static final String NEXMO_SMS_URL =

src/main/java/org/whispersystems/textsecuregcm/sms/TwilioSmsSender.java

@@ -16,22 +16,25 @@
  */
 package org.whispersystems.textsecuregcm.sms;
 
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.SharedMetricRegistries;
 import com.twilio.sdk.TwilioRestClient;
 import com.twilio.sdk.TwilioRestException;
 import com.twilio.sdk.resource.factory.CallFactory;
 import com.twilio.sdk.resource.factory.MessageFactory;
-import com.yammer.metrics.Metrics;
-import com.yammer.metrics.core.Meter;
 import org.apache.http.NameValuePair;
 import org.apache.http.message.BasicNameValuePair;
 import org.whispersystems.textsecuregcm.configuration.TwilioConfiguration;
+import org.whispersystems.textsecuregcm.util.Constants;
 
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
+
+import static com.codahale.metrics.MetricRegistry.name;
 
 public class TwilioSmsSender {
 
@@ -40,8 +43,9 @@ public class TwilioSmsSender {
                                          "    <Say voice=\"woman\" language=\"en\">" + SmsSender.VOX_VERIFICATION_TEXT + "%s</Say>\n" +
                                          "</Response>";
 
-  private final Meter smsMeter = Metrics.newMeter(TwilioSmsSender.class, "sms", "delivered", TimeUnit.MINUTES);
+  private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-  private final Meter voxMeter = Metrics.newMeter(TwilioSmsSender.class, "vox", "delivered", TimeUnit.MINUTES);
+  private final Meter          smsMeter       = metricRegistry.meter(name(getClass(), "sms", "delivered"));
+  private final Meter          voxMeter       = metricRegistry.meter(name(getClass(), "vox", "delivered"));
 
   private final String accountId;
   private final String accountToken;

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 " +
+  @SqlBatch("INSERT INTO keys (number, device_id, key_id, public_key, last_resort) VALUES " +
-      "(:number, :device_id, :key_id, :public_key, :identity_key, :last_resort)")
+      "(: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 " +
+  @SqlUpdate("INSERT INTO keys (number, device_id, key_id, public_key, last_resort) VALUES " +
-      "(:number, :device_id, :key_id, :public_key, :identity_key, :last_resort)")
+      "(: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);
     }
   }

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

@@ -134,6 +134,7 @@ public class PubSubManager {
     public void onSubscribe(String channel, int count) {
       try {
         WebsocketAddress address = new WebsocketAddress(channel);
+
         if (address.getAccountId() == 0 && address.getDeviceId() == 0) {
           synchronized (PubSubManager.this) {
             subscribed = true;

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

@@ -1,66 +0,0 @@
-/**
- * Copyright (C) 2014 Open WhisperSystems
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU Affero General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU Affero General Public License for more details.
- *
- * You should have received a copy of the GNU Affero General Public License
- * along with this program.  If not, see <http://www.gnu.org/licenses/>.
- */
-package org.whispersystems.textsecuregcm.storage;
-
-import org.whispersystems.textsecuregcm.entities.CryptoEncodingException;
-import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
-import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
-
-import java.util.List;
-
-public class StoredMessageManager {
-
-  private final StoredMessages storedMessages;
-  private final PubSubManager  pubSubManager;
-
-  public StoredMessageManager(StoredMessages storedMessages, PubSubManager pubSubManager) {
-    this.storedMessages = storedMessages;
-    this.pubSubManager  = pubSubManager;
-  }
-
-  public void storeMessage(Account account, Device device, EncryptedOutgoingMessage outgoingMessage)
-      throws CryptoEncodingException
-  {
-    storeMessage(account, device, outgoingMessage.serialize());
-  }
-
-  public void storeMessages(Account account, Device device, List<String> serializedMessages) {
-    for (String serializedMessage : serializedMessages) {
-      storeMessage(account, device, serializedMessage);
-    }
-  }
-
-  private void storeMessage(Account account, Device device, String serializedMessage) {
-    if (device.getFetchesMessages()) {
-      WebsocketAddress address       = new WebsocketAddress(account.getId(), device.getId());
-      PubSubMessage    pubSubMessage = new PubSubMessage(PubSubMessage.TYPE_DELIVER, serializedMessage);
-
-      if (!pubSubManager.publish(address, pubSubMessage)) {
-        storedMessages.insert(account.getId(), device.getId(), serializedMessage);
-        pubSubManager.publish(address, new PubSubMessage(PubSubMessage.TYPE_QUERY_DB, null));
-      }
-
-      return;
-    }
-
-    storedMessages.insert(account.getId(), device.getId(), serializedMessage);
-  }
-
-  public List<String> getOutgoingMessages(Account account, Device device) {
-    return storedMessages.getMessagesForDevice(account.getId(), device.getId());
-  }
-}

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

@@ -16,21 +16,71 @@
  */
 package org.whispersystems.textsecuregcm.storage;
 
-import org.skife.jdbi.v2.sqlobject.Bind;
+import com.codahale.metrics.Histogram;
-import org.skife.jdbi.v2.sqlobject.SqlBatch;
+import com.codahale.metrics.MetricRegistry;
-import org.skife.jdbi.v2.sqlobject.SqlQuery;
+import com.codahale.metrics.SharedMetricRegistries;
-import org.skife.jdbi.v2.sqlobject.SqlUpdate;
 
+import org.whispersystems.textsecuregcm.util.Constants;
+
+import java.util.LinkedList;
 import java.util.List;
 
-public interface StoredMessages {
+import static com.codahale.metrics.MetricRegistry.name;
+import redis.clients.jedis.Jedis;
+import redis.clients.jedis.JedisPool;
 
-  @SqlUpdate("INSERT INTO messages (account_id, device_id, encrypted_message) VALUES (:account_id, :device_id, :encrypted_message)")
+public class StoredMessages {
-  void insert(@Bind("account_id") long accountId, @Bind("device_id") long deviceId, @Bind("encrypted_message") String encryptedOutgoingMessage);
 
-  @SqlBatch("INSERT INTO messages (account_id, device_id, encrypted_message) VALUES (:account_id, :device_id, :encrypted_message)")
+  private final MetricRegistry metricRegistry     = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
-  void insert(@Bind("account_id") long accountId, @Bind("device_id") long deviceId, @Bind("encrypted_message") List<String> encryptedOutgoingMessages);
+  private final Histogram      queueSizeHistogram = metricRegistry.histogram(name(getClass(), "queue_size"));
+
+  private static final String QUEUE_PREFIX = "msgs";
+
+  private final JedisPool jedisPool;
+
+  public StoredMessages(JedisPool jedisPool) {
+    this.jedisPool = jedisPool;
+  }
+
+  public void insert(long accountId, long deviceId, String message) {
+    Jedis jedis = null;
+
+    try {
+      jedis = jedisPool.getResource();
+
+      long queueSize = jedis.lpush(getKey(accountId, deviceId), message);
+      queueSizeHistogram.update(queueSize);
+
+      if (queueSize > 1000) {
+        jedis.ltrim(getKey(accountId, deviceId), 0, 999);
+      }
+    } finally {
+      if (jedis != null)
+        jedisPool.returnResource(jedis);
+    }
+  }
+
+  public List<String> getMessagesForDevice(long accountId, long deviceId) {
+    List<String> messages = new LinkedList<>();
+    Jedis        jedis    = null;
+
+    try {
+      jedis = jedisPool.getResource();
+      String message;
+
+      while ((message = jedis.rpop(QUEUE_PREFIX + accountId + ":" + deviceId)) != null) {
+        messages.add(message);
+      }
+
+      return messages;
+    } finally {
+      if (jedis != null)
+        jedisPool.returnResource(jedis);
+    }
+  }
+
+  private String getKey(long accountId, long deviceId) {
+    return QUEUE_PREFIX + ":" + accountId + ":" + deviceId;
+  }
 
-  @SqlQuery("DELETE FROM messages WHERE account_id = :account_id AND device_id = :device_id RETURNING encrypted_message")
-  List<String> getMessagesForDevice(@Bind("account_id") long accountId, @Bind("device_id") long deviceId);
 }

src/main/java/org/whispersystems/textsecuregcm/util/Base64.java

@@ -1241,7 +1241,7 @@ public class Base64
      * @since 1.4
      */
     public static byte[] decode( String s ) throws java.io.IOException {
-        return decode( s, NO_OPTIONS );
+        return decode( s, DONT_GUNZIP );
     }
 
     

src/main/java/org/whispersystems/textsecuregcm/util/Constants.java

@@ -0,0 +1,7 @@
+package org.whispersystems.textsecuregcm.util;
+
+public class Constants {
+
+  public static final String METRICS_NAME = "textsecure";
+
+}

src/main/java/org/whispersystems/textsecuregcm/websocket/WebsocketControllerFactory.java

@@ -0,0 +1,47 @@
+package org.whispersystems.textsecuregcm.websocket;
+
+import org.eclipse.jetty.websocket.api.UpgradeRequest;
+import org.eclipse.jetty.websocket.api.UpgradeResponse;
+import org.eclipse.jetty.websocket.servlet.WebSocketCreator;
+import org.eclipse.jetty.websocket.servlet.WebSocketServlet;
+import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
+import org.whispersystems.textsecuregcm.controllers.WebsocketController;
+import org.whispersystems.textsecuregcm.push.PushSender;
+import org.whispersystems.textsecuregcm.push.WebsocketSender;
+import org.whispersystems.textsecuregcm.storage.PubSubManager;
+import org.whispersystems.textsecuregcm.storage.StoredMessages;
+
+
+public class WebsocketControllerFactory extends WebSocketServlet implements WebSocketCreator {
+
+  private final Logger logger = LoggerFactory.getLogger(WebsocketControllerFactory.class);
+
+  private final PushSender           pushSender;
+  private final StoredMessages       storedMessages;
+  private final PubSubManager        pubSubManager;
+  private final AccountAuthenticator accountAuthenticator;
+
+  public WebsocketControllerFactory(AccountAuthenticator accountAuthenticator,
+                                    PushSender           pushSender,
+                                    StoredMessages       storedMessages,
+                                    PubSubManager        pubSubManager)
+  {
+    this.accountAuthenticator = accountAuthenticator;
+    this.pushSender           = pushSender;
+    this.storedMessages       = storedMessages;
+    this.pubSubManager        = pubSubManager;
+  }
+
+  @Override
+  public void configure(WebSocketServletFactory factory) {
+    factory.setCreator(this);
+  }
+
+  @Override
+  public Object createWebSocket(UpgradeRequest upgradeRequest, UpgradeResponse upgradeResponse) {
+    return new WebsocketController(accountAuthenticator, pushSender, pubSubManager, storedMessages);
+  }
+}

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

@@ -16,13 +16,6 @@
  */
 package org.whispersystems.textsecuregcm.workers;
 
-import com.yammer.dropwizard.cli.ConfiguredCommand;
-import com.yammer.dropwizard.config.Bootstrap;
-import com.yammer.dropwizard.db.DatabaseConfiguration;
-import com.yammer.dropwizard.jdbi.ImmutableListContainerFactory;
-import com.yammer.dropwizard.jdbi.ImmutableSetContainerFactory;
-import com.yammer.dropwizard.jdbi.OptionalContainerFactory;
-import com.yammer.dropwizard.jdbi.args.OptionalArgumentFactory;
 import net.sourceforge.argparse4j.inf.Namespace;
 import net.spy.memcached.MemcachedClient;
 import org.skife.jdbi.v2.DBI;
@@ -36,6 +29,13 @@ import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.DirectoryManager;
 
+import io.dropwizard.cli.ConfiguredCommand;
+import io.dropwizard.db.DataSourceFactory;
+import io.dropwizard.jdbi.ImmutableListContainerFactory;
+import io.dropwizard.jdbi.ImmutableSetContainerFactory;
+import io.dropwizard.jdbi.OptionalContainerFactory;
+import io.dropwizard.jdbi.args.OptionalArgumentFactory;
+import io.dropwizard.setup.Bootstrap;
 import redis.clients.jedis.JedisPool;
 
 public class DirectoryCommand extends ConfiguredCommand<WhisperServerConfiguration> {
@@ -53,7 +53,7 @@ public class DirectoryCommand extends ConfiguredCommand<WhisperServerConfigurati
       throws Exception
   {
     try {
-      DatabaseConfiguration dbConfig = config.getDatabaseConfiguration();
+      DataSourceFactory dbConfig = config.getDataSourceFactory();
       DBI               dbi      = new DBI(dbConfig.getUrl(), dbConfig.getUser(), dbConfig.getPassword());
 
       dbi.registerArgumentFactory(new OptionalArgumentFactory(dbConfig.getDriverClass()));

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/AccountControllerTest.java

@@ -2,7 +2,8 @@ package org.whispersystems.textsecuregcm.tests.controllers;
 
 import com.google.common.base.Optional;
 import com.sun.jersey.api.client.ClientResponse;
-import com.yammer.dropwizard.testing.ResourceTest;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.controllers.AccountController;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
@@ -16,11 +17,12 @@ import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 
 import javax.ws.rs.core.MediaType;
 
+import io.dropwizard.testing.junit.ResourceTestRule;
 import static org.fest.assertions.api.Assertions.assertThat;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.*;
 
-public class AccountControllerTest extends ResourceTest {
+public class AccountControllerTest {
 
   private static final String SENDER = "+14152222222";
 
@@ -30,26 +32,29 @@ public class AccountControllerTest extends ResourceTest {
   private RateLimiter            rateLimiter            = mock(RateLimiter.class           );
   private SmsSender              smsSender              = mock(SmsSender.class             );
 
-  @Override
+  @Rule
-  protected void setUpResources() throws Exception {
+  public final ResourceTestRule resources = ResourceTestRule.builder()
-    addProvider(AuthHelper.getAuthenticator());
+                                                            .addProvider(AuthHelper.getAuthenticator())
+                                                            .addResource(new AccountController(pendingAccountsManager,
+                                                                                               accountsManager,
+                                                                                               rateLimiters,
+                                                                                               smsSender))
+                                                            .build();
 
+
+  @Before
+  public void setup() throws Exception {
     when(rateLimiters.getSmsDestinationLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getVoiceDestinationLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getVerifyLimiter()).thenReturn(rateLimiter);
 
     when(pendingAccountsManager.getCodeForNumber(SENDER)).thenReturn(Optional.of("1234"));
-
-    addResource(new AccountController(pendingAccountsManager,
-                                      accountsManager,
-                                      rateLimiters,
-                                      smsSender));
   }
 
   @Test
   public void testSendCode() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/accounts/sms/code/%s", SENDER))
+        resources.client().resource(String.format("/v1/accounts/sms/code/%s", SENDER))
             .get(ClientResponse.class);
 
     assertThat(response.getStatus()).isEqualTo(200);
@@ -60,7 +65,7 @@ public class AccountControllerTest extends ResourceTest {
   @Test
   public void testVerifyCode() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/accounts/code/%s", "1234"))
+        resources.client().resource(String.format("/v1/accounts/code/%s", "1234"))
             .header("Authorization", AuthHelper.getAuthHeader(SENDER, "bar"))
             .entity(new AccountAttributes("keykeykeykey", false, false, 2222))
             .type(MediaType.APPLICATION_JSON_TYPE)
@@ -68,13 +73,13 @@ public class AccountControllerTest extends ResourceTest {
 
     assertThat(response.getStatus()).isEqualTo(204);
 
-    verify(accountsManager).create(isA(Account.class));
+    verify(accountsManager, times(1)).create(isA(Account.class));
   }
 
   @Test
   public void testVerifyBadCode() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/accounts/code/%s", "1111"))
+        resources.client().resource(String.format("/v1/accounts/code/%s", "1111"))
             .header("Authorization", AuthHelper.getAuthHeader(SENDER, "bar"))
             .entity(new AccountAttributes("keykeykeykey", false, false, 3333))
             .type(MediaType.APPLICATION_JSON_TYPE)

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

@@ -17,7 +17,8 @@
 package org.whispersystems.textsecuregcm.tests.controllers;
 
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.testing.ResourceTest;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.controllers.DeviceController;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
@@ -33,10 +34,11 @@ import org.whispersystems.textsecuregcm.util.VerificationCode;
 import javax.ws.rs.Path;
 import javax.ws.rs.core.MediaType;
 
+import io.dropwizard.testing.junit.ResourceTestRule;
 import static org.fest.assertions.api.Assertions.assertThat;
 import static org.mockito.Mockito.*;
 
-public class DeviceControllerTest extends ResourceTest {
+public class DeviceControllerTest {
   @Path("/v1/devices")
   static class DumbVerificationDeviceController extends DeviceController {
     public DumbVerificationDeviceController(PendingDevicesManager pendingDevices, AccountsManager accounts, RateLimiters rateLimiters) {
@@ -55,10 +57,17 @@ public class DeviceControllerTest extends ResourceTest {
   private RateLimiter           rateLimiter           = mock(RateLimiter.class           );
   private Account               account               = mock(Account.class               );
 
-  @Override
+  @Rule
-  protected void setUpResources() throws Exception {
+  public final ResourceTestRule resources = ResourceTestRule.builder()
-    addProvider(AuthHelper.getAuthenticator());
+                                                            .addProvider(AuthHelper.getAuthenticator())
+                                                            .addResource(new DumbVerificationDeviceController(pendingDevicesManager,
+                                                                                                              accountsManager,
+                                                                                                              rateLimiters))
+                                                            .build();
 
+
+  @Before
+  public void setup() throws Exception {
     when(rateLimiters.getSmsDestinationLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getVoiceDestinationLimiter()).thenReturn(rateLimiter);
     when(rateLimiters.getVerifyLimiter()).thenReturn(rateLimiter);
@@ -69,19 +78,17 @@ public class DeviceControllerTest extends ResourceTest {
 
     when(pendingDevicesManager.getCodeForNumber(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of("5678901"));
     when(accountsManager.get(AuthHelper.VALID_NUMBER)).thenReturn(Optional.of(account));
-
-    addResource(new DumbVerificationDeviceController(pendingDevicesManager, accountsManager, rateLimiters));
   }
 
   @Test
   public void validDeviceRegisterTest() throws Exception {
-    VerificationCode deviceCode = client().resource("/v1/devices/provisioning_code")
+    VerificationCode deviceCode = resources.client().resource("/v1/devices/provisioning_code")
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
         .get(VerificationCode.class);
 
     assertThat(deviceCode).isEqualTo(new VerificationCode(5678901));
 
-    DeviceResponse response = client().resource("/v1/devices/5678901")
+    DeviceResponse response = resources.client().resource("/v1/devices/5678901")
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, "password1"))
         .entity(new AccountAttributes("keykeykeykey", false, true, 1234))
         .type(MediaType.APPLICATION_JSON_TYPE)

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

@@ -1,7 +1,9 @@
 package org.whispersystems.textsecuregcm.tests.controllers;
 
 import com.sun.jersey.api.client.ClientResponse;
-import com.yammer.dropwizard.testing.ResourceTest;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
@@ -17,21 +19,28 @@ import javax.ws.rs.core.MediaType;
 import java.util.LinkedList;
 import java.util.List;
 
+import io.dropwizard.testing.junit.ResourceTestRule;
 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 {
+public class DirectoryControllerTest {
 
-  private RateLimiters     rateLimiters     = mock(RateLimiters.class    );
+  private final RateLimiters     rateLimiters     = mock(RateLimiters.class    );
-  private RateLimiter      rateLimiter      = mock(RateLimiter.class     );
+  private final RateLimiter      rateLimiter      = mock(RateLimiter.class     );
-  private DirectoryManager directoryManager = mock(DirectoryManager.class);
+  private final DirectoryManager directoryManager = mock(DirectoryManager.class);
 
-  @Override
+  @Rule
-  protected void setUpResources() throws Exception {
+  public final ResourceTestRule resources = ResourceTestRule.builder()
-    addProvider(AuthHelper.getAuthenticator());
+                                                            .addProvider(AuthHelper.getAuthenticator())
+                                                            .addResource(new DirectoryController(rateLimiters,
+                                                                                                 directoryManager))
+                                                            .build();
 
+
+  @Before
+  public void setup() throws Exception {
     when(rateLimiters.getContactsLimiter()).thenReturn(rateLimiter);
     when(directoryManager.get(anyList())).thenAnswer(new Answer<List<byte[]>>() {
       @Override
@@ -42,8 +51,6 @@ public class DirectoryControllerTest extends ResourceTest {
         return response;
       }
     });
-
-    addResource(new DirectoryController(rateLimiters, directoryManager));
   }
 
   @Test
@@ -58,7 +65,7 @@ public class DirectoryControllerTest extends ResourceTest {
     expectedResponse.remove(0);
 
     ClientResponse response =
-        client().resource("/v1/directory/tokens/")
+        resources.client().resource("/v1/directory/tokens/")
                  .entity(new ClientContactTokens(tokens))
                  .type(MediaType.APPLICATION_JSON_TYPE)
                  .header("Authorization",
@@ -67,7 +74,6 @@ public class DirectoryControllerTest extends ResourceTest {
                  .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

@@ -4,7 +4,8 @@ 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.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.controllers.FederationController;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
@@ -23,17 +24,16 @@ import javax.ws.rs.core.MediaType;
 import java.util.LinkedList;
 import java.util.List;
 
-import static com.yammer.dropwizard.testing.JsonHelpers.jsonFixture;
+import io.dropwizard.testing.junit.ResourceTestRule;
 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.*;
-import static org.mockito.Mockito.verify;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.jsonFixture;
-import static org.mockito.Mockito.when;
 
-public class FederatedControllerTest extends ResourceTest {
+public class FederatedControllerTest {
 
   private static final String SINGLE_DEVICE_RECIPIENT = "+14151111111";
   private static final String MULTI_DEVICE_RECIPIENT  = "+14152222222";
@@ -46,10 +46,20 @@ public class FederatedControllerTest extends ResourceTest {
 
   private final ObjectMapper mapper = new ObjectMapper();
 
-  @Override
+  private final MessageController messageController = new MessageController(rateLimiters, pushSender, accountsManager, federatedClientManager);
-  protected void setUpResources() throws Exception {
-    addProvider(AuthHelper.getAuthenticator());
 
+  @Rule
+  public final ResourceTestRule resources = ResourceTestRule.builder()
+                                                            .addProvider(AuthHelper.getAuthenticator())
+                                                            .addResource(new FederationController(accountsManager,
+                                                                                                  null, null,
+                                                                                                  messageController))
+                                                            .build();
+
+
+
+  @Before
+  public void setup() throws Exception {
     List<Device> singleDeviceList = new LinkedList<Device>() {{
       add(new Device(1, "foo", "bar", "baz", "isgcm", null, false, 111));
     }};
@@ -66,15 +76,12 @@ public class FederatedControllerTest extends ResourceTest {
     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))
+        resources.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)

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

@@ -2,9 +2,12 @@ package org.whispersystems.textsecuregcm.tests.controllers;
 
 import com.google.common.base.Optional;
 import com.sun.jersey.api.client.ClientResponse;
-import com.yammer.dropwizard.testing.ResourceTest;
+import org.junit.Before;
+import org.junit.Rule;
 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;
@@ -22,32 +25,36 @@ import javax.ws.rs.core.MediaType;
 import java.util.LinkedList;
 import java.util.List;
 
+import io.dropwizard.testing.junit.ResourceTestRule;
 import static org.fest.assertions.api.Assertions.assertThat;
 import static org.mockito.Mockito.*;
 
-public class KeyControllerTest extends ResourceTest {
+public class KeyControllerTest {
 
-  private final String EXISTS_NUMBER     = "+14152222222";
+  private static final String EXISTS_NUMBER     = "+14152222222";
-  private final String NOT_EXISTS_NUMBER = "+14152222220";
+  private static String NOT_EXISTS_NUMBER = "+14152222220";
 
-  private final int SAMPLE_REGISTRATION_ID  =  999;
+  private static int SAMPLE_REGISTRATION_ID  =  999;
-  private final int SAMPLE_REGISTRATION_ID2 = 1002;
+  private static 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);
   private final Account         existsAccount = mock(Account.class        );
 
+  private RateLimiters          rateLimiters  = mock(RateLimiters.class);
+  private RateLimiter           rateLimiter   = mock(RateLimiter.class );
 
-  @Override
+  @Rule
-  protected void setUpResources() {
+  public final ResourceTestRule resources = ResourceTestRule.builder()
-    addProvider(AuthHelper.getAuthenticator());
+                                                            .addProvider(AuthHelper.getAuthenticator())
-
+                                                            .addResource(new KeysController(rateLimiters, keys, accounts, null))
-    RateLimiters    rateLimiters = mock(RateLimiters.class);
+                                                            .build();
-    RateLimiter     rateLimiter  = mock(RateLimiter.class );
 
+  @Before
+  public void setup() {
     Device  sampleDevice  = mock(Device.class );
     Device  sampleDevice2 = mock(Device.class);
     Device  sampleDevice3 = mock(Device.class);
@@ -63,31 +70,41 @@ 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());
 
+    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(SAMPLE_KEY);
+        allKeys.add(cloneKey(SAMPLE_KEY));
-    allKeys.add(SAMPLE_KEY2);
+        allKeys.add(cloneKey(SAMPLE_KEY2));
-    allKeys.add(SAMPLE_KEY3);
+        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);
-
-    addResource(new KeysController(rateLimiters, keys, accounts, null));
   }
 
   @Test
   public void validKeyStatusTest() throws Exception {
-    PreKeyStatus result = client().resource("/v1/keys")
+    PreKeyStatus result = resources.client().resource("/v1/keys")
         .header("Authorization",
                 AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
         .get(PreKeyStatus.class);
@@ -99,13 +116,13 @@ public class KeyControllerTest extends ResourceTest {
 
   @Test
   public void validLegacyRequestTest() throws Exception {
-    PreKey result = client().resource(String.format("/v1/keys/%s", EXISTS_NUMBER))
+    PreKey result = resources.client().resource(String.format("/v1/keys/%s", EXISTS_NUMBER))
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
         .get(PreKey.class);
 
     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);
@@ -116,7 +133,7 @@ public class KeyControllerTest extends ResourceTest {
 
   @Test
   public void validMultiRequestTest() throws Exception {
-    UnstructuredPreKeyList results = client().resource(String.format("/v1/keys/%s/*", EXISTS_NUMBER))
+    UnstructuredPreKeyList results = resources.client().resource(String.format("/v1/keys/%s/*", EXISTS_NUMBER))
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
         .get(UnstructuredPreKeyList.class);
 
@@ -126,7 +143,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 +152,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);
@@ -148,27 +165,27 @@ public class KeyControllerTest extends ResourceTest {
 
   @Test
   public void invalidRequestTest() throws Exception {
-    ClientResponse response = client().resource(String.format("/v1/keys/%s", NOT_EXISTS_NUMBER))
+    ClientResponse response = resources.client().resource(String.format("/v1/keys/%s", NOT_EXISTS_NUMBER))
         .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
         .get(ClientResponse.class);
 
-    assertThat(response.getClientResponseStatus().getStatusCode()).isEqualTo(404);
+    assertThat(response.getStatusInfo().getStatusCode()).isEqualTo(404);
   }
 
   @Test
   public void unauthorizedRequestTest() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/keys/%s", NOT_EXISTS_NUMBER))
+        resources.client().resource(String.format("/v1/keys/%s", NOT_EXISTS_NUMBER))
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.INVALID_PASSWORD))
             .get(ClientResponse.class);
 
-    assertThat(response.getClientResponseStatus().getStatusCode()).isEqualTo(401);
+    assertThat(response.getStatusInfo().getStatusCode()).isEqualTo(401);
 
     response =
-        client().resource(String.format("/v1/keys/%s", NOT_EXISTS_NUMBER))
+        resources.client().resource(String.format("/v1/keys/%s", NOT_EXISTS_NUMBER))
             .get(ClientResponse.class);
 
-    assertThat(response.getClientResponseStatus().getStatusCode()).isEqualTo(401);
+    assertThat(response.getStatusInfo().getStatusCode()).isEqualTo(401);
   }
 
   @Test
@@ -185,7 +202,7 @@ public class KeyControllerTest extends ResourceTest {
     preKeyList.setLastResortKey(lastResortKey);
 
     ClientResponse response =
-        client().resource("/v1/keys")
+        resources.client().resource("/v1/keys")
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
             .type(MediaType.APPLICATION_JSON_TYPE)
             .put(ClientResponse.class, preKeyList);
@@ -209,4 +226,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());
+  }
+
 }

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

@@ -3,10 +3,11 @@ 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.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
-import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.entities.IncomingMessageList;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.entities.MismatchedDevices;
@@ -24,35 +25,39 @@ import javax.ws.rs.core.MediaType;
 import java.util.LinkedList;
 import java.util.List;
 
-import static com.yammer.dropwizard.testing.JsonHelpers.asJson;
+import io.dropwizard.testing.junit.ResourceTestRule;
-import static com.yammer.dropwizard.testing.JsonHelpers.jsonFixture;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.is;
 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.*;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.asJson;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.jsonFixture;
 
-public class MessageControllerTest extends ResourceTest {
+public class MessageControllerTest {
 
   private static final String SINGLE_DEVICE_RECIPIENT = "+14151111111";
   private static final String MULTI_DEVICE_RECIPIENT  = "+14152222222";
 
-  private PushSender             pushSender             = mock(PushSender.class            );
+  private  final PushSender             pushSender             = mock(PushSender.class            );
-  private FederatedClientManager federatedClientManager = mock(FederatedClientManager.class);
+  private  final FederatedClientManager federatedClientManager = mock(FederatedClientManager.class);
-  private AccountsManager        accountsManager        = mock(AccountsManager.class       );
+  private  final AccountsManager        accountsManager        = mock(AccountsManager.class       );
-  private RateLimiters           rateLimiters           = mock(RateLimiters.class          );
+  private  final RateLimiters           rateLimiters           = mock(RateLimiters.class          );
-  private RateLimiter            rateLimiter            = mock(RateLimiter.class           );
+  private  final RateLimiter            rateLimiter            = mock(RateLimiter.class           );
 
   private  final ObjectMapper mapper = new ObjectMapper();
 
-  @Override
+  @Rule
-  protected void setUpResources() throws Exception {
+  public final ResourceTestRule resources = ResourceTestRule.builder()
-    addProvider(AuthHelper.getAuthenticator());
+                                                            .addProvider(AuthHelper.getAuthenticator())
+                                                            .addResource(new MessageController(rateLimiters, pushSender, accountsManager,
+                                                                                               federatedClientManager))
+                                                            .build();
 
+
+  @Before
+  public void setup() throws Exception {
     List<Device> singleDeviceList = new LinkedList<Device>() {{
       add(new Device(1, "foo", "bar", "baz", "isgcm", null, false, 111));
     }};
@@ -69,15 +74,12 @@ public class MessageControllerTest extends ResourceTest {
     when(accountsManager.get(eq(MULTI_DEVICE_RECIPIENT))).thenReturn(Optional.of(multiDeviceAccount));
 
     when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter);
-
-    addResource(new MessageController(rateLimiters, pushSender, accountsManager,
-                                      federatedClientManager));
   }
 
   @Test
-  public void testSingleDeviceLegacy() throws Exception {
+  public synchronized void testSingleDeviceLegacy() throws Exception {
     ClientResponse response =
-        client().resource("/v1/messages/")
+        resources.client().resource("/v1/messages/")
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
             .entity(mapper.readValue(jsonFixture("fixtures/legacy_message_single_device.json"), IncomingMessageList.class))
             .type(MediaType.APPLICATION_JSON_TYPE)
@@ -85,13 +87,13 @@ public class MessageControllerTest extends ResourceTest {
 
     assertThat("Good Response", response.getStatus(), is(equalTo(200)));
 
-    verify(pushSender).sendMessage(any(Account.class), any(Device.class), any(MessageProtos.OutgoingMessageSignal.class));
+    verify(pushSender, times(1)).sendMessage(any(Account.class), any(Device.class), any(MessageProtos.OutgoingMessageSignal.class));
   }
 
   @Test
-  public void testSingleDeviceCurrent() throws Exception {
+  public synchronized void testSingleDeviceCurrent() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/messages/%s", SINGLE_DEVICE_RECIPIENT))
+        resources.client().resource(String.format("/v1/messages/%s", SINGLE_DEVICE_RECIPIENT))
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
         .entity(mapper.readValue(jsonFixture("fixtures/current_message_single_device.json"), IncomingMessageList.class))
         .type(MediaType.APPLICATION_JSON_TYPE)
@@ -99,13 +101,13 @@ public class MessageControllerTest extends ResourceTest {
 
     assertThat("Good Response", response.getStatus(), is(equalTo(204)));
 
-    verify(pushSender).sendMessage(any(Account.class), any(Device.class), any(MessageProtos.OutgoingMessageSignal.class));
+    verify(pushSender, times(1)).sendMessage(any(Account.class), any(Device.class), any(MessageProtos.OutgoingMessageSignal.class));
   }
 
   @Test
-  public void testMultiDeviceMissing() throws Exception {
+  public synchronized void testMultiDeviceMissing() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
+        resources.client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
             .entity(mapper.readValue(jsonFixture("fixtures/current_message_single_device.json"), IncomingMessageList.class))
             .type(MediaType.APPLICATION_JSON_TYPE)
@@ -121,9 +123,9 @@ public class MessageControllerTest extends ResourceTest {
   }
 
   @Test
-  public void testMultiDeviceExtra() throws Exception {
+  public synchronized void testMultiDeviceExtra() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
+        resources.client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
             .entity(mapper.readValue(jsonFixture("fixtures/current_message_extra_device.json"), IncomingMessageList.class))
             .type(MediaType.APPLICATION_JSON_TYPE)
@@ -139,9 +141,9 @@ public class MessageControllerTest extends ResourceTest {
   }
 
   @Test
-  public void testMultiDevice() throws Exception {
+  public synchronized void testMultiDevice() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
+        resources.client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
             .entity(mapper.readValue(jsonFixture("fixtures/current_message_multi_device.json"), IncomingMessageList.class))
             .type(MediaType.APPLICATION_JSON_TYPE)
@@ -153,9 +155,9 @@ public class MessageControllerTest extends ResourceTest {
   }
 
   @Test
-  public void testRegistrationIdMismatch() throws Exception {
+  public synchronized void testRegistrationIdMismatch() throws Exception {
     ClientResponse response =
-        client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
+        resources.client().resource(String.format("/v1/messages/%s", MULTI_DEVICE_RECIPIENT))
             .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD))
             .entity(mapper.readValue(jsonFixture("fixtures/current_message_registration_id.json"), IncomingMessageList.class))
             .type(MediaType.APPLICATION_JSON_TYPE)

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

@@ -2,25 +2,28 @@ package org.whispersystems.textsecuregcm.tests.controllers;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Optional;
-import com.yammer.dropwizard.auth.basic.BasicCredentials;
+import org.eclipse.jetty.websocket.api.CloseStatus;
-import org.eclipse.jetty.websocket.WebSocket;
+import org.eclipse.jetty.websocket.api.RemoteEndpoint;
+import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.api.UpgradeRequest;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.controllers.WebsocketController;
-import org.whispersystems.textsecuregcm.controllers.WebsocketControllerFactory;
 import org.whispersystems.textsecuregcm.entities.AcknowledgeWebsocketMessage;
+import org.whispersystems.textsecuregcm.entities.EncryptedOutgoingMessage;
+import org.whispersystems.textsecuregcm.push.PushSender;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.PubSubManager;
-import org.whispersystems.textsecuregcm.storage.StoredMessageManager;
+import org.whispersystems.textsecuregcm.storage.StoredMessages;
 import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
+import org.whispersystems.textsecuregcm.websocket.WebsocketControllerFactory;
 
-import javax.servlet.http.HttpServletRequest;
+import java.util.HashMap;
-
 import java.util.LinkedList;
 import java.util.List;
 
-import static org.fest.assertions.api.Assertions.assertThat;
+import io.dropwizard.auth.basic.BasicCredentials;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.*;
 
@@ -34,13 +37,14 @@ public class WebsocketControllerTest {
   private static final String VALID_PASSWORD   = "secure";
   private static final String INVALID_PASSWORD = "insecure";
 
-  private final StoredMessageManager storedMessageManager = mock(StoredMessageManager.class);
+  private static final StoredMessages       storedMessages       = mock(StoredMessages.class);
-  private final AccountAuthenticator accountAuthenticator = mock(AccountAuthenticator.class);
+  private static final AccountAuthenticator accountAuthenticator = mock(AccountAuthenticator.class);
-  private final PubSubManager        pubSubManager        = mock(PubSubManager.class       );
+  private static final PubSubManager        pubSubManager        = mock(PubSubManager.class       );
-  private final Account              account              = mock(Account.class             );
+  private static final Account              account              = mock(Account.class             );
-  private final Device               device               = mock(Device.class              );
+  private static final Device               device               = mock(Device.class              );
-  private final HttpServletRequest   request              = mock(HttpServletRequest.class  );
+  private static final UpgradeRequest       upgradeRequest       = mock(UpgradeRequest.class      );
-  private final WebSocket.Connection connection           = mock(WebSocket.Connection.class);
+  private static final Session              session              = mock(Session.class             );
+  private static final PushSender           pushSender           = mock(PushSender.class);
 
   @Test
   public void testCredentials() throws Exception {
@@ -50,23 +54,35 @@ public class WebsocketControllerTest {
     when(accountAuthenticator.authenticate(eq(new BasicCredentials(INVALID_USER, INVALID_PASSWORD))))
         .thenReturn(Optional.<Account>absent());
 
-    WebsocketControllerFactory factory = new WebsocketControllerFactory(accountAuthenticator,
+    when(session.getUpgradeRequest()).thenReturn(upgradeRequest);
-                                                                        storedMessageManager,
-                                                                        pubSubManager);
 
-    when(request.getParameter(eq("user"))).thenReturn(VALID_USER);
+    WebsocketController controller = new WebsocketController(accountAuthenticator, pushSender, pubSubManager, storedMessages);
-    when(request.getParameter(eq("password"))).thenReturn(VALID_PASSWORD);
 
-    assertThat(factory.checkOrigin(request, "foobar")).isEqualTo(true);
+    when(upgradeRequest.getParameterMap()).thenReturn(new HashMap<String, String[]>() {{
+      put("login", new String[] {VALID_USER});
+      put("password", new String[] {VALID_PASSWORD});
+    }});
 
-    when(request.getParameter(eq("user"))).thenReturn(INVALID_USER);
+    controller.onWebSocketConnect(session);
-    when(request.getParameter(eq("password"))).thenReturn(INVALID_PASSWORD);
 
-    assertThat(factory.checkOrigin(request, "foobar")).isEqualTo(false);
+    verify(session, never()).close();
+    verify(session, never()).close(any(CloseStatus.class));
+    verify(session, never()).close(anyInt(), anyString());
+
+    when(upgradeRequest.getParameterMap()).thenReturn(new HashMap<String, String[]>() {{
+      put("login", new String[] {INVALID_USER});
+      put("password", new String[] {INVALID_PASSWORD});
+    }});
+
+    controller.onWebSocketConnect(session);
+
+    verify(session).close(any(CloseStatus.class));
   }
 
   @Test
   public void testOpen() throws Exception {
+    RemoteEndpoint remote = mock(RemoteEndpoint.class);
+
     List<String> outgoingMessages = new LinkedList<String>() {{
       add("first");
       add("second");
@@ -76,36 +92,36 @@ public class WebsocketControllerTest {
     when(device.getId()).thenReturn(2L);
     when(account.getId()).thenReturn(31337L);
     when(account.getAuthenticatedDevice()).thenReturn(Optional.of(device));
+    when(session.getRemote()).thenReturn(remote);
+    when(session.getUpgradeRequest()).thenReturn(upgradeRequest);
 
-    when(request.getParameter(eq("user"))).thenReturn(VALID_USER);
+    when(upgradeRequest.getParameterMap()).thenReturn(new HashMap<String, String[]>() {{
-    when(request.getParameter(eq("password"))).thenReturn(VALID_PASSWORD);
+      put("login", new String[] {VALID_USER});
+      put("password", new String[] {VALID_PASSWORD});
+    }});
 
     when(accountAuthenticator.authenticate(eq(new BasicCredentials(VALID_USER, VALID_PASSWORD))))
         .thenReturn(Optional.of(account));
 
-    when(storedMessageManager.getOutgoingMessages(eq(account), eq(device))).thenReturn(outgoingMessages);
+    when(storedMessages.getMessagesForDevice(account.getId(), device.getId())).thenReturn(outgoingMessages);
 
-    WebsocketControllerFactory factory = new WebsocketControllerFactory(accountAuthenticator,
+    WebsocketControllerFactory factory    = new WebsocketControllerFactory(accountAuthenticator, pushSender, storedMessages, pubSubManager);
-                                                                        storedMessageManager,
+    WebsocketController        controller = (WebsocketController) factory.createWebSocket(null, null);
-                                                                        pubSubManager);
 
-    assertThat(factory.checkOrigin(request, "foobar")).isEqualTo(true);
+    controller.onWebSocketConnect(session);
 
-    WebsocketController socket = (WebsocketController)factory.doWebSocketConnect(request, "foo");
+    verify(pubSubManager).subscribe(eq(new WebsocketAddress(31337L, 2L)), eq((controller)));
-    socket.onOpen(connection);
+    verify(remote, times(3)).sendStringByFuture(anyString());
 
-    verify(pubSubManager).subscribe(eq(new WebsocketAddress(31337L, 2L)), eq((socket)));
+    controller.onWebSocketText(mapper.writeValueAsString(new AcknowledgeWebsocketMessage(1)));
-    verify(connection, times(3)).sendMessage(anyString());
+    controller.onWebSocketClose(1000, "Closed");
-
-    socket.onMessage(mapper.writeValueAsString(new AcknowledgeWebsocketMessage(1)));
-    socket.onClose(1000, "Closed");
 
     List<String> pending = new LinkedList<String>() {{
       add("first");
       add("third");
     }};
 
-    verify(storedMessageManager).storeMessages(eq(account), eq(device), eq(pending));
+    verify(pushSender, times(2)).sendMessage(eq(account), eq(device), any(EncryptedOutgoingMessage.class));
   }
 
 }

src/test/java/org/whispersystems/textsecuregcm/tests/entities/ClientContactTest.java

@@ -1,13 +1,16 @@
 package org.whispersystems.textsecuregcm.tests.entities;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.Test;
 import org.whispersystems.textsecuregcm.entities.ClientContact;
 import org.whispersystems.textsecuregcm.util.Util;
 
-import static com.yammer.dropwizard.testing.JsonHelpers.*;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.asJson;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.fromJson;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.jsonFixture;
 
 public class ClientContactTest {
 
@@ -41,4 +44,5 @@ public class ClientContactTest {
                is(contact));
   }
 
+
 }

src/test/java/org/whispersystems/textsecuregcm/tests/entities/PreKeyTest.java

@@ -5,10 +5,10 @@ import org.whispersystems.textsecuregcm.entities.ClientContact;
 import org.whispersystems.textsecuregcm.entities.PreKey;
 import org.whispersystems.textsecuregcm.util.Util;
 
-import static com.yammer.dropwizard.testing.JsonHelpers.*;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.whispersystems.textsecuregcm.tests.util.JsonHelpers.*;
 
 public class PreKeyTest {
 

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

@@ -0,0 +1,27 @@
+package org.whispersystems.textsecuregcm.tests.util;
+
+import com.fasterxml.jackson.core.JsonParseException;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import java.io.IOException;
+
+import static io.dropwizard.testing.FixtureHelpers.fixture;
+
+public class JsonHelpers {
+
+  private static final ObjectMapper objectMapper = new ObjectMapper();
+
+  public static String asJson(Object object) throws JsonProcessingException {
+    return objectMapper.writeValueAsString(object);
+  }
+
+  public static <T> T fromJson(String value, Class<T> clazz) throws IOException {
+    return objectMapper.readValue(value, clazz);
+  }
+
+  public static String jsonFixture(String filename) throws IOException {
+    return objectMapper.writeValueAsString(objectMapper.readValue(fixture(filename), JsonNode.class));
+  }
+}

src/test/resources/fixtures/contact.relay.json

@@ -1,4 +1,4 @@
 {
-    "token" : "BQVVHxMt5zAFXA",
+    "relay" : "whisper",
-    "relay" : "whisper"
+    "token" : "BQVVHxMt5zAFXA"
 }

src/test/resources/fixtures/contact.relay.sms.json

@@ -1,5 +1,5 @@
 {
-    "token" : "BQVVHxMt5zAFXA",
     "relay" : "whisper",
-    "supportsSms"   : true
+    "supportsSms"   : true,
+    "token" : "BQVVHxMt5zAFXA"
 }