From 28076335e074afb309be4bf876e90cefbfa7feff Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Mon, 1 Aug 2022 17:29:58 -0400
Subject: [PATCH] Generate APNs payloads using a payload builder

---
 .../textsecuregcm/push/APNSender.java         | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/push/APNSender.java b/service/src/main/java/org/whispersystems/textsecuregcm/push/APNSender.java
index 34711da1a..24813ffb7 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/push/APNSender.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/push/APNSender.java
@@ -9,6 +9,7 @@ import com.eatthepath.pushy.apns.ApnsClientBuilder;
 import com.eatthepath.pushy.apns.DeliveryPriority;
 import com.eatthepath.pushy.apns.PushType;
 import com.eatthepath.pushy.apns.auth.ApnsSigningKey;
+import com.eatthepath.pushy.apns.util.SimpleApnsPayloadBuilder;
 import com.eatthepath.pushy.apns.util.SimpleApnsPushNotification;
 import com.google.common.annotations.VisibleForTesting;
 import io.dropwizard.lifecycle.Managed;
@@ -28,16 +29,16 @@ public class APNSender implements Managed, PushNotificationSender {
   private final ApnsClient apnsClient;
 
   @VisibleForTesting
-  static final String APN_VOIP_NOTIFICATION_PAYLOAD    = "{\"aps\":{\"sound\":\"default\",\"alert\":{\"loc-key\":\"APN_Message\"}}}";
+  static final String APN_VOIP_NOTIFICATION_PAYLOAD = new SimpleApnsPayloadBuilder()
+      .setSound("default")
+      .setLocalizedAlertMessage("APN_Message")
+      .build();
 
   @VisibleForTesting
-  static final String APN_NSE_NOTIFICATION_PAYLOAD     = "{\"aps\":{\"mutable-content\":1,\"alert\":{\"loc-key\":\"APN_Message\"}}}";
-
-  @VisibleForTesting
-  static final String APN_CHALLENGE_PAYLOAD            = "{\"aps\":{\"sound\":\"default\",\"alert\":{\"loc-key\":\"APN_Message\"}}, \"challenge\" : \"%s\"}";
-
-  @VisibleForTesting
-  static final String APN_RATE_LIMIT_CHALLENGE_PAYLOAD = "{\"aps\":{\"sound\":\"default\",\"alert\":{\"loc-key\":\"APN_Message\"}}, \"rateLimitChallenge\" : \"%s\"}";
+  static final String APN_NSE_NOTIFICATION_PAYLOAD = new SimpleApnsPayloadBuilder()
+      .setMutableContent(true)
+      .setLocalizedAlertMessage("APN_Message")
+      .build();
 
   @VisibleForTesting
   static final Instant MAX_EXPIRATION = Instant.ofEpochMilli(Integer.MAX_VALUE * 1000L);
@@ -76,8 +77,18 @@ public class APNSender implements Managed, PushNotificationSender {
 
     final String payload = switch (notification.notificationType()) {
       case NOTIFICATION -> isVoip ? APN_VOIP_NOTIFICATION_PAYLOAD : APN_NSE_NOTIFICATION_PAYLOAD;
-      case CHALLENGE -> String.format(APN_CHALLENGE_PAYLOAD, notification.data());
-      case RATE_LIMIT_CHALLENGE -> String.format(APN_RATE_LIMIT_CHALLENGE_PAYLOAD, notification.data());
+
+      case CHALLENGE -> new SimpleApnsPayloadBuilder()
+          .setSound("default")
+          .setLocalizedAlertMessage("APN_Message")
+          .addCustomProperty("challenge", notification.data())
+          .build();
+
+      case RATE_LIMIT_CHALLENGE -> new SimpleApnsPayloadBuilder()
+          .setSound("default")
+          .setLocalizedAlertMessage("APN_Message")
+          .addCustomProperty("rateLimitChallenge", notification.data())
+          .build();
     };
 
     final String collapseId =