Upgrade to dropwizard 0.7.

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 Logger         logger            = LoggerFactory.getLogger(DirectoryController.class);
+  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 {
 

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,9 +1,14 @@
 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.IncomingWebsocketMessage;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -22,111 +27,94 @@ 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 AccountAuthenticator accountAuthenticator;
   private final StoredMessageManager storedMessageManager;
   private final PubSubManager        pubSubManager;
 
-  private final Account account;
-  private final Device  device;
+  private Account account;
+  private Device  device;
+  private Session session;
 
-  private Connection connection;
-  private long       pendingMessageSequence;
+  private long pendingMessageSequence;
 
-  public WebsocketController(StoredMessageManager storedMessageManager,
-                             PubSubManager pubSubManager,
-                             Account account)
+  public WebsocketController(AccountAuthenticator accountAuthenticator,
+                             StoredMessageManager storedMessageManager,
+                             PubSubManager pubSubManager)
   {
+    this.accountAuthenticator = accountAuthenticator;
     this.storedMessageManager = storedMessageManager;
     this.pubSubManager        = pubSubManager;
-    this.account              = account;
-    this.device               = account.getAuthenticatedDevice().get();
-  }
-
-
-  @Override
-  public void onOpen(Connection connection) {
-    this.connection = connection;
-    pubSubManager.subscribe(new WebsocketAddress(this.account.getId(), this.device.getId()), this);
-    handleQueryDatabase();
   }
 
   @Override
-  public void onClose(int i, String s) {
-    handleClose();
+  public void onWebSocketConnect(Session session) {
+    try {
+      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(401, "Unauthorized"));
+        return;
+      }
+
+      BasicCredentials  credentials = new BasicCredentials(usernames[0], passwords[0]);
+      Optional<Account> account     = accountAuthenticator.authenticate(credentials);
+
+      if (!account.isPresent()) {
+        session.close(new CloseStatus(401, "Unauthorized"));
+        return;
+      }
+
+      this.account = account.get();
+      this.device  = account.get().getAuthenticatedDevice().get();
+      this.session = session;
+
+      this.session.setIdleTimeout(10 * 60 * 1000);
+      this.pubSubManager.subscribe(new WebsocketAddress(this.account.getId(),
+                                                        this.device.getId()), this);
+
+      handleQueryDatabase();
+    } catch (AuthenticationException e) {
+      try { session.close(500, "Server Error");} catch (IOException e1) {}
+    } catch (IOException ioe) {
+      logger.info("Abrupt session close.");
+    }
   }
 
   @Override
