Add fallback KBS and migrate to current version

lib/src/main/java/org/asamk/signal/manager/RegistrationManagerImpl.java

@@ -92,7 +92,15 @@ class RegistrationManagerImpl implements RegistrationManager {
                 serviceEnvironmentConfig.getKeyBackupConfig().getServiceId(),
                 serviceEnvironmentConfig.getKeyBackupConfig().getMrenclave(),
                 10);
-        this.pinHelper = new PinHelper(keyBackupService);
+        final var fallbackKeyBackupServices = serviceEnvironmentConfig.getFallbackKeyBackupConfigs()
+                .stream()
+                .map(config -> accountManager.getKeyBackupService(ServiceConfig.getIasKeyStore(),
+                        config.getEnclaveName(),
+                        config.getServiceId(),
+                        config.getMrenclave(),
+                        10))
+                .toList();
+        this.pinHelper = new PinHelper(keyBackupService, fallbackKeyBackupServices);
     }
 
     @Override

lib/src/main/java/org/asamk/signal/manager/SignalDependencies.java

@@ -22,6 +22,7 @@ import org.whispersystems.signalservice.api.util.UptimeSleepTimer;
 import org.whispersystems.signalservice.api.websocket.WebSocketFactory;
 import org.whispersystems.signalservice.internal.websocket.WebSocketConnection;
 
+import java.util.Collection;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.function.Supplier;
@@ -183,6 +184,17 @@ public class SignalDependencies {
                         10));
     }
 
