Add new upload-for-copy backup endpoint

service/src/test/java/org/whispersystems/textsecuregcm/backup/BackupManagerTest.java

@@ -13,6 +13,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
@@ -60,7 +61,11 @@ import org.signal.libsignal.protocol.ecc.ECKeyPair;
 import org.signal.libsignal.zkgroup.GenericServerSecretParams;
 import org.signal.libsignal.zkgroup.VerificationFailedException;
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialPresentation;
+import org.whispersystems.textsecuregcm.attachments.TusAttachmentGenerator;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedBackupUser;
+import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import org.whispersystems.textsecuregcm.limits.RateLimiter;
+import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
@@ -79,6 +84,8 @@ public class BackupManagerTest {
 
   private final TestClock testClock = TestClock.now();
   private final BackupAuthTestUtil backupAuthTestUtil = new BackupAuthTestUtil(testClock);
+  private final RateLimiter mediaUploadLimiter = mock(RateLimiter.class);
+  private final TusAttachmentGenerator tusAttachmentGenerator = mock(TusAttachmentGenerator.class);
   private final Cdn3BackupCredentialGenerator tusCredentialGenerator = mock(Cdn3BackupCredentialGenerator.class);
   private final RemoteStorageManager remoteStorageManager = mock(RemoteStorageManager.class);
   private final byte[] backupKey = TestRandomUtil.nextBytes(32);
@@ -90,8 +97,12 @@ public class BackupManagerTest {
 
   @BeforeEach
   public void setup() {
-    reset(tusCredentialGenerator);
+    reset(tusCredentialGenerator, mediaUploadLimiter);
     testClock.unpin();
+
+    final RateLimiters rateLimiters = mock(RateLimiters.class);
+    when(rateLimiters.forDescriptor(RateLimiters.For.BACKUP_ATTACHMENT)).thenReturn(mediaUploadLimiter);
+
     this.backupsDb = new BackupsDb(
         DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
         DynamoDbExtensionSchema.Tables.BACKUPS.tableName(),
@@ -99,6 +110,8 @@ public class BackupManagerTest {
     this.backupManager = new BackupManager(
         backupsDb,
         backupAuthTestUtil.params,
+        rateLimiters,
+        tusAttachmentGenerator,
         tusCredentialGenerator,
         remoteStorageManager,
         Map.of(3, "cdn3.example.org/attachments"),
@@ -127,6 +140,28 @@ public class BackupManagerTest {
     checkExpectedExpirations(now, backupTier == BackupTier.MEDIA ? now : null, backupUser);
   }
 
+  @Test
+  public void createTemporaryMediaAttachmentRateLimited() throws RateLimitExceededException {
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupTier.MEDIA);
+    doThrow(new RateLimitExceededException(null, true))
+        .when(mediaUploadLimiter)
+        .validate(eq(BackupManager.rateLimitKey(backupUser)));
+
+    assertThatExceptionOfType(RateLimitExceededException.class)
+        .isThrownBy(() -> backupManager.createTemporaryAttachmentUploadDescriptor(backupUser))
+        .satisfies(e -> assertThat(e.isLegacy()).isFalse());
+  }
+
+  @Test
+  public void createTemporaryMediaAttachmentWrongTier() throws RateLimitExceededException {
+    final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupTier.MESSAGES);
+    assertThatExceptionOfType(StatusRuntimeException.class)
+        .isThrownBy(() -> backupManager.createTemporaryAttachmentUploadDescriptor(backupUser))
+        .extracting(StatusRuntimeException::getStatus)
+        .extracting(Status::getCode)
+        .isEqualTo(Status.Code.PERMISSION_DENIED);
+  }
+
   @ParameterizedTest
   @EnumSource(mode = EnumSource.Mode.EXCLUDE, names = {"NONE"})
   public void ttlRefresh(final BackupTier backupTier) {
@@ -317,7 +352,7 @@ public class BackupManagerTest {
   public void copySuccess() {
     final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupTier.MEDIA);
     when(tusCredentialGenerator.generateUpload(any()))
-        .thenReturn(new MessageBackupUploadDescriptor(3, "def", Collections.emptyMap(), ""));
+        .thenReturn(new BackupUploadDescriptor(3, "def", Collections.emptyMap(), ""));
     when(remoteStorageManager.copy(eq(URI.create("cdn3.example.org/attachments/abc")), eq(100), any(), any()))
         .thenReturn(CompletableFuture.completedFuture(null));
     final MediaEncryptionParameters encryptionParams = new MediaEncryptionParameters(
@@ -343,7 +378,7 @@ public class BackupManagerTest {
   public void copyFailure() {
     final AuthenticatedBackupUser backupUser = backupUser(TestRandomUtil.nextBytes(16), BackupTier.MEDIA);
     when(tusCredentialGenerator.generateUpload(any()))
-        .thenReturn(new MessageBackupUploadDescriptor(3, "def", Collections.emptyMap(), ""));
+        .thenReturn(new BackupUploadDescriptor(3, "def", Collections.emptyMap(), ""));
     when(remoteStorageManager.copy(eq(URI.create("cdn3.example.org/attachments/abc")), eq(100), any(), any()))
         .thenReturn(CompletableFuture.failedFuture(new SourceObjectNotFoundException()));
 

service/src/test/java/org/whispersystems/textsecuregcm/backup/Cdn3BackupCredentialGeneratorTest.java

@@ -22,7 +22,7 @@ public class Cdn3BackupCredentialGeneratorTest {
         new SecretBytes(TestRandomUtil.nextBytes(32)),
         "https://example.org/upload"));
 
-    final MessageBackupUploadDescriptor messageBackupUploadDescriptor = generator.generateUpload("subdir/key");
+    final BackupUploadDescriptor messageBackupUploadDescriptor = generator.generateUpload("subdir/key");
     assertThat(messageBackupUploadDescriptor.signedUploadLocation()).isEqualTo("https://example.org/upload/backups");
     assertThat(messageBackupUploadDescriptor.key()).isEqualTo("subdir/key");
     assertThat(messageBackupUploadDescriptor.headers()).containsKey("Authorization");

service/src/test/java/org/whispersystems/textsecuregcm/backup/Cdn3RemoteStorageManagerTest.java

@@ -124,7 +124,7 @@ public class Cdn3RemoteStorageManagerTest {
             URI.create(wireMock.url("/cdn" + sourceCdn + "/source/small")),
             expectedSource.length(),
             encryptionParameters,
-            new MessageBackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
+            new BackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
         .toCompletableFuture().join();
 
     final byte[] destBody = wireMock.findAll(postRequestedFor(urlEqualTo("/cdn3/dest"))).get(0).getBody();
@@ -148,7 +148,7 @@ public class Cdn3RemoteStorageManagerTest {
             URI.create(wireMock.url("/cdn3/source/large")),
             LARGE.length(),
             params,
-            new MessageBackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
+            new BackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
         .toCompletableFuture().join();
 
     final byte[] destBody = wireMock.findAll(postRequestedFor(urlEqualTo("/cdn3/dest"))).get(0).getBody();
@@ -165,7 +165,7 @@ public class Cdn3RemoteStorageManagerTest {
                 URI.create(wireMock.url("/cdn3/source/small")),
                 SMALL_CDN3.length() - 1,
                 new MediaEncryptionParameters(AES_KEY, HMAC_KEY, IV),
-                new MessageBackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
+                new BackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
             .toCompletableFuture());
   }
 
@@ -176,7 +176,7 @@ public class Cdn3RemoteStorageManagerTest {
                 URI.create(wireMock.url("/cdn3/source/missing")),
                 1,
                 new MediaEncryptionParameters(AES_KEY, HMAC_KEY, IV),
-                new MessageBackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
+                new BackupUploadDescriptor(3, "test", Collections.emptyMap(), wireMock.url("/cdn3/dest")))
             .toCompletableFuture());
   }
 

service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java

@@ -26,6 +26,7 @@ import java.time.temporal.ChronoUnit;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -36,7 +37,6 @@ import javax.ws.rs.client.Invocation;
 import javax.ws.rs.client.WebTarget;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
-import org.checkerframework.checker.units.qual.A;
 import org.glassfish.jersey.server.ServerProperties;
 import org.glassfish.jersey.test.grizzly.GrizzlyWebTestContainerFactory;
 import org.junit.jupiter.api.BeforeEach;
@@ -68,6 +68,7 @@ import org.whispersystems.textsecuregcm.backup.BackupManager;
 import org.whispersystems.textsecuregcm.backup.BackupTier;
 import org.whispersystems.textsecuregcm.backup.InvalidLengthException;
 import org.whispersystems.textsecuregcm.backup.SourceObjectNotFoundException;
+import org.whispersystems.textsecuregcm.backup.BackupUploadDescriptor;
 import org.whispersystems.textsecuregcm.mappers.CompletionExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.GrpcStatusRuntimeExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
@@ -108,6 +109,7 @@ public class ArchiveControllerTest {
       GET,    v1/archives/auth/read,
       GET,    v1/archives/,
       GET,    v1/archives/upload/form,
+      GET,    v1/archives/media/upload/form,
       POST,   v1/archives/,
       PUT,    v1/archives/keys, '{"backupIdPublicKey": "aaaaa"}'
       PUT,    v1/archives/media, '{
@@ -531,6 +533,38 @@ public class ArchiveControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
   }
 
+  @Test
+  public void mediaUploadForm() throws RateLimitExceededException, VerificationFailedException {
+    final BackupAuthCredentialPresentation presentation =
+        backupAuthTestUtil.getPresentation(BackupTier.MEDIA, backupKey, aci);
+    when(backupManager.authenticateBackupUser(any(), any()))
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupTier.MEDIA)));
+    when(backupManager.createTemporaryAttachmentUploadDescriptor(any()))
+        .thenReturn(new BackupUploadDescriptor(3, "abc", Map.of("k", "v"), "example.org"));
+    final ArchiveController.UploadDescriptorResponse desc = resources.getJerseyTest()
+        .target("v1/archives/media/upload/form")
+        .request()
+        .header("X-Signal-ZK-Auth", Base64.getEncoder().encodeToString(presentation.serialize()))
+        .header("X-Signal-ZK-Auth-Signature", "aaa")
+        .get(ArchiveController.UploadDescriptorResponse.class);
+    assertThat(desc.cdn()).isEqualTo(3);
+    assertThat(desc.key()).isEqualTo("abc");
+    assertThat(desc.headers()).containsExactlyEntriesOf(Map.of("k", "v"));
+    assertThat(desc.signedUploadLocation()).isEqualTo("example.org");
+
+    // rate limit
+    when(backupManager.createTemporaryAttachmentUploadDescriptor(any()))
+        .thenThrow(new RateLimitExceededException(null, false));
+    final Response response = resources.getJerseyTest()
+        .target("v1/archives/media/upload/form")
+        .request()
+        .header("X-Signal-ZK-Auth", Base64.getEncoder().encodeToString(presentation.serialize()))
+        .header("X-Signal-ZK-Auth-Signature", "aaa")
+        .get();
+    assertThat(response.getStatus()).isEqualTo(429);
+  }
+
+
   private static AuthenticatedBackupUser backupUser(byte[] backupId, BackupTier backupTier) {
     return new AuthenticatedBackupUser(backupId, backupTier, "myBackupDir", "myMediaDir");
   }