Add the ability to migrate to new KBS enclaves.

app/build.gradle

@@ -132,9 +132,10 @@ android {
         buildConfigField "int", "CONTENT_PROXY_PORT", "443"
         buildConfigField "String", "SIGNAL_AGENT", "\"OWA\""
         buildConfigField "String", "CDS_MRENCLAVE", "\"c98e00a4e3ff977a56afefe7362a27e4961e4f19e211febfbb19b897e6b80b15\""
-        buildConfigField "String", "KBS_ENCLAVE_NAME", "\"fe7c1bfae98f9b073d220366ea31163ee82f6d04bead774f71ca8e5c40847bfe\""
+        buildConfigField "KbsEnclave", "KBS_ENCLAVE", "new KbsEnclave(\"fe7c1bfae98f9b073d220366ea31163ee82f6d04bead774f71ca8e5c40847bfe\"," +
-        buildConfigField "String", "KBS_SERVICE_ID", "\"fe7c1bfae98f9b073d220366ea31163ee82f6d04bead774f71ca8e5c40847bfe\""
+                                                                     "\"fe7c1bfae98f9b073d220366ea31163ee82f6d04bead774f71ca8e5c40847bfe\", " +
-        buildConfigField "String", "KBS_MRENCLAVE", "\"a3baab19ef6ce6f34ab9ebb25ba722725ae44a8872dc0ff08ad6d83a9489de87\""
+                                                                     "\"a3baab19ef6ce6f34ab9ebb25ba722725ae44a8872dc0ff08ad6d83a9489de87\")";
+        buildConfigField "KbsEnclave[]", "KBS_FALLBACKS", "new KbsEnclave[0]"
         buildConfigField "String", "UNIDENTIFIED_SENDER_TRUST_ROOT", "\"BXu6QIKVz5MA8gstzfOgRQGqyLqOwNKHL6INkv3IHWMF\""
         buildConfigField "String", "ZKGROUP_SERVER_PUBLIC_PARAMS", "\"AMhf5ywVwITZMsff/eCyudZx9JDmkkkbV6PInzG4p8x3VqVJSFiMvnvlEKWuRob/1eaIetR31IYeAbm0NdOuHH8Qi+Rexi1wLlpzIo1gstHWBfZzy1+qHRV5A4TqPp15YzBPm0WSggW6PbSn+F4lf57VCnHF7p8SvzAA2ZZJPYJURt8X7bbg+H3i+PEjH9DXItNEqs2sNcug37xZQDLm7X0=\""
         buildConfigField "String[]", "LANGUAGES", "new String[]{\"" + autoResConfig().collect { s -> s.replace('-r', '_') }.join('", "') + '"}'
@@ -214,8 +215,10 @@ android {
             buildConfigField "String", "SIGNAL_CONTACT_DISCOVERY_URL", "\"https://api-staging.directory.signal.org\""
             buildConfigField "String", "SIGNAL_KEY_BACKUP_URL", "\"https://api-staging.backup.signal.org\""
             buildConfigField "String", "CDS_MRENCLAVE", "\"bd123560b01c8fa92935bc5ae15cd2064e5c45215f23f0bd40364d521329d2ad\""
-            buildConfigField "String", "KBS_ENCLAVE_NAME", "\"823a3b2c037ff0cbe305cc48928cfcc97c9ed4a8ca6d49af6f7d6981fb60a4e9\""
+            buildConfigField "KbsEnclave", "KBS_ENCLAVE", "new KbsEnclave(\"823a3b2c037ff0cbe305cc48928cfcc97c9ed4a8ca6d49af6f7d6981fb60a4e9\", " +
-            buildConfigField "String", "KBS_SERVICE_ID", "\"038c40bbbacdc873caa81ac793bb75afde6dfe436a99ab1f15e3f0cbb7434ced\""
+                                                                         "\"038c40bbbacdc873caa81ac793bb75afde6dfe436a99ab1f15e3f0cbb7434ced\", " +
+                                                                         "\"a3baab19ef6ce6f34ab9ebb25ba722725ae44a8872dc0ff08ad6d83a9489de87\")"
+            buildConfigField "KbsEnclave[]", "KBS_FALLBACKS", "new KbsEnclave[0]"
             buildConfigField "String", "UNIDENTIFIED_SENDER_TRUST_ROOT", "\"BbqY1DzohE4NUZoVF+L18oUPrK3kILllLEJh2UnPSsEx\""
             buildConfigField "String", "ZKGROUP_SERVER_PUBLIC_PARAMS", "\"ABSY21VckQcbSXVNCGRYJcfWHiAMZmpTtTELcDmxgdFbtp/bWsSxZdMKzfCp8rvIs8ocCU3B37fT3r4Mi5qAemeGeR2X+/YmOGR5ofui7tD5mDQfstAI9i+4WpMtIe8KC3wU5w3Inq3uNWVmoGtpKndsNfwJrCg0Hd9zmObhypUnSkfYn2ooMOOnBpfdanRtrvetZUayDMSC5iSRcXKpdls=\""
         }

app/src/main/java/org/thoughtcrime/securesms/KbsEnclave.java

@@ -0,0 +1,49 @@
+package org.thoughtcrime.securesms;
+
+import androidx.annotation.NonNull;
+
+import java.util.Objects;
+
+/**
+ * Used in our {@link BuildConfig} to tie together the various attributes of a KBS instance. This
+ * is sitting in the root directory so it can be accessed by the build config.
+ */
+public final class KbsEnclave {
+
+  private final String enclaveName;
+  private final String serviceId;
+  private final String mrEnclave;
+
+  public KbsEnclave(@NonNull String enclaveName, @NonNull String serviceId, @NonNull String mrEnclave) {
+    this.enclaveName = enclaveName;
+    this.serviceId   = serviceId;
+    this.mrEnclave   = mrEnclave;
+  }
+
+  public @NonNull String getMrEnclave() {
+    return mrEnclave;
+  }
+
+  public @NonNull String getEnclaveName() {
+    return enclaveName;
+  }
+
+  public @NonNull String getServiceId() {
+    return serviceId;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    KbsEnclave enclave = (KbsEnclave) o;
+    return enclaveName.equals(enclave.enclaveName) &&
+           serviceId.equals(enclave.serviceId)     &&
+           mrEnclave.equals(enclave.mrEnclave);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(enclaveName, serviceId, mrEnclave);
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/dependencies/ApplicationDependencies.java

@@ -6,6 +6,7 @@ import androidx.annotation.MainThread;
 import androidx.annotation.NonNull;
 
 import org.thoughtcrime.securesms.BuildConfig;
+import org.thoughtcrime.securesms.KbsEnclave;
 import org.thoughtcrime.securesms.messages.IncomingMessageProcessor;
 import org.thoughtcrime.securesms.messages.BackgroundMessageRetriever;
 import org.thoughtcrime.securesms.groups.GroupsV2AuthorizationMemoryValueCache;
@@ -15,6 +16,7 @@ import org.thoughtcrime.securesms.keyvalue.KeyValueStore;
 import org.thoughtcrime.securesms.keyvalue.SignalStore;
 import org.thoughtcrime.securesms.megaphone.MegaphoneRepository;
 import org.thoughtcrime.securesms.notifications.MessageNotifier;
+import org.thoughtcrime.securesms.pin.KbsEnclaves;
 import org.thoughtcrime.securesms.push.SignalServiceNetworkAccess;
 import org.thoughtcrime.securesms.recipients.LiveRecipientCache;
 import org.thoughtcrime.securesms.messages.IncomingMessageObserver;
@@ -111,11 +113,11 @@ public class ApplicationDependencies {
     return groupsV2Operations;
   }
 
-  public static synchronized @NonNull KeyBackupService getKeyBackupService() {
+  public static synchronized @NonNull KeyBackupService getKeyBackupService(@NonNull KbsEnclave enclave) {
     return getSignalServiceAccountManager().getKeyBackupService(IasKeyStore.getIasKeyStore(application),
-                                                                BuildConfig.KBS_ENCLAVE_NAME,
+                                                                enclave.getEnclaveName(),
-                                                                Hex.fromStringOrThrow(BuildConfig.KBS_SERVICE_ID),
+                                                                Hex.fromStringOrThrow(enclave.getServiceId()),
-                                                                BuildConfig.KBS_MRENCLAVE,
+                                                                enclave.getMrEnclave(),
                                                                 10);
   }
 

app/src/main/java/org/thoughtcrime/securesms/jobs/ClearFallbackKbsEnclaveJob.java

@@ -0,0 +1,102 @@
+package org.thoughtcrime.securesms.jobs;
+
+import androidx.annotation.NonNull;
+
+import org.thoughtcrime.securesms.KbsEnclave;
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
+import org.thoughtcrime.securesms.jobmanager.Data;
+import org.thoughtcrime.securesms.jobmanager.Job;
+import org.thoughtcrime.securesms.jobmanager.JobManager;
+import org.thoughtcrime.securesms.jobmanager.impl.NetworkConstraint;
+import org.thoughtcrime.securesms.logging.Log;
+import org.thoughtcrime.securesms.pin.KbsEnclaves;
+import org.whispersystems.signalservice.api.push.exceptions.PushNetworkException;
+import org.whispersystems.signalservice.internal.contacts.crypto.UnauthenticatedResponseException;
+
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Clears data from an old KBS enclave.
+ */
+public class ClearFallbackKbsEnclaveJob extends BaseJob {
+
+  public static final String KEY = "ClearFallbackKbsEnclaveJob";
+
+  private static final String TAG = Log.tag(ClearFallbackKbsEnclaveJob.class);
+
+  private static final String KEY_ENCLAVE_NAME = "enclaveName";
+  private static final String KEY_SERVICE_ID   = "serviceId";
+  private static final String KEY_MR_ENCLAVE   = "mrEnclave";
+
+  private final KbsEnclave enclave;
+
+  ClearFallbackKbsEnclaveJob(@NonNull KbsEnclave enclave) {
+    this(new Parameters.Builder()
+                       .addConstraint(NetworkConstraint.KEY)
+                       .setLifespan(TimeUnit.DAYS.toMillis(90))
+                       .setMaxAttempts(Parameters.UNLIMITED)
+                       .setQueue("ClearFallbackKbsEnclaveJob")
+                       .build(),
+        enclave);
+  }
+
+  public static void clearAll() {
+    if (KbsEnclaves.fallbacks().isEmpty()) {
+      Log.i(TAG, "No fallbacks!");
+      return;
+    }
+
+    JobManager jobManager = ApplicationDependencies.getJobManager();
+
+    for (KbsEnclave enclave : KbsEnclaves.fallbacks()) {
+      jobManager.add(new ClearFallbackKbsEnclaveJob(enclave));
+    }
+  }
+
+  private ClearFallbackKbsEnclaveJob(@NonNull Parameters parameters, @NonNull KbsEnclave enclave) {
+    super(parameters);
+    this.enclave = enclave;
+  }
+
+  @Override
+  public @NonNull String getFactoryKey() {
+    return KEY;
+  }
+
+  @Override
+  public @NonNull Data serialize() {
+    return new Data.Builder().putString(KEY_ENCLAVE_NAME, enclave.getEnclaveName())
+                             .putString(KEY_SERVICE_ID, enclave.getServiceId())
+                             .putString(KEY_MR_ENCLAVE, enclave.getMrEnclave())
+                             .build();
+  }
+
+  @Override
+  public void onRun() throws IOException, UnauthenticatedResponseException {
+    Log.i(TAG, "Preparing to delete data from " + enclave.getEnclaveName());
+    ApplicationDependencies.getKeyBackupService(enclave).newPinChangeSession().removePin();
+    Log.i(TAG, "Successfully deleted the data from " + enclave.getEnclaveName());
+  }
+
+  @Override
+  public boolean onShouldRetry(@NonNull Exception e) {
+    return true;
+  }
+
+  @Override
+  public void onFailure() {
+    throw new AssertionError("This job should never fail. " + getClass().getSimpleName());
+  }
+
+  public static class Factory implements Job.Factory<ClearFallbackKbsEnclaveJob> {
+    @Override
+    public @NonNull ClearFallbackKbsEnclaveJob create(@NonNull Parameters parameters, @NonNull Data data) {
+      KbsEnclave enclave = new KbsEnclave(data.getString(KEY_ENCLAVE_NAME),
+                                          data.getString(KEY_SERVICE_ID),
+                                          data.getString(KEY_MR_ENCLAVE));
+
+      return new ClearFallbackKbsEnclaveJob(parameters, enclave);
+    }
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/jobs/JobManagerFactories.java

@@ -30,6 +30,7 @@ import org.thoughtcrime.securesms.migrations.AvatarIdRemovalMigrationJob;
 import org.thoughtcrime.securesms.migrations.AvatarMigrationJob;
 import org.thoughtcrime.securesms.migrations.CachedAttachmentsMigrationJob;
 import org.thoughtcrime.securesms.migrations.DatabaseMigrationJob;
+import org.thoughtcrime.securesms.migrations.KbsEnclaveMigrationJob;
 import org.thoughtcrime.securesms.migrations.LegacyMigrationJob;
 import org.thoughtcrime.securesms.migrations.MigrationCompleteJob;
 import org.thoughtcrime.securesms.migrations.PassingMigrationJob;
@@ -62,9 +63,11 @@ public final class JobManagerFactories {
       put(AvatarGroupsV1DownloadJob.KEY,             new AvatarGroupsV1DownloadJob.Factory());
       put(AvatarGroupsV2DownloadJob.KEY,             new AvatarGroupsV2DownloadJob.Factory());
       put(CleanPreKeysJob.KEY,                       new CleanPreKeysJob.Factory());
+      put(ClearFallbackKbsEnclaveJob.KEY,            new ClearFallbackKbsEnclaveJob.Factory());
       put(CreateSignedPreKeyJob.KEY,                 new CreateSignedPreKeyJob.Factory());
       put(DirectoryRefreshJob.KEY,                   new DirectoryRefreshJob.Factory());
       put(FcmRefreshJob.KEY,                         new FcmRefreshJob.Factory());
+      put(KbsEnclaveMigrationWorkerJob.KEY,          new KbsEnclaveMigrationWorkerJob.Factory());
       put(LeaveGroupJob.KEY,                         new LeaveGroupJob.Factory());
       put(LocalBackupJob.KEY,                        new LocalBackupJob.Factory());
       put(MmsDownloadJob.KEY,                        new MmsDownloadJob.Factory());
@@ -132,6 +135,7 @@ public final class JobManagerFactories {
       put(AvatarMigrationJob.KEY,                    new AvatarMigrationJob.Factory());
       put(CachedAttachmentsMigrationJob.KEY,         new CachedAttachmentsMigrationJob.Factory());
       put(DatabaseMigrationJob.KEY,                  new DatabaseMigrationJob.Factory());
+      put(KbsEnclaveMigrationJob.KEY,                new KbsEnclaveMigrationJob.Factory());
       put(LegacyMigrationJob.KEY,                    new LegacyMigrationJob.Factory());
       put(MigrationCompleteJob.KEY,                  new MigrationCompleteJob.Factory());
       put(PinOptOutMigration.KEY,                    new PinOptOutMigration.Factory());

app/src/main/java/org/thoughtcrime/securesms/jobs/KbsEnclaveMigrationWorkerJob.java

@@ -0,0 +1,85 @@
+package org.thoughtcrime.securesms.jobs;
+
+import androidx.annotation.NonNull;
+
+import org.thoughtcrime.securesms.jobmanager.Data;
+import org.thoughtcrime.securesms.jobmanager.Job;
+import org.thoughtcrime.securesms.jobmanager.impl.NetworkConstraint;
+import org.thoughtcrime.securesms.keyvalue.SignalStore;
+import org.thoughtcrime.securesms.logging.Log;
+import org.thoughtcrime.securesms.migrations.KbsEnclaveMigrationJob;
+import org.thoughtcrime.securesms.pin.PinState;
+import org.whispersystems.signalservice.internal.contacts.crypto.UnauthenticatedResponseException;
+
+import java.io.IOException;
+
+/**
+ * Should only be enqueued by {@link KbsEnclaveMigrationJob}. Does the actual work of migrating KBS
+ * data to the new enclave and deleting it from the old enclave(s).
+ */
+public class KbsEnclaveMigrationWorkerJob extends BaseJob {
+
+  public static final String KEY = "KbsEnclaveMigrationWorkerJob";
+
+  private static final String TAG = Log.tag(KbsEnclaveMigrationWorkerJob.class);
+
+  public KbsEnclaveMigrationWorkerJob() {
+    this(new Parameters.Builder()
+                       .addConstraint(NetworkConstraint.KEY)
+                       .setLifespan(Parameters.IMMORTAL)
+                       .setMaxAttempts(Parameters.UNLIMITED)
+                       .setQueue("KbsEnclaveMigrationWorkerJob")
+                       .setMaxInstances(1)
+                       .build());
+  }
+
+  private KbsEnclaveMigrationWorkerJob(@NonNull Parameters parameters) {
+    super(parameters);
+  }
+
+  @Override
+  public @NonNull String getFactoryKey() {
+    return KEY;
+  }
+
+  @Override
+  public @NonNull Data serialize() {
+    return Data.EMPTY;
+  }
+
+  @Override
+  public void onRun() throws IOException, UnauthenticatedResponseException {
+    String pin = SignalStore.kbsValues().getPin();
+
+    if (SignalStore.kbsValues().hasOptedOut()) {
+      Log.w(TAG, "Opted out of KBS! Nothing to migrate.");
+      return;
+    }
+
+    if (pin == null) {
+      Log.w(TAG, "No PIN available! Can't migrate!");
+      return;
+    }
+
+    PinState.onMigrateToNewEnclave(pin);
+    Log.i(TAG, "Migration successful!");
+  }
+
+  @Override
+  public boolean onShouldRetry(@NonNull Exception e) {
+    return e instanceof IOException ||
+           e instanceof UnauthenticatedResponseException;
+  }
+
+  @Override
+  public void onFailure() {
+    throw new AssertionError("This job should never fail. " + getClass().getSimpleName());
+  }
+
+  public static class Factory implements Job.Factory<KbsEnclaveMigrationWorkerJob> {
+    @Override
+    public @NonNull KbsEnclaveMigrationWorkerJob create(@NonNull Parameters parameters, @NonNull Data data) {
+      return new KbsEnclaveMigrationWorkerJob(parameters);
+    }
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/keyvalue/KbsValues.java

@@ -137,6 +137,10 @@ public final class KbsValues extends SignalStoreValues {
     }
   }
 
+  public synchronized @Nullable String getPin() {
+    return getString(PIN, null);
+  }
+
   public synchronized @Nullable String getLocalPinHash() {
     return getString(LOCK_LOCAL_PIN_HASH, null);
   }

app/src/main/java/org/thoughtcrime/securesms/migrations/KbsEnclaveMigrationJob.java

@@ -0,0 +1,53 @@
+package org.thoughtcrime.securesms.migrations;
+
+import androidx.annotation.NonNull;
+
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
+import org.thoughtcrime.securesms.jobmanager.Data;
+import org.thoughtcrime.securesms.jobmanager.Job;
+import org.thoughtcrime.securesms.jobs.KbsEnclaveMigrationWorkerJob;
+
+/**
+ * A job to be run whenever we add a new KBS enclave. In order to prevent this moderately-expensive
+ * task from blocking the network for too long, this task simply enqueues another non-migration job,
+ * {@link KbsEnclaveMigrationWorkerJob}, to do the heavy lifting.
+ */
+public class KbsEnclaveMigrationJob extends MigrationJob {
+
+  public static final String KEY = "KbsEnclaveMigrationJob";
+
+  KbsEnclaveMigrationJob() {
+    this(new Parameters.Builder().build());
+  }
+
+  private KbsEnclaveMigrationJob(@NonNull Parameters parameters) {
+    super(parameters);
+  }
+
+  @Override
+  public boolean isUiBlocking() {
+    return false;
+  }
+
+  @Override
+  public @NonNull String getFactoryKey() {
+    return KEY;
+  }
+
+  @Override
+  public void performMigration() {
+    ApplicationDependencies.getJobManager().add(new KbsEnclaveMigrationWorkerJob());
+  }
+
+  @Override
+  boolean shouldRetry(@NonNull Exception e) {
+    return false;
+  }
+
+  public static class Factory implements Job.Factory<KbsEnclaveMigrationJob> {
+    @Override
+    public @NonNull KbsEnclaveMigrationJob create(@NonNull Parameters parameters, @NonNull Data data) {
+      return new KbsEnclaveMigrationJob(parameters);
+    }
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/pin/KbsEnclaves.java

@@ -0,0 +1,26 @@
+package org.thoughtcrime.securesms.pin;
+
+import androidx.annotation.NonNull;
+
+import org.thoughtcrime.securesms.BuildConfig;
+import org.thoughtcrime.securesms.KbsEnclave;
+import org.thoughtcrime.securesms.util.Util;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+public final class KbsEnclaves {
+
+  public static @NonNull KbsEnclave current() {
+    return BuildConfig.KBS_ENCLAVE;
+  }
+
+  public static @NonNull List<KbsEnclave> all() {
+    return Util.join(Collections.singletonList(BuildConfig.KBS_ENCLAVE), fallbacks());
+  }
+
+  public static @NonNull List<KbsEnclave> fallbacks() {
+    return Arrays.asList(BuildConfig.KBS_FALLBACKS);
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/pin/PinRestoreRepository.java

@@ -1,11 +1,15 @@
 package org.thoughtcrime.securesms.pin;
 
-import androidx.annotation.NonNull;
+import android.os.Parcel;
+import android.os.Parcelable;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import org.thoughtcrime.securesms.KbsEnclave;
 import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
 import org.thoughtcrime.securesms.jobs.StorageAccountRestoreJob;
 import org.thoughtcrime.securesms.jobs.StorageSyncJob;
-import org.thoughtcrime.securesms.keyvalue.SignalStore;
 import org.thoughtcrime.securesms.logging.Log;
 import org.thoughtcrime.securesms.registration.service.KeyBackupSystemWrongPinException;
 import org.thoughtcrime.securesms.util.Stopwatch;
@@ -21,32 +25,58 @@ import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 
-class PinRestoreRepository {
+public class PinRestoreRepository {
 
   private static final String TAG = Log.tag(PinRestoreRepository.class);
 
-  private final Executor         executor = SignalExecutors.UNBOUNDED;
+  private final Executor executor = SignalExecutors.UNBOUNDED;
-  private final KeyBackupService kbs      = ApplicationDependencies.getKeyBackupService();
 
   void getToken(@NonNull Callback<Optional<TokenData>> callback) {
     executor.execute(() -> {
       try {
-        String        authorization = kbs.getAuthorization();
+        callback.onComplete(Optional.fromNullable(getTokenSync(null)));
-        TokenResponse token         = kbs.getToken(authorization);
-        TokenData     tokenData     = new TokenData(authorization, token);
-        callback.onComplete(Optional.of(tokenData));
       } catch (IOException e) {
         callback.onComplete(Optional.absent());
       }
     });
   }
 
+  /**
+   * @param authorization If this is being called before the user is registered (i.e. as part of
+   *                      reglock), you must pass in an authorization token that can be used to
+   *                      retrieve a backup. Otherwise, pass in null and we'll fetch one.
+   */
+  public @NonNull TokenData getTokenSync(@Nullable String authorization) throws IOException {
+    TokenData firstKnownTokenData = null;
+
+    for (KbsEnclave enclave : KbsEnclaves.all()) {
+      KeyBackupService kbs = ApplicationDependencies.getKeyBackupService(enclave);
+
+      authorization = authorization == null ? kbs.getAuthorization() : authorization;
+
+      TokenResponse token     = kbs.getToken(authorization);
+      TokenData     tokenData = new TokenData(enclave, authorization, token);
+
+      if (tokenData.getTriesRemaining() > 0) {
+        Log.i(TAG, "Found data! " + enclave.getEnclaveName());
+        return tokenData;
+      } else if (firstKnownTokenData == null) {
+        Log.i(TAG, "No data, but storing as the first response. " + enclave.getEnclaveName());
+        firstKnownTokenData = tokenData;
+      } else {
+        Log.i(TAG, "No data, and we already have a 'first response'. " + enclave.getEnclaveName());
+      }
+    }
+
+    return Objects.requireNonNull(firstKnownTokenData);
+  }
+
   void submitPin(@NonNull String pin, @NonNull TokenData tokenData, @NonNull Callback<PinResultData> callback) {
     executor.execute(() -> {
       try {
         Stopwatch stopwatch = new Stopwatch("PinSubmission");
 
-        KbsPinData kbsData = PinState.restoreMasterKey(pin, tokenData.basicAuth, tokenData.tokenResponse);
+        KbsPinData kbsData = PinState.restoreMasterKey(pin, tokenData.getEnclave(), tokenData.getBasicAuth(), tokenData.getTokenResponse());
         PinState.onSignalPinRestore(ApplicationDependencies.getApplication(), Objects.requireNonNull(kbsData), pin);
         stopwatch.split("MasterKey");
 
@@ -64,7 +94,7 @@ class PinRestoreRepository {
       } catch (KeyBackupSystemNoDataException e) {
         callback.onComplete(new PinResultData(PinResult.LOCKED, tokenData));
       } catch (KeyBackupSystemWrongPinException e) {
-        callback.onComplete(new PinResultData(PinResult.INCORRECT, new TokenData(tokenData.basicAuth, e.getTokenResponse())));
+        callback.onComplete(new PinResultData(PinResult.INCORRECT, TokenData.withResponse(tokenData, e.getTokenResponse())));
       }
     });
   }
@@ -73,18 +103,81 @@ class PinRestoreRepository {
     void onComplete(@NonNull T value);
   }
 
-  static class TokenData {
+  public static class TokenData implements Parcelable {
+    private final KbsEnclave    enclave;
     private final String        basicAuth;
     private final TokenResponse tokenResponse;
 
-    TokenData(@NonNull String basicAuth, @NonNull TokenResponse tokenResponse) {
+    TokenData(@NonNull KbsEnclave enclave, @NonNull String basicAuth, @NonNull TokenResponse tokenResponse) {
+      this.enclave       = enclave;
       this.basicAuth     = basicAuth;
       this.tokenResponse = tokenResponse;
     }
 
-    int getTriesRemaining() {
+    private TokenData(Parcel in) {
+      //noinspection ConstantConditions
+      this.enclave   = new KbsEnclave(in.readString(), in.readString(), in.readString());
+      this.basicAuth = in.readString();
+
+      byte[] backupId = new byte[0];
+      byte[] token    = new byte[0];
+
+      in.readByteArray(backupId);
+      in.readByteArray(token);
+
+      this.tokenResponse = new TokenResponse(backupId, token, in.readInt());
+    }
+
+    public static @NonNull TokenData withResponse(@NonNull TokenData data, @NonNull TokenResponse response) {
+      return new TokenData(data.getEnclave(), data.getBasicAuth(), response);
+    }
+
+    public int getTriesRemaining() {
       return tokenResponse.getTries();
     }
+
+    public @NonNull String getBasicAuth() {
+      return basicAuth;
+    }
+
+    public @NonNull TokenResponse getTokenResponse() {
+      return tokenResponse;
+    }
+
+    public @NonNull KbsEnclave getEnclave() {
+      return enclave;
+    }
+
+    @Override
+    public int describeContents() {
+      return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+      dest.writeString(enclave.getEnclaveName());
+      dest.writeString(enclave.getServiceId());
+      dest.writeString(enclave.getMrEnclave());
+
+      dest.writeString(basicAuth);
+
+      dest.writeByteArray(tokenResponse.getBackupId());
+      dest.writeByteArray(tokenResponse.getToken());
+      dest.writeInt(tokenResponse.getTries());
+    }
+
+    public static final Creator<TokenData> CREATOR = new Creator<TokenData>() {
+      @Override
+      public TokenData createFromParcel(Parcel in) {
+        return new TokenData(in);
+      }
+
+      @Override
+      public TokenData[] newArray(int size) {
+        return new TokenData[size];
+      }
+    };
+
   }
 
   static class PinResultData {
@@ -92,7 +185,7 @@ class PinRestoreRepository {
     private final TokenData tokenData;
 
     PinResultData(@NonNull PinResult result, @NonNull TokenData tokenData) {
-      this.result = result;
+      this.result    = result;
       this.tokenData = tokenData;
     }
 

app/src/main/java/org/thoughtcrime/securesms/pin/PinState.java

@@ -6,8 +6,11 @@ import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
 
+import org.thoughtcrime.securesms.BuildConfig;
+import org.thoughtcrime.securesms.KbsEnclave;
 import org.thoughtcrime.securesms.dependencies.ApplicationDependencies;
 import org.thoughtcrime.securesms.jobmanager.JobTracker;
+import org.thoughtcrime.securesms.jobs.ClearFallbackKbsEnclaveJob;
 import org.thoughtcrime.securesms.jobs.RefreshAttributesJob;
 import org.thoughtcrime.securesms.jobs.StorageForcePushJob;
 import org.thoughtcrime.securesms.keyvalue.KbsValues;
@@ -26,12 +29,14 @@ import org.whispersystems.signalservice.api.KeyBackupServicePinException;
 import org.whispersystems.signalservice.api.KeyBackupSystemNoDataException;
 import org.whispersystems.signalservice.api.kbs.HashedPin;
 import org.whispersystems.signalservice.api.kbs.MasterKey;
+import org.whispersystems.signalservice.api.push.exceptions.NonSuccessfulResponseCodeException;
 import org.whispersystems.signalservice.internal.contacts.crypto.UnauthenticatedResponseException;
 import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
 
 import java.io.IOException;
 import java.security.SecureRandom;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.TimeUnit;
 
@@ -46,6 +51,7 @@ public final class PinState {
    * Does not affect {@link PinState}.
    */
   public static synchronized @Nullable KbsPinData restoreMasterKey(@Nullable String pin,
+                                                                   @NonNull KbsEnclave enclave,
                                                                    @Nullable String basicStorageCredentials,
                                                                    @NonNull TokenResponse tokenResponse)
     throws IOException, KeyBackupSystemWrongPinException, KeyBackupSystemNoDataException
@@ -58,20 +64,31 @@ public final class PinState {
       throw new AssertionError("Cannot restore KBS key, no storage credentials supplied");
     }
 
-    KeyBackupService keyBackupService = ApplicationDependencies.getKeyBackupService();
+    Log.i(TAG, "Preparing to restore from " + enclave.getEnclaveName());
+    return restoreMasterKeyFromEnclave(enclave, pin, basicStorageCredentials, tokenResponse);
+  }
 
-    Log.i(TAG, "Opening key backup service session");
+  private static @NonNull KbsPinData restoreMasterKeyFromEnclave(@NonNull KbsEnclave enclave,
-    KeyBackupService.RestoreSession session = keyBackupService.newRegistrationSession(basicStorageCredentials, tokenResponse);
+                                                                 @NonNull String pin,
+                                                                 @NonNull String basicStorageCredentials,
+                                                                 @NonNull TokenResponse tokenResponse)
+      throws IOException, KeyBackupSystemWrongPinException, KeyBackupSystemNoDataException
+  {
+    KeyBackupService                keyBackupService = ApplicationDependencies.getKeyBackupService(enclave);
+    KeyBackupService.RestoreSession session          = keyBackupService.newRegistrationSession(basicStorageCredentials, tokenResponse);
 
     try {
       Log.i(TAG, "Restoring pin from KBS");
+
       HashedPin  hashedPin = PinHashing.hashPin(pin, session);
       KbsPinData kbsData   = session.restorePin(hashedPin);
+
       if (kbsData != null) {
         Log.i(TAG, "Found registration lock token on KBS.");
       } else {
         throw new AssertionError("Null not expected");
       }
+
       return kbsData;
     } catch (UnauthenticatedResponseException e) {
       Log.w(TAG, "Failed to restore key", e);
@@ -90,7 +107,7 @@ public final class PinState {
                                                  @Nullable String pin,
                                                  boolean hasPinToRestore)
   {
-    Log.i(TAG, "onNewRegistration()");
+    Log.i(TAG, "onRegistration()");
 
     TextSecurePreferences.setV1RegistrationLockPin(context, pin);
 
@@ -106,7 +123,8 @@ public final class PinState {
       SignalStore.kbsValues().setV2RegistrationLockEnabled(true);
       SignalStore.kbsValues().setKbsMasterKey(kbsData, pin);
       SignalStore.pinValues().resetPinReminders();
-      resetPinRetryCount(context, pin, kbsData);
+      resetPinRetryCount(context, pin);
+      ClearFallbackKbsEnclaveJob.clearAll();
     } else if (hasPinToRestore) {
       Log.i(TAG, "Has a PIN to restore.");
       SignalStore.kbsValues().clearRegistrationLockAndPin();
@@ -131,7 +149,8 @@ public final class PinState {
     SignalStore.kbsValues().setV2RegistrationLockEnabled(false);
     SignalStore.pinValues().resetPinReminders();
     SignalStore.storageServiceValues().setNeedsAccountRestore(false);
-    resetPinRetryCount(context, pin, kbsData);
+    resetPinRetryCount(context, pin);
+    ClearFallbackKbsEnclaveJob.clearAll();
 
     updateState(buildInferredStateFromOtherFields());
   }
@@ -158,7 +177,7 @@ public final class PinState {
     KbsValues                         kbsValues        = SignalStore.kbsValues();
     boolean                           isFirstPin       = !kbsValues.hasPin() || kbsValues.hasOptedOut();
     MasterKey                         masterKey        = kbsValues.getOrCreateMasterKey();
-    KeyBackupService                  keyBackupService = ApplicationDependencies.getKeyBackupService();
+    KeyBackupService                  keyBackupService = ApplicationDependencies.getKeyBackupService(KbsEnclaves.current());
     KeyBackupService.PinChangeSession pinChangeSession = keyBackupService.newPinChangeSession();
     HashedPin                         hashedPin        = PinHashing.hashPin(pin, pinChangeSession);
     KbsPinData                        kbsData          = pinChangeSession.setPin(hashedPin, masterKey);
@@ -217,7 +236,7 @@ public final class PinState {
     assertState(State.PIN_WITH_REGISTRATION_LOCK_DISABLED);
 
     SignalStore.kbsValues().setV2RegistrationLockEnabled(false);
-    ApplicationDependencies.getKeyBackupService()
+    ApplicationDependencies.getKeyBackupService(KbsEnclaves.current())
                            .newPinChangeSession(SignalStore.kbsValues().getRegistrationLockTokenResponse())
                            .enableRegistrationLock(SignalStore.kbsValues().getOrCreateMasterKey());
     SignalStore.kbsValues().setV2RegistrationLockEnabled(true);
@@ -240,7 +259,7 @@ public final class PinState {
     assertState(State.PIN_WITH_REGISTRATION_LOCK_ENABLED);
 
     SignalStore.kbsValues().setV2RegistrationLockEnabled(true);
-    ApplicationDependencies.getKeyBackupService()
+    ApplicationDependencies.getKeyBackupService(KbsEnclaves.current())
                            .newPinChangeSession(SignalStore.kbsValues().getRegistrationLockTokenResponse())
                            .disableRegistrationLock();
     SignalStore.kbsValues().setV2RegistrationLockEnabled(false);
@@ -259,7 +278,7 @@ public final class PinState {
 
     KbsValues                         kbsValues        = SignalStore.kbsValues();
     MasterKey                         masterKey        = kbsValues.getOrCreateMasterKey();
-    KeyBackupService                  keyBackupService = ApplicationDependencies.getKeyBackupService();
+    KeyBackupService                  keyBackupService = ApplicationDependencies.getKeyBackupService(KbsEnclaves.current());
     KeyBackupService.PinChangeSession pinChangeSession = keyBackupService.newPinChangeSession();
     HashedPin                         hashedPin        = PinHashing.hashPin(pin, pinChangeSession);
     KbsPinData                        kbsData          = pinChangeSession.setPin(hashedPin, masterKey);
@@ -272,6 +291,22 @@ public final class PinState {
     updateState(buildInferredStateFromOtherFields());
   }
 
+  /**
+   * Should only be called by {@link org.thoughtcrime.securesms.jobs.KbsEnclaveMigrationWorkerJob}.
+   */
+  @WorkerThread
+  public static synchronized void onMigrateToNewEnclave(@NonNull String pin)
+      throws IOException, UnauthenticatedResponseException
+  {
+    Log.i(TAG, "onMigrateToNewEnclave()");
+    assertState(State.PIN_WITH_REGISTRATION_LOCK_DISABLED, State.PIN_WITH_REGISTRATION_LOCK_ENABLED);
+
+    Log.i(TAG, "Migrating to enclave " + KbsEnclaves.current().getEnclaveName());
+    setPinOnEnclave(KbsEnclaves.current(), pin, SignalStore.kbsValues().getOrCreateMasterKey());
+
+    ClearFallbackKbsEnclaveJob.clearAll();
+  }
+
   @WorkerThread
   private static void bestEffortRefreshAttributes() {
     Optional<JobTracker.JobState> result = ApplicationDependencies.getJobManager().runSynchronously(new RefreshAttributesJob(), TimeUnit.SECONDS.toMillis(10));
@@ -301,23 +336,14 @@ public final class PinState {
   }
 
   @WorkerThread
-  private static void resetPinRetryCount(@NonNull Context context, @Nullable String pin, @NonNull KbsPinData kbsData) {
+  private static void resetPinRetryCount(@NonNull Context context, @Nullable String pin) {
     if (pin == null) {
       return;
     }
 
-    KeyBackupService keyBackupService = ApplicationDependencies.getKeyBackupService();
-
     try {
-      KbsValues                         kbsValues        = SignalStore.kbsValues();
+      setPinOnEnclave(KbsEnclaves.current(), pin, SignalStore.kbsValues().getOrCreateMasterKey());
-      MasterKey                         masterKey        = kbsValues.getOrCreateMasterKey();
-      KeyBackupService.PinChangeSession pinChangeSession = keyBackupService.newPinChangeSession(kbsData.getTokenResponse());
-      HashedPin                         hashedPin        = PinHashing.hashPin(pin, pinChangeSession);
-      KbsPinData                        newData          = pinChangeSession.setPin(hashedPin, masterKey);
-
-      kbsValues.setKbsMasterKey(newData, pin);
       TextSecurePreferences.clearRegistrationLockV1(context);
-
       Log.i(TAG, "Pin set/attempts reset on KBS");
     } catch (IOException e) {
       Log.w(TAG, "May have failed to reset pin attempts!", e);
@@ -326,6 +352,20 @@ public final class PinState {
     }
   }
 
+  @WorkerThread
+  private static @NonNull KbsPinData setPinOnEnclave(@NonNull KbsEnclave enclave, @NonNull String pin, @NonNull MasterKey masterKey)
+      throws IOException, UnauthenticatedResponseException
+  {
+    KeyBackupService                  kbs              = ApplicationDependencies.getKeyBackupService(enclave);
+    KeyBackupService.PinChangeSession pinChangeSession = kbs.newPinChangeSession();
+    HashedPin                         hashedPin        = PinHashing.hashPin(pin, pinChangeSession);
+    KbsPinData                        newData          = pinChangeSession.setPin(hashedPin, masterKey);
+
+    SignalStore.kbsValues().setKbsMasterKey(newData, pin);
+
+    return newData;
+  }
+
   @WorkerThread
   private static void optOutOfPin() {
     SignalStore.kbsValues().optOut();

app/src/main/java/org/thoughtcrime/securesms/registration/fragments/EnterCodeFragment.java

@@ -22,6 +22,7 @@ import org.thoughtcrime.securesms.components.registration.CallMeCountDownView;
 import org.thoughtcrime.securesms.components.registration.VerificationCodeView;
 import org.thoughtcrime.securesms.components.registration.VerificationPinKeyboard;
 import org.thoughtcrime.securesms.logging.Log;
+import org.thoughtcrime.securesms.pin.PinRestoreRepository;
 import org.thoughtcrime.securesms.registration.ReceivedSmsEvent;
 import org.thoughtcrime.securesms.registration.service.CodeVerificationRequest;
 import org.thoughtcrime.securesms.registration.service.RegistrationCodeRequest;
@@ -30,7 +31,6 @@ import org.thoughtcrime.securesms.registration.viewmodel.RegistrationViewModel;
 import org.thoughtcrime.securesms.util.CommunicationActions;
 import org.thoughtcrime.securesms.util.SupportEmailUtil;
 import org.thoughtcrime.securesms.util.concurrent.AssertedSuccessListener;
-import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -107,7 +107,7 @@ public final class EnterCodeFragment extends BaseRegistrationFragment {
 
       RegistrationService registrationService = RegistrationService.getInstance(model.getNumber().getE164Number(), model.getRegistrationSecret());
 
-      registrationService.verifyAccount(requireActivity(), model.getFcmToken(), code, null, null, null,
+      registrationService.verifyAccount(requireActivity(), model.getFcmToken(), code, null, null,
         new CodeVerificationRequest.VerifyCallback() {
 
           @Override
@@ -133,10 +133,9 @@ public final class EnterCodeFragment extends BaseRegistrationFragment {
           }
 
           @Override
-          public void onKbsRegistrationLockPinRequired(long timeRemaining, @NonNull TokenResponse tokenResponse, @NonNull String kbsStorageCredentials) {
+          public void onKbsRegistrationLockPinRequired(long timeRemaining, @NonNull PinRestoreRepository.TokenData tokenData, @NonNull String kbsStorageCredentials) {
             model.setLockedTimeRemaining(timeRemaining);
-            model.setStorageCredentials(kbsStorageCredentials);
+            model.setKeyBackupTokenData(tokenData);
-            model.setKeyBackupCurrentToken(tokenResponse);
             keyboard.displayLocked().addListener(new AssertedSuccessListener<Boolean>() {
               @Override
               public void onSuccess(Boolean r) {
@@ -147,7 +146,7 @@ public final class EnterCodeFragment extends BaseRegistrationFragment {
           }
 
           @Override
-          public void onIncorrectKbsRegistrationLockPin(@NonNull TokenResponse tokenResponse) {
+          public void onIncorrectKbsRegistrationLockPin(@NonNull PinRestoreRepository.TokenData tokenData) {
             throw new AssertionError("Unexpected, user has made no pin guesses");
           }
 

app/src/main/java/org/thoughtcrime/securesms/registration/fragments/RegistrationLockFragment.java

@@ -25,12 +25,12 @@ import org.thoughtcrime.securesms.jobs.StorageAccountRestoreJob;
 import org.thoughtcrime.securesms.keyvalue.SignalStore;
 import org.thoughtcrime.securesms.lock.v2.PinKeyboardType;
 import org.thoughtcrime.securesms.logging.Log;
+import org.thoughtcrime.securesms.pin.PinRestoreRepository.TokenData;
 import org.thoughtcrime.securesms.registration.service.CodeVerificationRequest;
 import org.thoughtcrime.securesms.registration.service.RegistrationService;
 import org.thoughtcrime.securesms.registration.viewmodel.RegistrationViewModel;
 import org.thoughtcrime.securesms.util.ServiceUtil;
 import org.thoughtcrime.securesms.util.concurrent.SimpleTask;
-import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
 
 import java.util.concurrent.TimeUnit;
 
@@ -106,10 +106,10 @@ public final class RegistrationLockFragment extends BaseRegistrationFragment {
     getModel().getLockedTimeRemaining()
               .observe(getViewLifecycleOwner(), t -> timeRemaining = t);
 
-    TokenResponse keyBackupCurrentToken = getModel().getKeyBackupCurrentToken();
+    TokenData keyBackupCurrentToken = getModel().getKeyBackupCurrentToken();
 
     if (keyBackupCurrentToken != null) {
-      int triesRemaining = keyBackupCurrentToken.getTries();
+      int triesRemaining = keyBackupCurrentToken.getTriesRemaining();
       if (triesRemaining <= 3) {
         int daysRemaining = getLockoutDays(timeRemaining);
 
@@ -158,8 +158,7 @@ public final class RegistrationLockFragment extends BaseRegistrationFragment {
 
     RegistrationViewModel model                   = getModel();
     RegistrationService   registrationService     = RegistrationService.getInstance(model.getNumber().getE164Number(), model.getRegistrationSecret());
-    TokenResponse         tokenResponse           = model.getKeyBackupCurrentToken();
+    TokenData             tokenData               = model.getKeyBackupCurrentToken();
-    String                basicStorageCredentials = model.getBasicStorageCredentials();
 
     setSpinning(pinButton);
 
@@ -167,8 +166,7 @@ public final class RegistrationLockFragment extends BaseRegistrationFragment {
                                       model.getFcmToken(),
                                       model.getTextCodeEntered(),
                                       pin,
-                                      basicStorageCredentials,
+                                      tokenData,
-                                      tokenResponse,
 
       new CodeVerificationRequest.VerifyCallback() {
 
@@ -189,19 +187,19 @@ public final class RegistrationLockFragment extends BaseRegistrationFragment {
         }
 
         @Override
-        public void onKbsRegistrationLockPinRequired(long timeRemaining, @NonNull TokenResponse kbsTokenResponse, @NonNull String kbsStorageCredentials) {
+        public void onKbsRegistrationLockPinRequired(long timeRemaining, @NonNull TokenData kbsTokenData, @NonNull String kbsStorageCredentials) {
           throw new AssertionError("Not expected after a pin guess");
         }
 
         @Override
-        public void onIncorrectKbsRegistrationLockPin(@NonNull TokenResponse tokenResponse) {
+        public void onIncorrectKbsRegistrationLockPin(@NonNull TokenData tokenData) {
           cancelSpinning(pinButton);
           pinEntry.getText().clear();
           enableAndFocusPinEntry();
 
-          model.setKeyBackupCurrentToken(tokenResponse);
+          model.setKeyBackupTokenData(tokenData);
 
-          int triesRemaining = tokenResponse.getTries();
+          int triesRemaining = tokenData.getTriesRemaining();
 
           if (triesRemaining == 0) {
             Log.w(TAG, "Account locked. User out of attempts on KBS.");

app/src/main/java/org/thoughtcrime/securesms/registration/service/CodeVerificationRequest.java

@@ -22,6 +22,8 @@ import org.thoughtcrime.securesms.jobs.DirectoryRefreshJob;
 import org.thoughtcrime.securesms.jobs.RotateCertificateJob;
 import org.thoughtcrime.securesms.keyvalue.SignalStore;
 import org.thoughtcrime.securesms.logging.Log;
+import org.thoughtcrime.securesms.pin.PinRestoreRepository;
+import org.thoughtcrime.securesms.pin.PinRestoreRepository.TokenData;
 import org.thoughtcrime.securesms.pin.PinState;
 import org.thoughtcrime.securesms.push.AccountManagerFactory;
 import org.thoughtcrime.securesms.recipients.Recipient;
@@ -65,41 +67,40 @@ public final class CodeVerificationRequest {
   /**
    * Asynchronously verify the account via the code.
    *
-   * @param fcmToken         The FCM token for the device.
+   * @param fcmToken     The FCM token for the device.
-   * @param code             The code that was delivered to the user.
+   * @param code         The code that was delivered to the user.
-   * @param pin              The users registration pin.
+   * @param pin          The users registration pin.
-   * @param callback         Exactly one method on this callback will be called.
+   * @param callback     Exactly one method on this callback will be called.
-   * @param kbsTokenResponse By keeping the token, on failure, a newly returned token will be reused in subsequent pin
+   * @param kbsTokenData By keeping the token, on failure, a newly returned token will be reused in subsequent pin
-   *                         attempts, preventing certain attacks, we can also track the attempts making missing replies easier to spot.
+   *                     attempts, preventing certain attacks, we can also track the attempts making missing replies easier to spot.
    */
   static void verifyAccount(@NonNull Context context,
                             @NonNull Credentials credentials,
                             @Nullable String fcmToken,
                             @NonNull String code,
                             @Nullable String pin,
-                            @Nullable String basicStorageCredentials,
+                            @Nullable TokenData kbsTokenData,
-                            @Nullable TokenResponse kbsTokenResponse,
                             @NonNull VerifyCallback callback)
   {
     new AsyncTask<Void, Void, Result>() {
 
       private volatile LockedException lockedException;
-      private volatile TokenResponse   kbsToken;
+      private volatile TokenData       tokenData;
 
       @Override
       protected Result doInBackground(Void... voids) {
         final boolean pinSupplied = pin != null;
-        final boolean tryKbs      = kbsTokenResponse != null;
+        final boolean tryKbs      = tokenData != null;
 
         try {
-          kbsToken = kbsTokenResponse;
+          this.tokenData = kbsTokenData;
-          verifyAccount(context, credentials, code, pin, kbsTokenResponse, basicStorageCredentials, fcmToken);
+          verifyAccount(context, credentials, code, pin, tokenData, fcmToken);
           return Result.SUCCESS;
         } catch (KeyBackupSystemNoDataException e) {
           Log.w(TAG, "No data found on KBS");
           return Result.KBS_ACCOUNT_LOCKED;
         } catch (KeyBackupSystemWrongPinException e) {
-          kbsToken = e.getTokenResponse();
+          tokenData = TokenData.withResponse(tokenData, e.getTokenResponse());
           return Result.KBS_WRONG_PIN;
         } catch (LockedException e) {
           if (pinSupplied && tryKbs) {
@@ -110,8 +111,8 @@ public final class CodeVerificationRequest {
           lockedException = e;
           if (e.getBasicStorageCredentials() != null) {
             try {
-              kbsToken = getToken(e.getBasicStorageCredentials());
+              tokenData = getToken(e.getBasicStorageCredentials());
-              if (kbsToken == null || kbsToken.getTries() == 0) {
+              if (tokenData == null || tokenData.getTriesRemaining() == 0) {
                 return Result.KBS_ACCOUNT_LOCKED;
               }
             } catch (IOException ex) {
@@ -137,12 +138,12 @@ public final class CodeVerificationRequest {
             callback.onSuccessfulRegistration();
             break;
           case PIN_LOCKED:
-            if (kbsToken != null) {
+            if (tokenData != null) {
               if (lockedException.getBasicStorageCredentials() == null) {
                 throw new AssertionError("KBS Token set, but no storage credentials supplied.");
               }
               Log.w(TAG, "Reg Locked: V2 pin needed for registration");
-              callback.onKbsRegistrationLockPinRequired(lockedException.getTimeRemaining(), kbsToken, lockedException.getBasicStorageCredentials());
+              callback.onKbsRegistrationLockPinRequired(lockedException.getTimeRemaining(), tokenData, lockedException.getBasicStorageCredentials());
             } else {
               Log.w(TAG, "Reg Locked: V1 pin needed for registration");
               callback.onV1RegistrationLockPinRequiredOrIncorrect(lockedException.getTimeRemaining());
@@ -156,7 +157,7 @@ public final class CodeVerificationRequest {
             break;
           case KBS_WRONG_PIN:
             Log.w(TAG, "KBS Pin was wrong");
-            callback.onIncorrectKbsRegistrationLockPin(kbsToken);
+            callback.onIncorrectKbsRegistrationLockPin(tokenData);
             break;
           case KBS_ACCOUNT_LOCKED:
             Log.w(TAG, "KBS Account is locked");
@@ -167,9 +168,9 @@ public final class CodeVerificationRequest {
     }.executeOnExecutor(SignalExecutors.UNBOUNDED);
   }
 
-  private static TokenResponse getToken(@Nullable String basicStorageCredentials) throws IOException {
+  private static TokenData getToken(@Nullable String basicStorageCredentials) throws IOException {
     if (basicStorageCredentials == null) return null;
-    return ApplicationDependencies.getKeyBackupService().getToken(basicStorageCredentials);
+    return new PinRestoreRepository().getTokenSync(basicStorageCredentials);
   }
 
   private static void handleSuccessfulRegistration(@NonNull Context context) {
@@ -185,12 +186,11 @@ public final class CodeVerificationRequest {
                                     @NonNull Credentials credentials,
                                     @NonNull String code,
                                     @Nullable String pin,
-                                    @Nullable TokenResponse kbsTokenResponse,
+                                    @Nullable TokenData kbsTokenData,
-                                    @Nullable String kbsStorageCredentials,
                                     @Nullable String fcmToken)
     throws IOException, KeyBackupSystemWrongPinException, KeyBackupSystemNoDataException
   {
-    boolean    isV2RegistrationLock        = kbsTokenResponse != null;
+    boolean    isV2RegistrationLock        = kbsTokenData != null;
     int        registrationId              = KeyHelper.generateRegistrationId(false);
     boolean    universalUnidentifiedAccess = TextSecurePreferences.isUniversalUnidentifiedAccess(context);
     ProfileKey profileKey                  = findExistingProfileKey(context, credentials.getE164number());
@@ -206,7 +206,7 @@ public final class CodeVerificationRequest {
     SessionUtil.archiveAllSessions(context);
 
     SignalServiceAccountManager accountManager     = AccountManagerFactory.createUnauthenticated(context, credentials.getE164number(), credentials.getPassword());
-    KbsPinData                  kbsData            = isV2RegistrationLock ? PinState.restoreMasterKey(pin, kbsStorageCredentials, kbsTokenResponse) : null;
+    KbsPinData                  kbsData            = isV2RegistrationLock ? PinState.restoreMasterKey(pin, kbsTokenData.getEnclave(), kbsTokenData.getBasicAuth(), kbsTokenData.getTokenResponse()) : null;
     String                      registrationLockV2 = kbsData != null ? kbsData.getMasterKey().deriveRegistrationLock() : null;
     String                      registrationLockV1 = isV2RegistrationLock ? null : pin;
     boolean                     hasFcm             = fcmToken != null;
@@ -292,14 +292,14 @@ public final class CodeVerificationRequest {
     /**
      * The account is locked with a V2 (KBS) pin. Called before any user pin guesses.
      */
-    void onKbsRegistrationLockPinRequired(long timeRemaining, @NonNull TokenResponse kbsTokenResponse, @NonNull String kbsStorageCredentials);
+    void onKbsRegistrationLockPinRequired(long timeRemaining, @NonNull TokenData kbsTokenData, @NonNull String kbsStorageCredentials);
 
     /**
      * The account is locked with a V2 (KBS) pin. Called after a user pin guess.
      * <p>
      * i.e. an attempt has likely been used.
      */
-    void onIncorrectKbsRegistrationLockPin(@NonNull TokenResponse kbsTokenResponse);
+    void onIncorrectKbsRegistrationLockPin(@NonNull TokenData kbsTokenResponse);
 
     /**
      * V2 (KBS) pin is set, but there is no data on KBS.

app/src/main/java/org/thoughtcrime/securesms/registration/service/RegistrationService.java

@@ -5,6 +5,7 @@ import android.app.Activity;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
+import org.thoughtcrime.securesms.pin.PinRestoreRepository;
 import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
 
 import java.io.IOException;
@@ -39,10 +40,9 @@ public final class RegistrationService {
                             @Nullable String fcmToken,
                             @NonNull String code,
                             @Nullable String pin,
-                            @Nullable String basicStorageCredentials,
+                            @Nullable PinRestoreRepository.TokenData tokenData,
-                            @Nullable TokenResponse tokenResponse,
                             @NonNull CodeVerificationRequest.VerifyCallback callback)
   {
-    CodeVerificationRequest.verifyAccount(activity, credentials, fcmToken, code, pin, basicStorageCredentials, tokenResponse, callback);
+    CodeVerificationRequest.verifyAccount(activity, credentials, fcmToken, code, pin, tokenData, callback);
   }
 }

app/src/main/java/org/thoughtcrime/securesms/registration/viewmodel/RegistrationViewModel.java

@@ -9,6 +9,7 @@ import androidx.lifecycle.SavedStateHandle;
 import androidx.lifecycle.ViewModel;
 
 import org.thoughtcrime.securesms.logging.Log;
+import org.thoughtcrime.securesms.pin.PinRestoreRepository.TokenData;
 import org.thoughtcrime.securesms.util.Util;
 import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
 import org.whispersystems.signalservice.internal.util.JsonUtil;
@@ -28,11 +29,10 @@ public final class RegistrationViewModel extends ViewModel {
   private final MutableLiveData<String>                      textCodeEntered;
   private final MutableLiveData<String>                      captchaToken;
   private final MutableLiveData<String>                      fcmToken;
-  private final MutableLiveData<String>                      basicStorageCredentials;
   private final MutableLiveData<Boolean>                     restoreFlowShown;
   private final MutableLiveData<Integer>                     successfulCodeRequestAttempts;
   private final MutableLiveData<LocalCodeRequestRateLimiter> requestLimiter;
-  private final MutableLiveData<String>                      keyBackupCurrentTokenJson;
+  private final MutableLiveData<TokenData>                   kbsTokenData;
   private final MutableLiveData<Long>                        lockedTimeRemaining;
   private final MutableLiveData<Long>                        canCallAtTime;
 
@@ -43,11 +43,10 @@ public final class RegistrationViewModel extends ViewModel {
     textCodeEntered               = savedStateHandle.getLiveData("TEXT_CODE_ENTERED", "");
     captchaToken                  = savedStateHandle.getLiveData("CAPTCHA");
     fcmToken                      = savedStateHandle.getLiveData("FCM_TOKEN");
-    basicStorageCredentials       = savedStateHandle.getLiveData("BASIC_STORAGE_CREDENTIALS");
     restoreFlowShown              = savedStateHandle.getLiveData("RESTORE_FLOW_SHOWN", false);
     successfulCodeRequestAttempts = savedStateHandle.getLiveData("SUCCESSFUL_CODE_REQUEST_ATTEMPTS", 0);
     requestLimiter                = savedStateHandle.getLiveData("REQUEST_RATE_LIMITER", new LocalCodeRequestRateLimiter(60_000));
-    keyBackupCurrentTokenJson     = savedStateHandle.getLiveData("KBS_TOKEN");
+    kbsTokenData                  = savedStateHandle.getLiveData("KBS_TOKEN");
     lockedTimeRemaining           = savedStateHandle.getLiveData("TIME_REMAINING", 0L);
     canCallAtTime                 = savedStateHandle.getLiveData("CAN_CALL_AT_TIME", 0L);
   }
@@ -158,28 +157,12 @@ public final class RegistrationViewModel extends ViewModel {
     requestLimiter.setValue(requestLimiter.getValue());
   }
 
-  public void setStorageCredentials(@Nullable String storageCredentials) {
+  public @Nullable TokenData getKeyBackupCurrentToken() {
-    basicStorageCredentials.setValue(storageCredentials);
+    return kbsTokenData.getValue();
   }
 
-  public @Nullable String getBasicStorageCredentials() {
+  public void setKeyBackupTokenData(TokenData tokenData) {
-    return basicStorageCredentials.getValue();
+    kbsTokenData.setValue(tokenData);
-  }
-
-  public @Nullable TokenResponse getKeyBackupCurrentToken() {
-    String json = keyBackupCurrentTokenJson.getValue();
-    if (json == null) return null;
-    try {
-      return JsonUtil.fromJson(json, TokenResponse.class);
-    } catch (IOException e) {
-      Log.w(TAG, e);
-      return null;
-    }
-  }
-
-  public void setKeyBackupCurrentToken(TokenResponse tokenResponse) {
-    String json = tokenResponse == null ? null : JsonUtil.toJson(tokenResponse);
-    keyBackupCurrentTokenJson.setValue(json);
   }
 
   public LiveData<Long> getLockedTimeRemaining() {

app/src/main/java/org/thoughtcrime/securesms/util/Util.java

@@ -116,6 +116,18 @@ public class Util {
     return join(boxed, delimeter);
   }
 
+  @SafeVarargs
+  public static @NonNull <E> List<E> join(@NonNull List<E>... lists) {
+    int     totalSize = Stream.of(lists).reduce(0, (sum, list) -> sum + list.size());
+    List<E> joined    = new ArrayList<>(totalSize);
+
+    for (List<E> list : lists) {
+      joined.addAll(list);
+    }
+
+    return joined;
+  }
+
   public static String join(List<Long> list, String delimeter) {
     StringBuilder sb = new StringBuilder();
 

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

@@ -17,6 +17,7 @@ import org.whispersystems.signalservice.internal.keybackup.protos.BackupResponse
 import org.whispersystems.signalservice.internal.keybackup.protos.RestoreResponse;
 import org.whispersystems.signalservice.internal.push.PushServiceSocket;
 import org.whispersystems.signalservice.internal.push.RemoteAttestationUtil;
+import org.whispersystems.signalservice.internal.util.Hex;
 import org.whispersystems.signalservice.internal.util.Util;
 
 import java.io.IOException;
@@ -215,6 +216,21 @@ public final class KeyBackupService {
       return new KbsPinData(masterKey, tokenResponse);
     }
 
+    @Override
+    public void removePin()
+        throws IOException, UnauthenticatedResponseException
+    {
+      try {
+        RemoteAttestation remoteAttestation = getAndVerifyRemoteAttestation();
+        KeyBackupRequest  request           = KeyBackupCipher.createKeyDeleteRequest(currentToken, remoteAttestation, serviceId);
+        KeyBackupResponse response          = pushServiceSocket.putKbsData(authorization, request, remoteAttestation.getCookies(), enclaveName);
+
+        KeyBackupCipher.getKeyDeleteResponseStatus(response, remoteAttestation);
+      } catch (InvalidCiphertextException e) {
+        throw new UnauthenticatedResponseException(e);
+      }
+    }
+
     @Override
     public void enableRegistrationLock(MasterKey masterKey) throws IOException {
       pushServiceSocket.setRegistrationLockV2(masterKey.deriveRegistrationLock());
@@ -266,6 +282,9 @@ public final class KeyBackupService {
     /** Creates a PIN. Does nothing to registration lock. */
     KbsPinData setPin(HashedPin hashedPin, MasterKey masterKey) throws IOException, UnauthenticatedResponseException;
 
+    /** Removes the PIN data from KBS. */
+    void removePin() throws IOException, UnauthenticatedResponseException;
+
     /** Enables registration lock. This assumes a PIN is set. */
     void enableRegistrationLock(MasterKey masterKey) throws IOException;
 

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

@@ -2,6 +2,6 @@ package org.whispersystems.signalservice.api;
 
 public final class KeyBackupSystemNoDataException extends Exception {
 
-  KeyBackupSystemNoDataException() {
+  public KeyBackupSystemNoDataException() {
   }
 }