From 4c3aae63d37596baaa76c11c68b728ae8d7e68ab Mon Sep 17 00:00:00 2001
From: Moxie Marlinspike <moxie@thoughtcrime.org>
Date: Tue, 11 Aug 2015 20:15:05 -0700
Subject: [PATCH] Trim old messages

// FREEBIE
---
 .../textsecuregcm/WhisperServerService.java   |  2 +
 .../textsecuregcm/storage/Messages.java       |  3 ++
 .../workers/TrimMessagesCommand.java          | 53 +++++++++++++++++++
 3 files changed, 58 insertions(+)
 create mode 100644 src/main/java/org/whispersystems/textsecuregcm/workers/TrimMessagesCommand.java

diff --git a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 9cc10c780..ea1980034 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -88,6 +88,7 @@ import org.whispersystems.textsecuregcm.websocket.DeadLetterHandler;
 import org.whispersystems.textsecuregcm.websocket.ProvisioningConnectListener;
 import org.whispersystems.textsecuregcm.websocket.WebSocketAccountAuthenticator;
 import org.whispersystems.textsecuregcm.workers.DirectoryCommand;
+import org.whispersystems.textsecuregcm.workers.TrimMessagesCommand;
 import org.whispersystems.textsecuregcm.workers.VacuumCommand;
 import org.whispersystems.websocket.WebSocketResourceProviderFactory;
 import org.whispersystems.websocket.setup.WebSocketEnvironment;
@@ -120,6 +121,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
   public void initialize(Bootstrap<WhisperServerConfiguration> bootstrap) {
     bootstrap.addCommand(new DirectoryCommand());
     bootstrap.addCommand(new VacuumCommand());
+    bootstrap.addCommand(new TrimMessagesCommand());
     bootstrap.addBundle(new NameableMigrationsBundle<WhisperServerConfiguration>("accountdb", "accountsdb.xml") {
       @Override
       public DataSourceFactory getDataSourceFactory(WhisperServerConfiguration configuration) {
diff --git a/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java b/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java
index b7ecda3a2..3ea9de700 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java
@@ -58,6 +58,9 @@ public abstract class Messages {
   @SqlUpdate("DELETE FROM messages WHERE " + DESTINATION + " = :destination")
   abstract void clear(@Bind("destination") String destination);
 
+  @SqlUpdate("DELETE FROM messages WHERE " + TIMESTAMP + " < :timestamp")
+  public abstract void removeOld(@Bind("timestamp") long timestamp);
+
   @SqlUpdate("VACUUM messages")
   public abstract void vacuum();
 
diff --git a/src/main/java/org/whispersystems/textsecuregcm/workers/TrimMessagesCommand.java b/src/main/java/org/whispersystems/textsecuregcm/workers/TrimMessagesCommand.java
new file mode 100644
index 000000000..727fa7ff2
--- /dev/null
+++ b/src/main/java/org/whispersystems/textsecuregcm/workers/TrimMessagesCommand.java
@@ -0,0 +1,53 @@
+package org.whispersystems.textsecuregcm.workers;
+
+import net.sourceforge.argparse4j.inf.Namespace;
+import org.skife.jdbi.v2.DBI;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
+import org.whispersystems.textsecuregcm.storage.Accounts;
+import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.Messages;
+import org.whispersystems.textsecuregcm.storage.PendingAccounts;
+
+import java.util.concurrent.TimeUnit;
+
+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;
+
+public class TrimMessagesCommand extends ConfiguredCommand<WhisperServerConfiguration> {
+  private final Logger logger = LoggerFactory.getLogger(VacuumCommand.class);
+
+  public TrimMessagesCommand() {
+    super("trim", "Trim Messages Database");
+  }
+
+  @Override
+  protected void run(Bootstrap<WhisperServerConfiguration> bootstrap,
+                     Namespace namespace,
+                     WhisperServerConfiguration config)
+      throws Exception
+  {
+    DataSourceFactory messageDbConfig = config.getMessageStoreConfiguration();
+    DBI               messageDbi      = new DBI(messageDbConfig.getUrl(), messageDbConfig.getUser(), messageDbConfig.getPassword());
+
+    messageDbi.registerArgumentFactory(new OptionalArgumentFactory(messageDbConfig.getDriverClass()));
+    messageDbi.registerContainerFactory(new ImmutableListContainerFactory());
+    messageDbi.registerContainerFactory(new ImmutableSetContainerFactory());
+    messageDbi.registerContainerFactory(new OptionalContainerFactory());
+
+    Messages messages  = messageDbi.onDemand(Messages.class);
+    long     timestamp = System.currentTimeMillis() - TimeUnit.DAYS.toMillis(60);
+
+    logger.info("Trimming old messages: " + timestamp + "...");
+    messages.removeOld(timestamp);
+
+    Thread.sleep(3000);
+    System.exit(0);
+  }
+}