-  public void onMessage(String body) {
+  public void onWebSocketText(String body) {
     try {
       IncomingWebsocketMessage incomingMessage = mapper.readValue(body, IncomingWebsocketMessage.class);
 
       switch (incomingMessage.getType()) {
-        case IncomingWebsocketMessage.TYPE_ACKNOWLEDGE_MESSAGE: handleMessageAck(body); break;
-        case IncomingWebsocketMessage.TYPE_PING_MESSAGE:        handlePing();           break;
-        default:                                                handleClose();          break;
+        case IncomingWebsocketMessage.TYPE_ACKNOWLEDGE_MESSAGE:
+          handleMessageAck(body);
+          break;
+        default:
+          close(new CloseStatus(410, "Unknown Type"));
       }
     } catch (IOException e) {
       logger.debug("Parse", e);
-      handleClose();
+      close(new CloseStatus(410, "Badly Formatted"));
     }
   }
 
   @Override
-  public void onPubSubMessage(PubSubMessage outgoingMessage) {
-    switch (outgoingMessage.getType()) {
-      case PubSubMessage.TYPE_DELIVER:  handleDeliverOutgoingMessage(outgoingMessage.getContents()); break;
-      case PubSubMessage.TYPE_QUERY_DB: handleQueryDatabase();                                       break;
-      default:
-        logger.warn("Unknown pubsub message: " + outgoingMessage.getType());
-    }
-  }
-
-  private void handleDeliverOutgoingMessage(String message) {
-    try {
-      long messageSequence;
-
-      synchronized (pendingMessages) {
-        messageSequence = pendingMessageSequence++;
-        pendingMessages.put(messageSequence, message);
-      }
-
-      connection.sendMessage(mapper.writeValueAsString(new WebsocketMessage(messageSequence, message)));
-    } catch (IOException e) {
-      logger.debug("Response failed", e);
-      handleClose();
-    }
-  }
-
-  private void handleMessageAck(String message) {
-    try {
-      AcknowledgeWebsocketMessage ack = mapper.readValue(message, AcknowledgeWebsocketMessage.class);
-
-      synchronized (pendingMessages) {
-        pendingMessages.remove(ack.getId());
-      }
-    } catch (IOException e) {
-      logger.warn("Mapping", e);
-    }
-  }
-
-  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() {
+  public void onWebSocketClose(int i, String s) {
     pubSubManager.unsubscribe(new WebsocketAddress(account.getId(), device.getId()), this);
-    connection.close();
 
     List<String> remainingMessages = new LinkedList<>();
 
@@ -144,6 +132,50 @@ public class WebsocketController implements WebSocket.OnTextMessage, PubSubListe
     storedMessageManager.storeMessages(account, device, remainingMessages);
   }
 
+
+  @Override
+  public void onPubSubMessage(PubSubMessage outgoingMessage) {
+    switch (outgoingMessage.getType()) {
+      case PubSubMessage.TYPE_DELIVER:
+        handleDeliverOutgoingMessage(outgoingMessage.getContents());
+        break;
+      case PubSubMessage.TYPE_QUERY_DB:
+        handleQueryDatabase();
+        break;
+      default:
+        logger.warn("Unknown pubsub message: " + outgoingMessage.getType());
+    }
+  }
+
+  private void handleDeliverOutgoingMessage(String message) {
+    try {
+      long messageSequence;
+
+      synchronized (pendingMessages) {
+        messageSequence = pendingMessageSequence++;
+        pendingMessages.put(messageSequence, message);
+      }
+
+      WebsocketMessage websocketMessage = new WebsocketMessage(messageSequence, message);
+      session.getRemote().sendStringByFuture(mapper.writeValueAsString(websocketMessage));
+    } catch (IOException e) {
+      logger.debug("Response failed", e);
+      close(null);
+    }
+  }
+
+  private void handleMessageAck(String message) {
+    try {
+      AcknowledgeWebsocketMessage ack = mapper.readValue(message, AcknowledgeWebsocketMessage.class);
+
+      synchronized (pendingMessages) {
+        pendingMessages.remove(ack.getId());
+      }
+    } catch (IOException e) {
+      logger.warn("Mapping", e);
+    }
+  }
+
   private void handleQueryDatabase() {
     List<String> messages = storedMessageManager.getOutgoingMessages(account, device);
 
@@ -152,4 +184,25 @@ public class WebsocketController implements WebSocket.OnTextMessage, PubSubListe
     }
   }
 
+  @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,23 +1,18 @@
 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.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.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 {
+public class WebsocketControllerFactory extends WebSocketServlet implements WebSocketCreator {
 
   private final Logger logger = LoggerFactory.getLogger(WebsocketControllerFactory.class);
 
@@ -25,14 +20,6 @@ public class WebsocketControllerFactory extends WebSocketServlet {
   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)
@@ -43,56 +30,12 @@ public class WebsocketControllerFactory extends WebSocketServlet {
   }
 
   @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);
-    }
+  public void configure(WebSocketServletFactory factory) {
+    factory.setCreator(this);
   }
 
   @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;
-    }
+  public Object createWebSocket(UpgradeRequest upgradeRequest, UpgradeResponse upgradeResponse) {
+    return new WebsocketController(accountAuthenticator, storedMessageManager, pubSubManager);
   }
 }