Bug fixes for the new registration flow.

app/src/main/java/org/thoughtcrime/securesms/registration/VerifyAccountRepository.kt

@@ -182,6 +182,12 @@ class VerifyAccountRepository(private val context: Application) {
     }.subscribeOn(Schedulers.io())
   }
 
+  fun getFcmToken(): Single<String> {
+    return Single.fromCallable {
+      return@fromCallable FcmUtil.getToken(context).orElse("")
+    }.subscribeOn(Schedulers.io())
+  }
+
   interface KbsPinDataProducer {
     @Throws(IOException::class, KeyBackupSystemWrongPinException::class, KeyBackupSystemNoDataException::class)
     fun produceKbsPinData(): KbsPinData

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

@@ -6,6 +6,7 @@ import android.widget.ScrollView;
 import android.widget.TextView;
 import android.widget.Toast;
 
+import androidx.activity.OnBackPressedCallback;
 import androidx.annotation.CallSuper;
 import androidx.annotation.LayoutRes;
 import androidx.annotation.NonNull;
@@ -124,6 +125,15 @@ public abstract class BaseEnterSmsCodeFragment<ViewModel extends BaseRegistratio
         new ContactSupportBottomSheetFragment(this::openTroubleshootingSteps, this::sendEmailToSupport).show(getChildFragmentManager(), "support_bottom_sheet");
       }
     });
+
+    requireActivity().getOnBackPressedDispatcher().addCallback(getViewLifecycleOwner(), new OnBackPressedCallback(true) {
+      @Override
+      public void handleOnBackPressed() {
+        viewModel.resetSession();
+        this.remove();
+        requireActivity().getOnBackPressedDispatcher().onBackPressed();
+      }
+    });
   }
 
   protected abstract ViewModel getViewModel();
