From 51b9affe90c09bce2485862f37a60e9bf155b40b Mon Sep 17 00:00:00 2001
From: Moxie Marlinspike <moxie@thoughtcrime.org>
Date: Sat, 19 Oct 2013 18:16:11 -0700
Subject: [PATCH] Add federation support for attachment retrieval.

---
 .../textsecure/push/PushServiceSocket.java            | 11 ++++++++---
 .../securesms/mms/IncomingMediaMessage.java           |  5 +++++
 .../securesms/service/PushDownloader.java             | 11 ++++++++---
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/library/src/org/whispersystems/textsecure/push/PushServiceSocket.java b/library/src/org/whispersystems/textsecure/push/PushServiceSocket.java
index 350c77a5ed..d90bd92383 100644
--- a/library/src/org/whispersystems/textsecure/push/PushServiceSocket.java
+++ b/library/src/org/whispersystems/textsecure/push/PushServiceSocket.java
@@ -172,9 +172,14 @@ public class PushServiceSocket {
     return new Gson().fromJson(response.second, AttachmentKey.class).getId();
   }
 
-  public File retrieveAttachment(long attachmentId) throws IOException {
-    Pair<String, String> response = makeRequestForResponseHeader(String.format(ATTACHMENT_PATH, String.valueOf(attachmentId)),
-                                                                 "GET", null, "Content-Location");
+  public File retrieveAttachment(String relay, long attachmentId) throws IOException {
+    String path = String.format(ATTACHMENT_PATH, String.valueOf(attachmentId));
+
+    if (relay != null) {
+      path = path + "?relay=" + relay;
+    }
+
+    Pair<String, String> response = makeRequestForResponseHeader(path, "GET", null, "Content-Location");
 
     Log.w("PushServiceSocket", "Attachment: " + attachmentId + " is at: " + response.first);
 
diff --git a/src/org/thoughtcrime/securesms/mms/IncomingMediaMessage.java b/src/org/thoughtcrime/securesms/mms/IncomingMediaMessage.java
index 29b23bdcec..c3f80bf87e 100644
--- a/src/org/thoughtcrime/securesms/mms/IncomingMediaMessage.java
+++ b/src/org/thoughtcrime/securesms/mms/IncomingMediaMessage.java
@@ -57,6 +57,11 @@ public class IncomingMediaMessage {
         media.setContentType(Util.toIsoBytes(attachment.getContentType()));
         media.setContentLocation(Util.toIsoBytes(String.valueOf(attachment.getId())));
         media.setContentDisposition(Util.toIsoBytes(Base64.encodeBytes(encryptedKey)));
+
+        if (message.getRelay() != null) {
+          media.setName(Util.toIsoBytes(message.getRelay()));
+        }
+
         media.setPendingPush(true);
 
         body.addPart(media);
diff --git a/src/org/thoughtcrime/securesms/service/PushDownloader.java b/src/org/thoughtcrime/securesms/service/PushDownloader.java
index c40656d571..f2ebfdae8c 100644
--- a/src/org/thoughtcrime/securesms/service/PushDownloader.java
+++ b/src/org/thoughtcrime/securesms/service/PushDownloader.java
@@ -69,8 +69,13 @@ public class PushDownloader {
       MasterCipher masterCipher    = new MasterCipher(masterSecret);
       long         contentLocation = Long.parseLong(Util.toIsoString(part.getContentLocation()));
       byte[]       key             = masterCipher.decryptBytes(Base64.decode(Util.toIsoString(part.getContentDisposition())));
+      String       relay           = null;
 
-      attachmentFile              = downloadAttachment(contentLocation);
+      if (part.getName() != null) {
+        relay = Util.toIsoString(part.getName());
+      }
+
+      attachmentFile              = downloadAttachment(relay, contentLocation);
       InputStream attachmentInput = new AttachmentCipherInputStream(attachmentFile, key);
 
       database.updateDownloadedPart(messageId, partId, part, attachmentInput);
@@ -97,9 +102,9 @@ public class PushDownloader {
     }
   }
 
-  private File downloadAttachment(long contentLocation) throws IOException {
+  private File downloadAttachment(String relay, long contentLocation) throws IOException {
     PushServiceSocket socket = new PushServiceSocket(context, TextSecurePushCredentials.getInstance());
-    return socket.retrieveAttachment(contentLocation);
+    return socket.retrieveAttachment(relay, contentLocation);
   }
 
 }