From 7d07d56fc3220ed816c51b14f5a0bd75bb9f7a7b Mon Sep 17 00:00:00 2001
From: Moxie Marlinspike <moxie@thoughtcrime.org>
Date: Mon, 22 Jul 2013 15:04:31 -0700
Subject: [PATCH] Fix for 'bad encrypted message' errors.

1) There was a regression in the outgoing multipart transport
   logic, such that the same 'identifier' byte would be used
   for all messages (0).  This now works correctly.

2) Added some additional heuristics on the receiving side.
   Now mutlipart containers are only valid for 1hr, and are
   considered invalid if the container size is different from
   the multipart message size.
---
 .../securesms/sms/MultipartSmsIdentifier.java | 32 +++++++++++++++++++
 .../sms/MultipartSmsMessageHandler.java       | 29 ++++++-----------
 ...MultipartSmsTransportMessageFragments.java | 19 ++++++++++-
 3 files changed, 59 insertions(+), 21 deletions(-)
 create mode 100644 src/org/thoughtcrime/securesms/sms/MultipartSmsIdentifier.java

diff --git a/src/org/thoughtcrime/securesms/sms/MultipartSmsIdentifier.java b/src/org/thoughtcrime/securesms/sms/MultipartSmsIdentifier.java
new file mode 100644
index 0000000000..0088159e2b
--- /dev/null
+++ b/src/org/thoughtcrime/securesms/sms/MultipartSmsIdentifier.java
@@ -0,0 +1,32 @@
+package org.thoughtcrime.securesms.sms;
+
+
+import java.util.HashMap;
+
+public class MultipartSmsIdentifier {
+
+  private static final MultipartSmsIdentifier instance = new MultipartSmsIdentifier();
+
+  public static MultipartSmsIdentifier getInstance() {
+    return instance;
+  }
+
+  private final HashMap<String, Integer>  idMap = new HashMap<String, Integer>();
+
+  public synchronized byte getIdForRecipient(String recipient) {
+    Integer currentId;
+
+    if (idMap.containsKey(recipient)) {
+      currentId = idMap.get(recipient);
+      idMap.remove(recipient);
+    } else {
+      currentId = 0;
+    }
+
+    byte id  = currentId.byteValue();
+    idMap.put(recipient, (currentId + 1) % 255);
+
+    return id;
+  }
+
+}
diff --git a/src/org/thoughtcrime/securesms/sms/MultipartSmsMessageHandler.java b/src/org/thoughtcrime/securesms/sms/MultipartSmsMessageHandler.java
index 3aaa21d262..4bd92e00dc 100644
--- a/src/org/thoughtcrime/securesms/sms/MultipartSmsMessageHandler.java
+++ b/src/org/thoughtcrime/securesms/sms/MultipartSmsMessageHandler.java
@@ -29,13 +29,17 @@ public class MultipartSmsMessageHandler {
   private final HashMap<String, MultipartSmsTransportMessageFragments> partialMessages =
       new HashMap<String, MultipartSmsTransportMessageFragments>();
 
-  private final HashMap<String, Integer>  idMap = new HashMap<String, Integer>();
-
   private IncomingTextMessage processMultipartMessage(MultipartSmsTransportMessage message) {
     Log.w("MultipartSmsMessageHandler", "Processing multipart message...");
+    Log.w("MultipartSmsMessageHandler", "Multipart Count: " + message.getMultipartCount());
+    Log.w("MultipartSmsMessageHandler", "Multipart ID: " + message.getIdentifier());
+    Log.w("MultipartSmsMessageHandler", "Multipart Key: " + message.getKey());
     MultipartSmsTransportMessageFragments container = partialMessages.get(message.getKey());
 
-    if (container == null) {
+    Log.w("MultipartSmsMessageHandler", "Found multipart container: " + container);
+
+    if (container == null || container.getSize() != message.getMultipartCount() || container.isExpired()) {
+      Log.w("MultipartSmsMessageHandler", "Constructing new container...");
       container = new MultipartSmsTransportMessageFragments(message.getMultipartCount());
       partialMessages.put(message.getKey(), container);
     }
@@ -81,24 +85,9 @@ public class MultipartSmsMessageHandler {
     }
   }
 
-  private byte getIdForRecipient(String recipient) {
-    Integer currentId;
-
-    if (idMap.containsKey(recipient)) {
-      currentId = idMap.get(recipient);
-      idMap.remove(recipient);
-    } else {
-      currentId = 0;
-    }
-
-    byte id  = currentId.byteValue();
-    idMap.put(recipient, (currentId + 1) % 255);
-
-    return id;
-  }
-
   public ArrayList<String> divideMessage(OutgoingTextMessage message) {
-    byte identifier = getIdForRecipient(message.getRecipients().getPrimaryRecipient().getNumber());
+    String number     = message.getRecipients().getPrimaryRecipient().getNumber();
+    byte   identifier = MultipartSmsIdentifier.getInstance().getIdForRecipient(number);
     return MultipartSmsTransportMessage.getEncoded(message, identifier);
   }
 }
diff --git a/src/org/thoughtcrime/securesms/sms/MultipartSmsTransportMessageFragments.java b/src/org/thoughtcrime/securesms/sms/MultipartSmsTransportMessageFragments.java
index b4ef13f2f4..810988a193 100644
--- a/src/org/thoughtcrime/securesms/sms/MultipartSmsTransportMessageFragments.java
+++ b/src/org/thoughtcrime/securesms/sms/MultipartSmsTransportMessageFragments.java
@@ -2,16 +2,28 @@ package org.thoughtcrime.securesms.sms;
 
 public class MultipartSmsTransportMessageFragments {
 
+  private static final long VALID_TIME = 60 * 60 * 1000; // 1 Hour
+
   private final byte[][] fragments;
+  private final long initializedTime;
 
   public MultipartSmsTransportMessageFragments(int count) {
-    this.fragments = new byte[count][];
+    this.fragments       = new byte[count][];
+    this.initializedTime = System.currentTimeMillis();
   }
 
   public void add(MultipartSmsTransportMessage fragment) {
     this.fragments[fragment.getMultipartIndex()] = fragment.getStrippedMessage();
   }
 
+  public int getSize() {
+    return this.fragments.length;
+  }
+
+  public boolean isExpired() {
+    return (System.currentTimeMillis() - initializedTime) >= VALID_TIME;
+  }
+
   public boolean isComplete() {
     for (int i=0;i<fragments.length;i++)
       if (fragments[i] == null) return false;
@@ -37,4 +49,9 @@ public class MultipartSmsTransportMessageFragments {
     return totalMessage;
   }
 
+  @Override
+  public String toString() {
+    return String.format("[Size: %d, Initialized: %d, Exipired: %s, Complete: %s]",
+                         fragments.length, initializedTime, isExpired()+"", isComplete()+"");
+  }
 }