@@ -324,8 +334,7 @@ public abstract class BaseEnterSmsCodeFragment<ViewModel extends BaseRegistratio
                                     } else if (processor.captchaRequired()) {
                                       navigateToCaptcha();
                                     } else if (processor.rateLimit()) {
-                                      long rateLimit = processor.getRateLimit();
+                                      handleRateLimited();
-                                      Toast.makeText(requireContext(), R.string.RegistrationActivity_rate_limited_to_service, Toast.LENGTH_LONG).show();
                                     } else {
                                       Log.w(TAG, "Unable to request phone code", processor.getError());
                                       Toast.makeText(requireContext(), R.string.RegistrationActivity_unable_to_connect_to_service, Toast.LENGTH_LONG).show();

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

@@ -156,8 +156,7 @@ public final class EnterPhoneNumberFragment extends LoggingFragment implements R
   }
 
   private void handleRegister(@NonNull Context context) {
-    final EditText countryCodeEditText = countryCode.getEditText();
+    if (viewModel.getNumber().getCountryCode() == 0) {
-    if (countryCodeEditText == null || TextUtils.isEmpty(countryCodeEditText.getText())) {
       showErrorDialog(context, getString(R.string.RegistrationActivity_you_must_specify_your_country_code));
       return;
     }
@@ -272,6 +271,7 @@ public final class EnterPhoneNumberFragment extends LoggingFragment implements R
                                   .observeOn(AndroidSchedulers.mainThread())
                                   .subscribe(processor -> {
                                     if (processor.verificationCodeRequestSuccess()) {
+                                      disposables.add(updateFcmTokenValue());
                                       SafeNavigation.safeNavigate(navController, EnterPhoneNumberFragmentDirections.actionEnterVerificationCode());
                                     } else if (processor.captchaRequired()) {
                                       Log.i(TAG, "Unable to request sms code due to captcha required");
@@ -301,6 +301,10 @@ public final class EnterPhoneNumberFragment extends LoggingFragment implements R
     disposables.add(request);
   }
 
+  private Disposable updateFcmTokenValue() {
+    return viewModel.updateFcmTokenValue().subscribe();
+  }
+
   private String formatMillisecondsToString(long milliseconds) {
     long totalSeconds = milliseconds / 1000;
     long HH = totalSeconds / 3600;
@@ -350,7 +354,12 @@ public final class EnterPhoneNumberFragment extends LoggingFragment implements R
                                   .observeOn(AndroidSchedulers.mainThread())
                                   .subscribe(processor -> {
                                     if (processor.hasResult() && processor.canSubmitProofImmediately()) {
-                                      SafeNavigation.safeNavigate(navController, EnterPhoneNumberFragmentDirections.actionEnterVerificationCode());
+                                      try {
+                                        viewModel.restorePhoneNumberStateFromE164(sessionE164);
+                                        SafeNavigation.safeNavigate(navController, EnterPhoneNumberFragmentDirections.actionEnterVerificationCode());
+                                      } catch (NumberParseException numberParseException) {
+                                        viewModel.resetSession();
+                                      }
                                     } else {
                                       viewModel.resetSession();
                                     }

app/src/main/java/org/thoughtcrime/securesms/registration/util/RegistrationNumberInputController.kt

@@ -32,6 +32,7 @@ class RegistrationNumberInputController(
     .map { CountryPrefix(it, PhoneNumberUtil.getInstance().getRegionCodeForCountryCode(it)) }
     .sortedBy { it.digits }
   private val spinnerAdapter: ArrayAdapter<CountryPrefix> = ArrayAdapter<CountryPrefix>(context, R.layout.registration_country_code_dropdown_item, supportedCountryPrefixes)
+  private val countryCodeEntryListener = CountryCodeEntryListener()
 
   private var countryFormatter: AsYouTypeFormatter? = null
   private var isUpdating = true
@@ -41,11 +42,13 @@ class RegistrationNumberInputController(
 
     spinnerView.threshold = 100
     spinnerView.setAdapter(spinnerAdapter)
-    spinnerView.addTextChangedListener(CountryCodeEntryListener())
+    spinnerView.addTextChangedListener(countryCodeEntryListener)
   }
 
   fun prepopulateCountryCode() {
-    spinnerView.setText(supportedCountryPrefixes[0].toString())
+    if (spinnerView.editableText.isBlank()) {
+      spinnerView.setText(supportedCountryPrefixes[0].toString())
+    }
   }
 
   private fun advanceToPhoneNumberInput() {

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

@@ -6,6 +6,10 @@ import androidx.lifecycle.LiveData;
 import androidx.lifecycle.SavedStateHandle;
 import androidx.lifecycle.ViewModel;
 
+import com.google.i18n.phonenumbers.NumberParseException;
+import com.google.i18n.phonenumbers.PhoneNumberUtil;
+import com.google.i18n.phonenumbers.Phonenumber;
+
 import org.thoughtcrime.securesms.keyvalue.SignalStore;
 import org.thoughtcrime.securesms.pin.KbsRepository;
 import org.thoughtcrime.securesms.pin.TokenData;
@@ -102,10 +106,17 @@ public abstract class BaseRegistrationViewModel extends ViewModel {
     return savedState.getLiveData(STATE_NUMBER);
   }
 
+  public void restorePhoneNumberStateFromE164(String e164) throws NumberParseException {
+    Phonenumber.PhoneNumber phoneNumber = PhoneNumberUtil.getInstance().parse(e164, null);
+    onCountrySelected(null, phoneNumber.getCountryCode());
+    setNationalNumber(String.valueOf(phoneNumber.getNationalNumber()));
+  }
+
   public void onCountrySelected(@Nullable String selectedCountryName, int countryCode) {
     setViewState(getNumber().toBuilder()
                             .selectedCountryDisplayName(selectedCountryName)
-                            .countryCode(countryCode).build());
+                            .countryCode(countryCode)
+                            .build());
   }
 
   public void setNationalNumber(String number) {

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

@@ -26,12 +26,12 @@ import org.thoughtcrime.securesms.registration.VerifyAccountRepository;
 import org.thoughtcrime.securesms.registration.VerifyResponse;
 import org.thoughtcrime.securesms.registration.VerifyResponseProcessor;
 import org.thoughtcrime.securesms.registration.VerifyResponseWithRegistrationLockProcessor;
-import org.thoughtcrime.securesms.registration.VerifyResponseWithSuccessfulKbs;
 import org.thoughtcrime.securesms.registration.VerifyResponseWithoutKbs;
 import org.thoughtcrime.securesms.util.FeatureFlags;
 import org.thoughtcrime.securesms.util.Util;
 import org.whispersystems.signalservice.api.KbsPinData;
 import org.whispersystems.signalservice.api.KeyBackupSystemNoDataException;
+import org.whispersystems.signalservice.api.push.exceptions.IncorrectCodeException;
 import org.whispersystems.signalservice.internal.ServiceResponse;
 import org.whispersystems.signalservice.internal.push.RegistrationSessionMetadataResponse;
 
@@ -123,13 +123,17 @@ public final class RegistrationViewModel extends BaseRegistrationViewModel {
                                   .map(RegistrationSessionProcessor.RegistrationSessionProcessorForVerification::new)
                                   .observeOn(AndroidSchedulers.mainThread())
                                   .doOnSuccess(processor -> {
-                                    setCanSmsAtTime(processor.getNextCodeViaSmsAttempt());
+                                    if (processor.hasResult()) {
-                                    setCanCallAtTime(processor.getNextCodeViaCallAttempt());
+                                      setCanSmsAtTime(processor.getNextCodeViaSmsAttempt());
+                                      setCanCallAtTime(processor.getNextCodeViaCallAttempt());
+                                    }
                                   })
                                   .observeOn(Schedulers.io())
                                   .flatMap(processor -> {
                                     if (processor.isAlreadyVerified() || (processor.hasResult() && processor.isVerified())) {
                                       return verifyAccountRepository.registerAccount(sessionId, getRegistrationData(), null, null);
+                                    } else if (processor.getError() == null) {
+                                      return Single.just(ServiceResponse.<VerifyResponse>forApplicationError(new IncorrectCodeException(), 403, null));
                                     } else {
                                       return Single.just(ServiceResponse.<VerifyResponse, RegistrationSessionMetadataResponse>coerceError(processor.getResponse()));
                                     }
@@ -320,6 +324,10 @@ public final class RegistrationViewModel extends BaseRegistrationViewModel {
                                  .observeOn(AndroidSchedulers.mainThread());
   }
 
+  public Single<String> updateFcmTokenValue() {
+    return verifyAccountRepository.getFcmToken().observeOn(AndroidSchedulers.mainThread()).doOnSuccess(this::setFcmToken);
+  }
+
   private void restoreFromStorageService() {
     SignalStore.onboarding().clearAll();
 

libsignal/service/src/main/java/org/whispersystems/signalservice/api/push/exceptions/IncorrectCodeException.kt

@@ -0,0 +1,3 @@
+package org.whispersystems.signalservice.api.push.exceptions
+
+class IncorrectCodeException : NonSuccessfulResponseCodeException(403)