From 16dba09b61e1e145a79f8a3b8b677da89dca187c Mon Sep 17 00:00:00 2001
From: Ehren Kret <ehren@signal.org>
Date: Wed, 15 Sep 2021 17:02:01 -0500
Subject: [PATCH] Handle merging badges when adding to account

---
 .../textsecuregcm/storage/Account.java        | 19 +++++++++++++-
 .../textsecuregcm/storage/AccountBadge.java   | 26 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
index 07914aba3..a00a9b3f9 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
@@ -324,7 +324,24 @@ public class Account {
   public void addBadge(AccountBadge badge) {
     requireNotStale();
 
-    badges.add(badge);
+    boolean added = false;
+    for (int i = 0; i < badges.size(); i++) {
+      AccountBadge badgeInList = badges.get(i);
+      if (Objects.equals(badgeInList.getId(), badge.getId())) {
+        if (added) {
+          badges.remove(i);
+          i--;
+        } else {
+          badges.set(i, badgeInList.mergeWith(badge));
+          added = true;
+        }
+      }
+    }
+
+    if (!added) {
+      badges.add(badge);
+    }
+
     purgeStaleBadges();
   }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountBadge.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountBadge.java
index 16c7bad03..5aa459fde 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountBadge.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountBadge.java
@@ -26,6 +26,32 @@ public class AccountBadge {
     this.visible = visible;
   }
 
+  /**
+   * Returns a new AccountBadge that is a merging of the two originals. IDs must match for this operation to make sense.
+   * The expiration will be the later of the two.
+   * Visibility will be set if either of the passed in objects is visible.
+   */
+  public AccountBadge mergeWith(AccountBadge other) {
+    if (!Objects.equals(other.id, id)) {
+      throw new IllegalArgumentException("merging badges should only take place for same id");
+    }
+
+    final Instant latestExpiration;
+    if (expiration == null || other.expiration == null) {
+      latestExpiration = null;
+    } else if (expiration.isAfter(other.expiration)) {
+      latestExpiration = expiration;
+    } else {
+      latestExpiration = other.expiration;
+    }
+
+    return new AccountBadge(
+        id,
+        latestExpiration,
+        visible || other.visible
+    );
+  }
+
   public String getId() {
     return id;
   }