Avoid some 401 errors during story sends.

app/src/main/java/org/thoughtcrime/securesms/crypto/UnidentifiedAccessUtil.java

@@ -96,7 +96,6 @@ public class UnidentifiedAccessUtil {
     Map<CertificateType, Integer> typeCounts = new HashMap<>();
 
     for (Recipient recipient : recipients) {
-      byte[]          theirUnidentifiedAccessKey       = getTargetUnidentifiedAccessKey(recipient, isForStory);
       CertificateType    certificateType                  = getUnidentifiedAccessCertificateType(recipient);
       byte[]             ourUnidentifiedAccessCertificate = SignalStore.certificateValues().getUnidentifiedAccessCertificate(certificateType);
 
@@ -104,17 +103,22 @@ public class UnidentifiedAccessUtil {
       typeCount++;
       typeCounts.put(certificateType, typeCount);
 
-      if (theirUnidentifiedAccessKey != null && ourUnidentifiedAccessCertificate != null) {
+      if (ourUnidentifiedAccessCertificate != null) {
         try {
-          access.add(Optional.of(new UnidentifiedAccessPair(new UnidentifiedAccess(theirUnidentifiedAccessKey,
+          UnidentifiedAccess theirAccess = getTargetUnidentifiedAccess(recipient, ourUnidentifiedAccessCertificate, isForStory);
-                                                                                   ourUnidentifiedAccessCertificate),
+          UnidentifiedAccess ourAccess   = new UnidentifiedAccess(ourUnidentifiedAccessKey, ourUnidentifiedAccessCertificate, false);
-                                                            new UnidentifiedAccess(ourUnidentifiedAccessKey,
+
-                                                                                   ourUnidentifiedAccessCertificate))));
+          if (theirAccess != null) {
+            access.add(Optional.of(new UnidentifiedAccessPair(theirAccess, ourAccess)));
+          } else {
+            access.add(Optional.empty());
+          }
         } catch (InvalidCertificateException e) {
-          Log.w(TAG, e);
+          Log.w(TAG, "Invalid unidentified access certificate!", e);
           access.add(Optional.empty());
         }
       } else {
+        Log.w(TAG, "Missing unidentified access certificate!");
         access.add(Optional.empty());
       }
     }
@@ -140,9 +144,11 @@ public class UnidentifiedAccessUtil {
 
       if (ourUnidentifiedAccessCertificate != null) {
         return Optional.of(new UnidentifiedAccessPair(new UnidentifiedAccess(ourUnidentifiedAccessKey,
-                                                                             ourUnidentifiedAccessCertificate),
+                                                                             ourUnidentifiedAccessCertificate,
+                                                                             false),
                                                       new UnidentifiedAccess(ourUnidentifiedAccessKey,
-                                                                             ourUnidentifiedAccessCertificate)));
+                                                                             ourUnidentifiedAccessCertificate,
+                                                                             false)));
       }
 
       return Optional.empty();
@@ -168,7 +174,7 @@ public class UnidentifiedAccessUtil {
                       .getUnidentifiedAccessCertificate(getUnidentifiedAccessCertificateType(recipient));
   }
 