+    public Collection<KeyBackupService> getFallbackKeyBackupServices() {
+        return serviceEnvironmentConfig.getFallbackKeyBackupConfigs()
+                .stream()
+                .map(config -> getAccountManager().getKeyBackupService(ServiceConfig.getIasKeyStore(),
+                        config.getEnclaveName(),
+                        config.getServiceId(),
+                        config.getMrenclave(),
+                        10))
+                .toList();
+    }
+
     public ProfileService getProfileService() {
         return getOrCreate(() -> profileService,
                 () -> profileService = new ProfileService(getClientZkProfileOperations(),

lib/src/main/java/org/asamk/signal/manager/config/LiveConfig.java

@@ -15,6 +15,7 @@ import org.whispersystems.signalservice.internal.configuration.SignalServiceUrl;
 import org.whispersystems.signalservice.internal.configuration.SignalStorageUrl;
 
 import java.util.Base64;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -33,6 +34,10 @@ class LiveConfig {
     private final static byte[] KEY_BACKUP_SERVICE_ID = Hex.decode(
             "3a485adb56e2058ef7737764c738c4069dd62bc457637eafb6bbce1ce29ddb89");
     private final static String KEY_BACKUP_MRENCLAVE = "45627094b2ea4a66f4cf0b182858a8dcf4b8479122c3820fe7fd0551a6d4cf5c";
+    private final static String FALLBACK_KEY_BACKUP_ENCLAVE_NAME = "0cedba03535b41b67729ce9924185f831d7767928a1d1689acb689bc079c375f";
+    private final static byte[] FALLBACK_KEY_BACKUP_SERVICE_ID = Hex.decode(
+            "187d2739d22be65e74b65f0055e74d31310e4267e5fac2b1246cc8beba81af39");
+    private final static String FALLBACK_KEY_BACKUP_MRENCLAVE = "ee19f1965b1eefa3dc4204eb70c04f397755f771b8c1909d080c04dad2a6a9ba";
 
     private final static String URL = "https://chat.signal.org";
     private final static String CDN_URL = "https://cdn.signal.org";
@@ -80,6 +85,12 @@ class LiveConfig {
         return new KeyBackupConfig(KEY_BACKUP_ENCLAVE_NAME, KEY_BACKUP_SERVICE_ID, KEY_BACKUP_MRENCLAVE);
     }
 
+    static Collection<KeyBackupConfig> createFallbackKeyBackupConfigs() {
+        return List.of(new KeyBackupConfig(FALLBACK_KEY_BACKUP_ENCLAVE_NAME,
+                FALLBACK_KEY_BACKUP_SERVICE_ID,
+                FALLBACK_KEY_BACKUP_MRENCLAVE));
+    }
+
     static String getCdsMrenclave() {
         return CDS_MRENCLAVE;
     }

lib/src/main/java/org/asamk/signal/manager/config/ServiceConfig.java

@@ -79,11 +79,13 @@ public class ServiceConfig {
                     LiveConfig.createDefaultServiceConfiguration(interceptors),
                     LiveConfig.getUnidentifiedSenderTrustRoot(),
                     LiveConfig.createKeyBackupConfig(),
+                    LiveConfig.createFallbackKeyBackupConfigs(),
                     LiveConfig.getCdsMrenclave());
             case STAGING -> new ServiceEnvironmentConfig(serviceEnvironment,
                     StagingConfig.createDefaultServiceConfiguration(interceptors),
                     StagingConfig.getUnidentifiedSenderTrustRoot(),
                     StagingConfig.createKeyBackupConfig(),
+                    StagingConfig.createFallbackKeyBackupConfigs(),
                     StagingConfig.getCdsMrenclave());
         };
     }

lib/src/main/java/org/asamk/signal/manager/config/ServiceEnvironmentConfig.java

@@ -3,6 +3,8 @@ package org.asamk.signal.manager.config;
 import org.signal.libsignal.protocol.ecc.ECPublicKey;
 import org.whispersystems.signalservice.internal.configuration.SignalServiceConfiguration;
 
+import java.util.Collection;
+
 public class ServiceEnvironmentConfig {
 
     private final ServiceEnvironment type;
@@ -11,6 +13,7 @@ public class ServiceEnvironmentConfig {
     private final ECPublicKey unidentifiedSenderTrustRoot;
 
     private final KeyBackupConfig keyBackupConfig;
+    private final Collection<KeyBackupConfig> fallbackKeyBackupConfigs;
 
     private final String cdsMrenclave;
 
@@ -19,12 +22,14 @@ public class ServiceEnvironmentConfig {
             final SignalServiceConfiguration signalServiceConfiguration,
             final ECPublicKey unidentifiedSenderTrustRoot,
             final KeyBackupConfig keyBackupConfig,
+            final Collection<KeyBackupConfig> fallbackKeyBackupConfigs,
             final String cdsMrenclave
     ) {
         this.type = type;
         this.signalServiceConfiguration = signalServiceConfiguration;
         this.unidentifiedSenderTrustRoot = unidentifiedSenderTrustRoot;
         this.keyBackupConfig = keyBackupConfig;
+        this.fallbackKeyBackupConfigs = fallbackKeyBackupConfigs;
         this.cdsMrenclave = cdsMrenclave;
     }
 
@@ -44,6 +49,10 @@ public class ServiceEnvironmentConfig {
         return keyBackupConfig;
     }
 
+    public Collection<KeyBackupConfig> getFallbackKeyBackupConfigs() {
+        return fallbackKeyBackupConfigs;
+    }
+
     public String getCdsMrenclave() {
         return cdsMrenclave;
     }

lib/src/main/java/org/asamk/signal/manager/config/StagingConfig.java

@@ -15,6 +15,7 @@ import org.whispersystems.signalservice.internal.configuration.SignalServiceUrl;
 import org.whispersystems.signalservice.internal.configuration.SignalStorageUrl;
 
 import java.util.Base64;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -33,6 +34,10 @@ class StagingConfig {
     private final static byte[] KEY_BACKUP_SERVICE_ID = Hex.decode(
             "9dbc6855c198e04f21b5cc35df839fdcd51b53658454dfa3f817afefaffc95ef");
     private final static String KEY_BACKUP_MRENCLAVE = "45627094b2ea4a66f4cf0b182858a8dcf4b8479122c3820fe7fd0551a6d4cf5c";
+    private final static String FALLBACK_KEY_BACKUP_ENCLAVE_NAME = "dd6f66d397d9e8cf6ec6db238e59a7be078dd50e9715427b9c89b409ffe53f99";
+    private final static byte[] FALLBACK_KEY_BACKUP_SERVICE_ID = Hex.decode(
+            "4200003414528c151e2dccafbc87aa6d3d66a5eb8f8c05979a6e97cb33cd493a");
+    private final static String FALLBACK_KEY_BACKUP_MRENCLAVE = "ee19f1965b1eefa3dc4204eb70c04f397755f771b8c1909d080c04dad2a6a9ba";
 
     private final static String URL = "https://chat.staging.signal.org";
     private final static String CDN_URL = "https://cdn-staging.signal.org";
@@ -80,6 +85,12 @@ class StagingConfig {
         return new KeyBackupConfig(KEY_BACKUP_ENCLAVE_NAME, KEY_BACKUP_SERVICE_ID, KEY_BACKUP_MRENCLAVE);
     }
 
+    static Collection<KeyBackupConfig> createFallbackKeyBackupConfigs() {
+        return List.of(new KeyBackupConfig(FALLBACK_KEY_BACKUP_ENCLAVE_NAME,
+                FALLBACK_KEY_BACKUP_SERVICE_ID,
+                FALLBACK_KEY_BACKUP_MRENCLAVE));
+    }
+
     static String getCdsMrenclave() {
         return CDS_MRENCLAVE;
     }

lib/src/main/java/org/asamk/signal/manager/helper/AccountHelper.java

@@ -60,6 +60,7 @@ public class AccountHelper {
             }
         }
         try {
+            updateAccountAttributes();
             context.getPreKeyHelper().refreshPreKeysIfNecessary();
             if (account.getAci() == null || account.getPni() == null) {
                 checkWhoAmiI();
@@ -67,7 +68,11 @@ public class AccountHelper {
             if (!account.isPrimaryDevice() && account.getPniIdentityKeyPair() == null) {
                 context.getSyncHelper().requestSyncPniIdentity();
             }
-            updateAccountAttributes();
+            if (account.getPreviousStorageVersion() < 4
+                    && account.isPrimaryDevice()
+                    && account.getRegistrationLockPin() != null) {
+                migrateRegistrationPin();
+            }
         } catch (AuthorizationFailedException e) {
             account.setRegistered(false);
             throw e;
@@ -171,6 +176,12 @@ public class AccountHelper {
         account.setMultiDevice(devices.size() > 1);
     }
 
+    public void migrateRegistrationPin() throws IOException {
+        var masterKey = account.getOrCreatePinMasterKey();
+
+        context.getPinHelper().migrateRegistrationLockPin(account.getRegistrationLockPin(), masterKey);
+    }
+
     public void setRegistrationPin(String pin) throws IOException {
         var masterKey = account.getOrCreatePinMasterKey();
 

lib/src/main/java/org/asamk/signal/manager/helper/Context.java

@@ -114,7 +114,9 @@ public class Context {
     }
 
     PinHelper getPinHelper() {
-        return getOrCreate(() -> pinHelper, () -> pinHelper = new PinHelper(dependencies.getKeyBackupService()));
+        return getOrCreate(() -> pinHelper,
+                () -> pinHelper = new PinHelper(dependencies.getKeyBackupService(),
+                        dependencies.getFallbackKeyBackupServices()));
     }
 
     public PreKeyHelper getPreKeyHelper() {

lib/src/main/java/org/asamk/signal/manager/helper/PinHelper.java

@@ -1,8 +1,11 @@
 package org.asamk.signal.manager.helper;
 
 import org.asamk.signal.manager.api.IncorrectPinException;
+import org.asamk.signal.manager.api.Pair;
 import org.asamk.signal.manager.util.PinHashing;
 import org.signal.libsignal.protocol.InvalidKeyException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.signalservice.api.KbsPinData;
 import org.whispersystems.signalservice.api.KeyBackupService;
 import org.whispersystems.signalservice.api.KeyBackupServicePinException;
@@ -13,13 +16,21 @@ import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse
 import org.whispersystems.signalservice.internal.push.LockedException;
 
 import java.io.IOException;
+import java.util.Collection;
+import java.util.stream.Stream;
 
 public class PinHelper {
 
-    private final KeyBackupService keyBackupService;
+    private final static Logger logger = LoggerFactory.getLogger(PinHelper.class);
 
-    public PinHelper(final KeyBackupService keyBackupService) {
+    private final KeyBackupService keyBackupService;
+    private final Collection<KeyBackupService> fallbackKeyBackupServices;
+
+    public PinHelper(
+            final KeyBackupService keyBackupService, final Collection<KeyBackupService> fallbackKeyBackupServices
+    ) {
         this.keyBackupService = keyBackupService;
+        this.fallbackKeyBackupServices = fallbackKeyBackupServices;
     }
 
     public void setRegistrationLockPin(
@@ -36,6 +47,19 @@ public class PinHelper {
         pinChangeSession.enableRegistrationLock(masterKey);
     }
 
+    public void migrateRegistrationLockPin(String pin, MasterKey masterKey) throws IOException {
+        setRegistrationLockPin(pin, masterKey);
+
+        for (final var keyBackupService : fallbackKeyBackupServices) {
+            try {
+                final var pinChangeSession = keyBackupService.newPinChangeSession();
+                pinChangeSession.removePin();
+            } catch (Exception e) {
+                logger.warn("Failed to remove PIN from fallback KBS: {}", e.getMessage());
+            }
+        }
+    }
+
     public void removeRegistrationLockPin() throws IOException {
         final var pinChangeSession = keyBackupService.newPinChangeSession();
         pinChangeSession.disableRegistrationLock();
@@ -66,20 +90,34 @@ public class PinHelper {
     private KbsPinData getRegistrationLockData(
             String pin, String basicStorageCredentials
     ) throws IOException, KeyBackupSystemNoDataException, KeyBackupServicePinException {
-        var tokenResponse = keyBackupService.getToken(basicStorageCredentials);
+        var tokenResponsePair = getTokenResponse(basicStorageCredentials);
-        if (tokenResponse == null || tokenResponse.getTries() == 0) {
+        final var tokenResponse = tokenResponsePair.first();
-            throw new IOException("KBS Account locked, maximum pin attempts reached.");
+        final var keyBackupService = tokenResponsePair.second();
-        }
 
-        var registrationLockData = restoreMasterKey(pin, basicStorageCredentials, tokenResponse);
+        var registrationLockData = restoreMasterKey(pin, basicStorageCredentials, tokenResponse, keyBackupService);
         if (registrationLockData == null) {
             throw new AssertionError("Failed to restore master key");
         }
         return registrationLockData;
     }
 
+    private Pair<TokenResponse, KeyBackupService> getTokenResponse(String basicStorageCredentials) throws IOException {
+        final var keyBackupServices = Stream.concat(Stream.of(keyBackupService), fallbackKeyBackupServices.stream())
+                .toList();
+        for (final var keyBackupService : keyBackupServices) {
+            var tokenResponse = keyBackupService.getToken(basicStorageCredentials);
+            if (tokenResponse != null && tokenResponse.getTries() > 0) {
+                return new Pair<>(tokenResponse, keyBackupService);
+            }
+        }
+        throw new IOException("KBS Account locked, maximum pin attempts reached.");
+    }
+
     private KbsPinData restoreMasterKey(
-            String pin, String basicStorageCredentials, TokenResponse tokenResponse
+            String pin,
+            String basicStorageCredentials,
+            TokenResponse tokenResponse,
+            final KeyBackupService keyBackupService
     ) throws IOException, KeyBackupSystemNoDataException, KeyBackupServicePinException {
         if (pin == null) return null;
 

lib/src/main/java/org/asamk/signal/manager/storage/SignalAccount.java

@@ -93,7 +93,7 @@ public class SignalAccount implements Closeable {
     private final static Logger logger = LoggerFactory.getLogger(SignalAccount.class);
 
     private static final int MINIMUM_STORAGE_VERSION = 1;
-    private static final int CURRENT_STORAGE_VERSION = 3;
+    private static final int CURRENT_STORAGE_VERSION = 4;
 
     private final Object LOCK = new Object();
 
@@ -997,6 +997,10 @@ public class SignalAccount implements Closeable {
         save();
     }
 
+    public int getPreviousStorageVersion() {
+        return previousStorageVersion;
+    }
+
     public SignalServiceDataStore getSignalServiceDataStore() {
         return new SignalServiceDataStore() {
             @Override
@@ -1277,6 +1281,10 @@ public class SignalAccount implements Closeable {
         save();
     }
 
+    public String getRegistrationLockPin() {
+        return registrationLockPin;
+    }
+
     public String getRegistrationLock() {
         final var masterKey = getPinBackedMasterKey();
         if (masterKey == null) {