From b5f9564e13f6c27035323d62eda2fbd5ddeb8e37 Mon Sep 17 00:00:00 2001
From: Ravi Khadiwala <ravi@signal.org>
Date: Wed, 14 Aug 2024 17:36:30 -0500
Subject: [PATCH] Validate that sourceAttachments are valid base64 strings

---
 .../controllers/ArchiveController.java        |  4 ++
 .../util/ValidBase64URLString.java            | 52 +++++++++++++++++++
 .../controllers/ArchiveControllerTest.java    | 21 ++++++++
 3 files changed, 77 insertions(+)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/util/ValidBase64URLString.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java
index 59a71fd99..4820b6e2c 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/ArchiveController.java
@@ -65,6 +65,7 @@ import org.whispersystems.textsecuregcm.util.ByteArrayBase64UrlAdapter;
 import org.whispersystems.textsecuregcm.util.ECPublicKeyAdapter;
 import org.whispersystems.textsecuregcm.util.ExactlySize;
 import org.whispersystems.textsecuregcm.util.Util;
+import org.whispersystems.textsecuregcm.util.ValidBase64URLString;
 import org.whispersystems.websocket.auth.Mutable;
 import org.whispersystems.websocket.auth.ReadOnly;
 import reactor.core.publisher.Mono;
@@ -464,13 +465,16 @@ public class ArchiveController {
       @Schema(description = "The attachment cdn")
       @NotNull
       Integer cdn,
+
       @NotBlank
+      @ValidBase64URLString
       @Schema(description = "The attachment key")
       String key) {}
 
   public record CopyMediaRequest(
       @Schema(description = "The object on the attachment CDN to copy")
       @NotNull
+      @Valid
       RemoteAttachment sourceAttachment,
 
       @Schema(description = "The length of the source attachment before the encryption applied by the copy operation")
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/ValidBase64URLString.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/ValidBase64URLString.java
new file mode 100644
index 000000000..a54daa3ef
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/ValidBase64URLString.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+import javax.validation.Constraint;
+import javax.validation.ConstraintValidator;
+import javax.validation.ConstraintValidatorContext;
+import javax.validation.Payload;
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+import java.util.Base64;
+import java.util.HexFormat;
+import java.util.Objects;
+
+import static java.lang.annotation.ElementType.*;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+/**
+ * Constraint annotation that requires annotated entity is a valid url-base64 encoded string.
+ */
+@Target({ FIELD, PARAMETER, METHOD })
+@Retention(RUNTIME)
+@Constraint(validatedBy = ValidBase64URLString.Validator.class)
+@Documented
+public @interface ValidBase64URLString {
+
+  String message() default "value is not a valid base64 string";
+
+  Class<?>[] groups() default { };
+
+  Class<? extends Payload>[] payload() default { };
+
+  class Validator implements ConstraintValidator<ValidBase64URLString, String> {
+
+    @Override
+    public boolean isValid(final String value, final ConstraintValidatorContext context) {
+      if (Objects.isNull(value)) {
+        return true;
+      }
+      try {
+        Base64.getUrlDecoder().decode(value);
+        return true;
+      } catch (IllegalArgumentException e) {
+        return false;
+      }
+    }
+  }
+}
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java
index 138ee5e4b..d427e1c72 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ArchiveControllerTest.java
@@ -594,6 +594,27 @@ public class ArchiveControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
   }
 
+  @Test
+  public void invalidSourceAttachmentKey() throws VerificationFailedException {
+    final BackupAuthCredentialPresentation presentation = backupAuthTestUtil.getPresentation(
+        BackupLevel.MEDIA, backupKey, aci);
+    when(backupManager.authenticateBackupUser(any(), any()))
+        .thenReturn(CompletableFuture.completedFuture(backupUser(presentation.getBackupId(), BackupLevel.MEDIA)));
+    final Response r = resources.getJerseyTest()
+        .target("v1/archives/media")
+        .request()
+        .header("X-Signal-ZK-Auth", Base64.getEncoder().encodeToString(presentation.serialize()))
+        .header("X-Signal-ZK-Auth-Signature", "aaa")
+        .put(Entity.json(new ArchiveController.CopyMediaRequest(
+                new ArchiveController.RemoteAttachment(3, "invalid/urlBase64"),
+                100,
+                TestRandomUtil.nextBytes(15),
+                TestRandomUtil.nextBytes(32),
+                TestRandomUtil.nextBytes(32),
+                TestRandomUtil.nextBytes(16))));
+    assertThat(r.getStatus()).isEqualTo(422);
+  }
+
   private static AuthenticatedBackupUser backupUser(byte[] backupId, BackupLevel backupLevel) {
     return new AuthenticatedBackupUser(backupId, backupLevel, "myBackupDir", "myMediaDir");
   }