-  private static @Nullable byte[] getTargetUnidentifiedAccessKey(@NonNull Recipient recipient, boolean isForStory) {
+  private static @Nullable UnidentifiedAccess getTargetUnidentifiedAccess(@NonNull Recipient recipient, @NonNull byte[] certificate, boolean isForStory) throws InvalidCertificateException {
     ProfileKey theirProfileKey = ProfileKeyUtil.profileKeyOrNull(recipient.resolve().getProfileKey());
 
     byte[] accessKey;
@@ -176,7 +182,11 @@ public class UnidentifiedAccessUtil {
     switch (recipient.resolve().getUnidentifiedAccessMode()) {
       case UNKNOWN:
         if (theirProfileKey == null) {
+          if (isForStory) {
+            accessKey = null;
+          } else {
             accessKey = UNRESTRICTED_KEY;
+          }
         } else {
           accessKey = UnidentifiedAccess.deriveAccessKeyFrom(theirProfileKey);
         }
@@ -199,10 +209,12 @@ public class UnidentifiedAccessUtil {
     }
 
     if (accessKey == null && isForStory) {
-      accessKey = UNRESTRICTED_KEY;
+      return new UnidentifiedAccess(UNRESTRICTED_KEY, certificate, true);
+    } else if (accessKey != null) {
+      return new UnidentifiedAccess(accessKey, certificate, false);
+    } else {
+      return null;
     }
-
-    return accessKey;
   }
 
   private enum CertificateValidatorHolder {

libsignal/service/src/main/java/org/whispersystems/signalservice/api/SignalServiceMessageSender.java

@@ -2378,9 +2378,15 @@ public class SignalServiceMessageSender {
 
   private List<PreKeyBundle> getPreKeys(SignalServiceAddress recipient, Optional<UnidentifiedAccess> unidentifiedAccess, int deviceId, boolean story) throws IOException {
     try {
+      // If it's only unrestricted because it's a story send, then we know it'll fail
+      if (story && unidentifiedAccess.isPresent() && unidentifiedAccess.get().isUnrestrictedForStory()) {
+        unidentifiedAccess = Optional.empty();
+      }
+
       return socket.getPreKeys(recipient, unidentifiedAccess, deviceId);
     } catch (NonSuccessfulResponseCodeException e) {
       if (e.getCode() == 401 && story) {
+        Log.d(TAG, "Got 401 when fetching prekey for story. Trying without UD.");
         return socket.getPreKeys(recipient, Optional.empty(), deviceId);
       } else {
         throw e;

libsignal/service/src/main/java/org/whispersystems/signalservice/api/crypto/UnidentifiedAccess.java

@@ -21,12 +21,19 @@ public class UnidentifiedAccess {
 
   private final byte[]            unidentifiedAccessKey;
   private final SenderCertificate unidentifiedCertificate;
+  private final boolean           isUnrestrictedForStory;
 
-  public UnidentifiedAccess(byte[] unidentifiedAccessKey, byte[] unidentifiedCertificate)
+  /**
+   * @param isUnrestrictedForStory When sending to a story, we always want to use sealed sender. Receivers will accept it for story messages. However, there are
+   *                               some situations where we need to know if this access key will be correct for non-story purposes. Set this flag to true if
+   *                               the access key is a synthetic one that would only be valid for story messages.
+   */
+  public UnidentifiedAccess(byte[] unidentifiedAccessKey, byte[] unidentifiedCertificate, boolean isUnrestrictedForStory)
       throws InvalidCertificateException
   {
     this.unidentifiedAccessKey   = unidentifiedAccessKey;
     this.unidentifiedCertificate = new SenderCertificate(unidentifiedCertificate);
+    this.isUnrestrictedForStory  = isUnrestrictedForStory;
   }
 
   public byte[] getUnidentifiedAccessKey() {
@@ -37,6 +44,10 @@ public class UnidentifiedAccess {
     return unidentifiedCertificate;
   }
 
+  public boolean isUnrestrictedForStory() {
+    return isUnrestrictedForStory;
+  }
+
   public static byte[] deriveAccessKeyFrom(ProfileKey profileKey) {
     try {
       byte[]         nonce  = new byte[12];

microbenchmark/src/androidTest/java/org/signal/util/SignalClient.kt

@@ -125,7 +125,7 @@ class SignalClient {
 
     val outgoingPushMessage: OutgoingPushMessage = cipher.encrypt(
       SignalProtocolAddress(to.serviceId.toString(), 1),
-      Optional.of(UnidentifiedAccess(to.unidentifiedAccessKey, senderCertificate.serialized)),
+      Optional.of(UnidentifiedAccess(to.unidentifiedAccessKey, senderCertificate.serialized, false)),
       EnvelopeContent.encrypted(content, ContentHint.RESENDABLE, Optional.empty())
     )