From d415a3365a0bfddcb937d9eefa3c2a01e6205c36 Mon Sep 17 00:00:00 2001
From: Jonathan Klabunde Tomer <jkt@signal.org>
Date: Wed, 29 Jul 2026 10:48:30 -0700
Subject: [PATCH] add optional getter for PNI identity key

---
 .../textsecuregcm/storage/Account.java        | 37 ++++++++++++++++++-
 1 file changed, 36 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 3696f6a93..91443490b 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Account.java
@@ -91,6 +91,7 @@ public class Account {
   @JsonProperty("pniIdentityKey")
   @JsonSerialize(using = IdentityKeyAdapter.Serializer.class)
   @JsonDeserialize(using = IdentityKeyAdapter.Deserializer.class)
+  @Nullable
   private IdentityKey phoneNumberIdentityKey;
 
   @JsonProperty("cpv")
@@ -377,16 +378,50 @@ public class Account {
     this.identityKey = identityKey;
   }
 
+  /// Returns an identity key for the given identity type for this account with the assumption that all accounts have
+  /// identity keys for all identity types.
+  ///
+  /// @param identityType the identity type for which to retrieve an identity key
+  ///
+  /// @return the identity key for the given identity type
+  ///
+  /// @throws NoSuchElementException if the account does not have an identifier (and therefore identity key) for the given identity type
+  ///
+  /// @deprecated Different identity types have significantly differing existence requirements/guarantees
+  /// for their respective identity keys. Please use [#getAccountIdentityKey()] or
+  /// [#getPhoneNumberIdentityKeyOptional()] instead.
   public IdentityKey getIdentityKey(final IdentityType identityType) {
     requireNotStale();
 
     return switch (identityType) {
       case ACI -> identityKey;
-      case PNI -> phoneNumberIdentityKey;
+      case PNI -> Optional.ofNullable(phoneNumberIdentityKey).orElseThrow(NoSuchElementException::new);
     };
   }
 
+  /// Returns an identity key for the ACI identity for this account.
+  public IdentityKey getAccountIdentityKey() {
+    requireNotStale();
+    return identityKey;
+  }
+
+  /// Returns an identity key for the phone-number identity for this account, if it has such an identity.
+  ///
+  /// @return the identity key for the PNI identity for the account if it has one, or an empty `Optional` otherwise.
+  public Optional<IdentityKey> getPhoneNumberIdentityKey() {
+    requireNotStale();
+    return Optional.ofNullable(phoneNumberIdentityKey);
+  }
+
+  /// Sets the identity key for the phone-number identity of this account.
+  ///
+  /// @throws IllegalStateException if the account does not have a phone number identifier.
   public void setPhoneNumberIdentityKey(final IdentityKey phoneNumberIdentityKey) {
+    requireNotStale();
+
+    if (this.phoneNumberIdentifier == null) {
+      throw new IllegalStateException();
+    }
     this.phoneNumberIdentityKey = phoneNumberIdentityKey;
   }