DI Setup

Bug: 301578742
Topic: cl/584012825
Test: m ManagedProvisioning
Change-Id: If44f8403af79681a47d09cbdc66b360e1f033d74

17 files changed, 509 insertions, 43 deletions

diff --git a/Android.bp b/Android.bp
index 976704098..38db1433e 100644
--- a/Android.bp
+++ b/Android.bp
@@ -26,6 +26,19 @@ java_library {
     },
 }
 
+aconfig_declarations {
+    name: "aconfig_managed_provisioning_flags",
+    package: "com.android.managedprovisioning.flags",
+    srcs: [
+      "aconfig/*.aconfig"
+    ],
+}
+
+java_aconfig_library {
+    name: "aconfig_managed_provisioning_flags_lib",
+    aconfig_declarations: "aconfig_managed_provisioning_flags",
+}
+
 android_library {
     name: "ManagedProvisioningLib",
     manifest: "AndroidManifest.xml",
@@ -36,14 +49,20 @@ android_library {
 
     static_libs: [
         "androidx.legacy_legacy-support-v4",
-        "setupcompat",
+        "androidx.annotation_annotation",
         "setupdesign",
         "managedprovisioning_protoslite",
         "androidx.webkit_webkit-nodeps",
         "setupdesign-lottie-loading-layout",
         "lottie",
         "guava",
-        "aconfig_managed_provisioning_flags_lib"
+        "aconfig_managed_provisioning_flags_lib",
+        "hilt_android",
+        "android_onboarding.contracts.provisioning",
+        "android_onboarding.contracts.annotations",
+        "android_onboarding.contracts.setupwizard",
+        "android_onboarding.flags_hilt",
+        "jsr330",
     ],
 
     plugins: ["auto_value_plugin"],
@@ -83,16 +102,3 @@ android_app {
         include_filter: ["com.android.managedprovisioning.*"],
     },
 }
-
-aconfig_declarations {
-    name: "aconfig_managed_provisioning_flags",
-    package: "com.android.managedprovisioning.flags",
-    srcs: [
-      "aconfig/*.aconfig"
-    ],
-}
-
-java_aconfig_library {
-    name: "aconfig_managed_provisioning_flags_lib",
-    aconfig_declarations: "aconfig_managed_provisioning_flags",
-}
diff --git a/aconfig/root.aconfig b/aconfig/root.aconfig
index 858dd5a7e..defd6594d 100644
--- a/aconfig/root.aconfig
+++ b/aconfig/root.aconfig
@@ -1,7 +1,7 @@
 package: "com.android.managedprovisioning.flags"
 
 flag {
-    name: "cosmic_ray"
+    name: "is_cosmic_ray_enabled"
     namespace: "enterprise"
     description: "Enables Cosmic Ray features"
     bug: "288413994"
diff --git a/proguard.flags b/proguard.flags
index 2f45d5450..af83ce119 100644
--- a/proguard.flags
+++ b/proguard.flags
@@ -1,6 +1,6 @@
 # Ensure exception attributes are kept to allow mocking of methods with
 # checked exceptions.
--keepattributes Exceptions
+-keepattributes Exception
 
 # It is not an error that these classes are referenced in
 # AndriodManifest.xml, but not present in Android master. They are
diff --git a/src/com/android/managedprovisioning/ManagedProvisioningBaseApplication.java b/src/com/android/managedprovisioning/ManagedProvisioningBaseApplication.java
index 54aae9f8c..e20e21050 100644
--- a/src/com/android/managedprovisioning/ManagedProvisioningBaseApplication.java
+++ b/src/com/android/managedprovisioning/ManagedProvisioningBaseApplication.java
@@ -23,6 +23,10 @@ import android.view.WindowManager;
 
 import com.android.managedprovisioning.preprovisioning.EncryptionController;
 
+import dagger.hilt.android.HiltAndroidApp;
+
+import javax.inject.Inject;
+
 /**
  * A base {@link Application} that is meant to be extended.
  *
@@ -32,9 +36,11 @@ import com.android.managedprovisioning.preprovisioning.EncryptionController;
  *
  * <p>By default, the existing {@code ManagedProvisioning} {@link Activity} classes are used.
  */
-public abstract class ManagedProvisioningBaseApplication extends Application {
-    private final ScreenManager mScreenManager =
-            new ScreenManager(ScreenManager.DEFAULT_SCREEN_TO_ACTIVITY_MAP);
+@HiltAndroidApp(Application.class)
+public abstract class ManagedProvisioningBaseApplication extends
+        Hilt_ManagedProvisioningBaseApplication {
+    @Inject
+    protected ScreenManager mScreenManager;
     private EncryptionController mEncryptionController;
     private boolean mKeepScreenOn;
 
@@ -70,8 +76,8 @@ public abstract class ManagedProvisioningBaseApplication extends Application {
      * <p>If no screens were set via {@link #setOverrideActivity(ManagedProvisioningScreens,
      * Class)}, the base ManagedProvisioning {@link Activity} implementation will be returned.
      */
-    public final Class<? extends Activity>
-            getActivityClassForScreen(ManagedProvisioningScreens screen) {
+    public final Class<? extends Activity> getActivityClassForScreen(
+            ManagedProvisioningScreens screen) {
         return mScreenManager.getActivityClassForScreen(screen);
     }
 
diff --git a/src/com/android/managedprovisioning/Module.kt b/src/com/android/managedprovisioning/Module.kt
new file mode 100644
index 000000000..86625063f
--- /dev/null
+++ b/src/com/android/managedprovisioning/Module.kt
@@ -0,0 +1,15 @@
+package com.android.managedprovisioning
+
+import dagger.Module
+import dagger.Provides
+import dagger.hilt.InstallIn
+import dagger.hilt.android.components.ActivityComponent
+import dagger.hilt.components.SingletonComponent
+
+@Module
+@InstallIn(ActivityComponent::class, SingletonComponent::class)
+class Module {
+    @Provides
+    fun defaultScreenManager(): ScreenManager =
+        ScreenManager(ScreenManager.DEFAULT_SCREEN_TO_ACTIVITY_MAP)
+}
\ No newline at end of file
diff --git a/src/com/android/managedprovisioning/annotations/LegacyApi.kt b/src/com/android/managedprovisioning/annotations/LegacyApi.kt
new file mode 100644
index 000000000..671ad2360
--- /dev/null
+++ b/src/com/android/managedprovisioning/annotations/LegacyApi.kt
@@ -0,0 +1,9 @@
+package com.android.managedprovisioning.annotations
+
+@RequiresOptIn(
+    message =
+    "Marks API meant only to facilitate legacy code integration. " +
+            "Should be avoided whenever possible.",
+    level = RequiresOptIn.Level.ERROR,
+)
+annotation class LegacyApi
\ No newline at end of file
diff --git a/src/com/android/managedprovisioning/common/CommonModule.kt b/src/com/android/managedprovisioning/common/CommonModule.kt
new file mode 100644
index 000000000..76a30bfc2
--- /dev/null
+++ b/src/com/android/managedprovisioning/common/CommonModule.kt
@@ -0,0 +1,13 @@
+package com.android.managedprovisioning.common
+
+import dagger.Binds
+import dagger.Module
+import dagger.hilt.InstallIn
+import dagger.hilt.android.components.ActivityComponent
+
+@Module
+@InstallIn(ActivityComponent::class)
+interface CommonModule {
+    @Binds
+    fun bind(impl: DefaultFlags): Flags
+}
diff --git a/src/com/android/managedprovisioning/common/Flags.kt b/src/com/android/managedprovisioning/common/Flags.kt
new file mode 100644
index 000000000..270f5065d
--- /dev/null
+++ b/src/com/android/managedprovisioning/common/Flags.kt
@@ -0,0 +1,33 @@
+package com.android.managedprovisioning.common
+
+import com.android.managedprovisioning.annotations.LegacyApi
+import com.android.onboarding.flags.DefaultOnboardingFlagsProvider
+import com.android.onboarding.flags.OnboardingFlagsProvider
+import javax.inject.Inject
+import com.android.managedprovisioning.flags.FeatureFlags as AconfigFlags
+import com.android.managedprovisioning.flags.FeatureFlagsImpl as DefaultAconfigFlags
+
+
+interface Flags : AconfigFlags {
+    @Deprecated(
+        "Static version is reserved for edge-cases where DI could not be " +
+                "easily achieved and as such should be avoided. " +
+                "Consider using injected version whenever possible."
+    )
+    @LegacyApi
+    companion object : Flags by DefaultFlags(DefaultOnboardingFlagsProvider())
+}
+
+class DefaultFlags(
+    private val onboardingFlags: OnboardingFlagsProvider,
+    private val aconfigFlags: AconfigFlags,
+) : Flags, AconfigFlags by aconfigFlags {
+    @Inject
+    constructor(
+        onboardingFlags: OnboardingFlagsProvider,
+    ) : this(onboardingFlags, DefaultAconfigFlags())
+
+    override fun isCosmicRayEnabled(): Boolean =
+        onboardingFlags.isDebug ||
+                (onboardingFlags.isContractEnabled && aconfigFlags.isCosmicRayEnabled)
+}
diff --git a/src/com/android/managedprovisioning/common/SetupLayoutActivity.java b/src/com/android/managedprovisioning/common/SetupLayoutActivity.java
index ea8e9320a..dd228fbcc 100644
--- a/src/com/android/managedprovisioning/common/SetupLayoutActivity.java
+++ b/src/com/android/managedprovisioning/common/SetupLayoutActivity.java
@@ -45,10 +45,13 @@ import com.android.managedprovisioning.common.ThemeHelper.DefaultNightModeChecke
 import com.android.managedprovisioning.common.ThemeHelper.DefaultSetupWizardBridge;
 import com.android.managedprovisioning.preprovisioning.EncryptionController;
 
+import dagger.hilt.android.AndroidEntryPoint;
+
 /**
  * Base class for setting up the layout.
  */
-public abstract class SetupLayoutActivity extends AppCompatActivity {
+@AndroidEntryPoint(AppCompatActivity.class)
+public abstract class SetupLayoutActivity extends Hilt_SetupLayoutActivity {
     protected final Utils mUtils;
     protected final SettingsFacade mSettingsFacade;
     private final ThemeHelper mThemeHelper;
@@ -176,9 +179,11 @@ public abstract class SetupLayoutActivity extends AppCompatActivity {
 
     /**
      * Constructs and shows a {@link DialogFragment} unless it is already displayed.
+     *
      * @param dialogBuilder Lightweight builder, that it is inexpensive to discard it if dialog
-     * already shown.
-     * @param tag The tag for this dialog, as per {@link FragmentTransaction#add(Fragment, String)}.
+     *                      already shown.
+     * @param tag           The tag for this dialog, as per
+     *                      {@link FragmentTransaction#add(Fragment, String)}.
      */
     protected void showDialog(DialogBuilder dialogBuilder, String tag) {
         FragmentManager fragmentManager = getFragmentManager();
@@ -189,6 +194,7 @@ public abstract class SetupLayoutActivity extends AppCompatActivity {
 
     /**
      * Checks whether the {@link DialogFragment} associated with the given tag is currently showing.
+     *
      * @param tag The tag for this dialog.
      */
     protected boolean isDialogAdded(String tag) {
diff --git a/src/com/android/managedprovisioning/contracts/Checksum.kt b/src/com/android/managedprovisioning/contracts/Checksum.kt
new file mode 100644
index 000000000..5ef4927f7
--- /dev/null
+++ b/src/com/android/managedprovisioning/contracts/Checksum.kt
@@ -0,0 +1,17 @@
+package com.android.managedprovisioning.contracts
+
+sealed interface Checksum {
+    val bytes: ByteArray
+
+    /**
+     * SHA-256 hash of the .apk file
+     */
+    @JvmInline
+    value class PackageChecksum(override val bytes: ByteArray) : Checksum
+
+    /**
+     * SHA-256 hash of the signature in the .apk file
+     */
+    @JvmInline
+    value class SignatureChecksum(override val bytes: ByteArray) : Checksum
+}
\ No newline at end of file
diff --git a/src/com/android/managedprovisioning/contracts/DownloadRoleHolderContract.kt b/src/com/android/managedprovisioning/contracts/DownloadRoleHolderContract.kt
new file mode 100644
index 000000000..53e26de0f
--- /dev/null
+++ b/src/com/android/managedprovisioning/contracts/DownloadRoleHolderContract.kt
@@ -0,0 +1,100 @@
+package com.android.managedprovisioning.contracts
+
+import android.content.Context
+import android.content.Intent
+import com.android.managedprovisioning.ManagedProvisioningScreens.DOWNLOAD_ROLE_HOLDER
+import com.android.managedprovisioning.ScreenManager
+import com.android.managedprovisioning.annotations.LegacyApi
+import com.android.managedprovisioning.contracts.Checksum.SignatureChecksum
+import com.android.managedprovisioning.model.PackageDownloadInfo
+import com.android.managedprovisioning.model.ProvisioningParams
+import com.android.onboarding.common.MANAGED_PROVISIONING
+import com.android.onboarding.contracts.IntentScope
+import com.android.onboarding.contracts.ScopedIntentSerializer
+import com.android.onboarding.contracts.VoidOnboardingActivityApiContract
+import com.android.onboarding.contracts.annotations.OnboardingNode
+import com.android.onboarding.contracts.provisioning.EXTRAS
+import com.android.onboarding.contracts.require
+import com.android.onboarding.contracts.setupwizard.SuwArguments
+import com.android.onboarding.contracts.setupwizard.SuwArgumentsSerializer
+import com.android.onboarding.contracts.setupwizard.WithSuwArguments
+import javax.inject.Inject
+
+/**
+ * @property location Url where the package (.apk) can be downloaded from. {@code null} if there is no download location specified.
+ * @property cookieHeader Cookie header for http request.
+ * @property checksum SHA-256 hash of the signature in the .apk file
+ * @property minVersion Minimum supported version code of the downloaded package.
+ */
+data class DownloadRoleHolderArguments(
+    override val suwArguments: SuwArguments,
+    override val provisioningParams: ProvisioningParams,
+    val location: String,
+    val checksum: SignatureChecksum,
+    val cookieHeader: String? = null,
+    val minVersion: Int = Int.MAX_VALUE,
+) : WithProvisioningParams, WithSuwArguments {
+    @LegacyApi
+    constructor(
+        suwArguments: SuwArguments,
+        provisioningParams: ProvisioningParams,
+    ) : this(
+        suwArguments = suwArguments,
+        provisioningParams = provisioningParams,
+        location = provisioningParams.roleHolderDownloadInfo?.location
+                ?.require("Download location must not be empty.", String::isNotBlank)
+            ?: error("Missing download location."),
+        checksum = provisioningParams.roleHolderDownloadInfo.signatureChecksum
+                .require("Checksum is missing", ByteArray::isNotEmpty)
+                .let(Checksum::SignatureChecksum),
+        cookieHeader = provisioningParams.roleHolderDownloadInfo.cookieHeader,
+        minVersion = provisioningParams.roleHolderDownloadInfo.minVersion
+    )
+}
+
+@OnboardingNode(
+    component = MANAGED_PROVISIONING,
+    name = "DownloadRoleHolder",
+    hasUi = OnboardingNode.HasUi.NO)
+class DownloadRoleHolderContract
+@Inject
+constructor(
+    private val screenManager: ScreenManager,
+    val suwArgumentsSerializer: SuwArgumentsSerializer
+) : VoidOnboardingActivityApiContract<DownloadRoleHolderArguments>(),
+    ScopedIntentSerializer<DownloadRoleHolderArguments> {
+
+    override fun performCreateIntent(context: Context, arg: DownloadRoleHolderArguments): Intent =
+        Intent(context, screenManager.getActivityClassForScreen(DOWNLOAD_ROLE_HOLDER))
+                .also { write(it, arg) }
+
+    override fun performExtractArgument(intent: Intent): DownloadRoleHolderArguments = read(intent)
+    override fun IntentScope.write(value: DownloadRoleHolderArguments) {
+        write(suwArgumentsSerializer, value.suwArguments)
+        this[EXTRAS.EXTRA_PROVISIONING_PARAMS] =
+            value.provisioningParams.toBuilder()
+                    .setRoleHolderDownloadInfo(
+                        PackageDownloadInfo.Builder()
+                                .setLocation(value.location)
+                                .setCookieHeader(value.cookieHeader)
+                                .setMinVersion(value.minVersion)
+                                .setSignatureChecksum(value.checksum.bytes)
+                                .build())
+                    .build()
+    }
+
+    override fun IntentScope.read(): DownloadRoleHolderArguments {
+        val provisioningParams =
+            parcelableExtra<ProvisioningParams>(EXTRAS.EXTRA_PROVISIONING_PARAMS)
+        val params = provisioningParams.roleHolderDownloadInfo
+            ?: error("Missing role holder extras")
+        return DownloadRoleHolderArguments(
+            suwArguments = read(suwArgumentsSerializer),
+            provisioningParams = provisioningParams,
+            location = params.location,
+            checksum = SignatureChecksum(params.signatureChecksum),
+            cookieHeader = params.cookieHeader,
+            minVersion = params.minVersion,
+        )
+    }
+}
\ No newline at end of file
diff --git a/src/com/android/managedprovisioning/contracts/ProvisioningArgumentsSerializer.kt b/src/com/android/managedprovisioning/contracts/ProvisioningArgumentsSerializer.kt
new file mode 100644
index 000000000..4bfdb928a
--- /dev/null
+++ b/src/com/android/managedprovisioning/contracts/ProvisioningArgumentsSerializer.kt
@@ -0,0 +1,84 @@
+package com.android.managedprovisioning.contracts
+
+import android.content.ComponentName
+import com.android.managedprovisioning.annotations.LegacyApi
+import com.android.managedprovisioning.model.PackageDownloadInfo
+import com.android.managedprovisioning.model.ProvisioningParams
+import com.android.managedprovisioning.model.WifiInfo
+import com.android.onboarding.contracts.IntentScope
+import com.android.onboarding.contracts.ScopedIntentSerializer
+import com.android.onboarding.contracts.provisioning.EXTRAS
+import javax.inject.Inject
+
+interface WithProvisioningParams : WithOptionalProvisioningParams {
+    override val provisioningParams: ProvisioningParams
+}
+
+interface WithOptionalProvisioningParams {
+    /** Only needed to preserve data parity when passing provisioning params bundle along */
+    val provisioningParams: ProvisioningParams?
+}
+
+enum class FlowType {
+    Unspecified,
+    Legacy,
+    AdminIntegrated
+}
+
+data class BaseProvisioningArguments(
+    override val provisioningParams: ProvisioningParams,
+    val flowType: FlowType,
+    val deviceAdminDownloadInfo: PackageDownloadInfo?,
+    val deviceAdminComponentName: ComponentName?,
+    val wifiInfo: WifiInfo?,
+    val useMobileData: Boolean,
+    val isNfc: Boolean,
+    val isQr: Boolean,
+) : WithProvisioningParams {
+    @LegacyApi
+    constructor(
+        provisioningParams: ProvisioningParams
+    ) : this(
+        provisioningParams = provisioningParams,
+        flowType = FlowType.values()[provisioningParams.flowType],
+        deviceAdminDownloadInfo = provisioningParams.deviceAdminDownloadInfo,
+        deviceAdminComponentName = provisioningParams.deviceAdminComponentName,
+        wifiInfo = provisioningParams.wifiInfo,
+        useMobileData = provisioningParams.useMobileData,
+        isNfc = provisioningParams.isNfc,
+        isQr = provisioningParams.isQrProvisioning,
+    )
+}
+
+class ProvisioningArgumentsSerializer @Inject constructor() :
+    ScopedIntentSerializer<BaseProvisioningArguments> {
+
+    override fun IntentScope.write(value: BaseProvisioningArguments) {
+        val params =
+            value.provisioningParams
+                    .toBuilder()
+                    .setFlowType(value.flowType.ordinal)
+                    .setDeviceAdminDownloadInfo(value.deviceAdminDownloadInfo)
+                    .setDeviceAdminComponentName(value.deviceAdminComponentName)
+                    .setWifiInfo(value.wifiInfo)
+                    .setUseMobileData(value.useMobileData)
+                    .setIsNfc(value.isNfc)
+                    .setIsQrProvisioning(value.isQr)
+                    .build()
+        this[EXTRAS.EXTRA_PROVISIONING_PARAMS] = params
+    }
+
+    override fun IntentScope.read(): BaseProvisioningArguments =
+        parcelableExtra<ProvisioningParams>(EXTRAS.EXTRA_PROVISIONING_PARAMS).let {
+            BaseProvisioningArguments(
+                provisioningParams = it,
+                flowType = FlowType.entries[it.flowType],
+                deviceAdminDownloadInfo = it.deviceAdminDownloadInfo,
+                deviceAdminComponentName = it.deviceAdminComponentName,
+                wifiInfo = it.wifiInfo,
+                useMobileData = it.useMobileData,
+                isNfc = it.isNfc,
+                isQr = it.isQrProvisioning,
+            )
+        }
+}
diff --git a/src/com/android/managedprovisioning/preprovisioning/DownloadRoleHolderActivity.java b/src/com/android/managedprovisioning/preprovisioning/DownloadRoleHolderActivity.java
index 6d6c27e09..0ddefe138 100644
--- a/src/com/android/managedprovisioning/preprovisioning/DownloadRoleHolderActivity.java
+++ b/src/com/android/managedprovisioning/preprovisioning/DownloadRoleHolderActivity.java
@@ -26,11 +26,17 @@ import com.android.managedprovisioning.ManagedProvisioningBaseApplication;
 import com.android.managedprovisioning.R;
 import com.android.managedprovisioning.common.ErrorDialogUtils;
 import com.android.managedprovisioning.common.ErrorWrapper;
+import com.android.managedprovisioning.common.Flags;
 import com.android.managedprovisioning.common.RoleHolderProvider;
 import com.android.managedprovisioning.common.SetupGlifLayoutActivity;
+import com.android.managedprovisioning.contracts.DownloadRoleHolderContract;
 import com.android.managedprovisioning.model.ProvisioningParams;
 import com.android.managedprovisioning.preprovisioning.DownloadRoleHolderViewModel.DownloadRoleHolderViewModelFactory;
 
+import dagger.hilt.android.AndroidEntryPoint;
+
+import javax.inject.Inject;
+
 /**
  * Spinner which takes care of network connectivity if needed, and downloading of the role holder.
  *
@@ -43,12 +49,22 @@ import com.android.managedprovisioning.preprovisioning.DownloadRoleHolderViewMod
  * ErrorDialogUtils#EXTRA_FACTORY_RESET_REQUIRED} which can be used to display in a user-visible
  * dialog.
  */
-public class DownloadRoleHolderActivity extends SetupGlifLayoutActivity {
+@AndroidEntryPoint(SetupGlifLayoutActivity.class)
+public class DownloadRoleHolderActivity extends Hilt_DownloadRoleHolderActivity {
     private DownloadRoleHolderViewModel mViewModel;
 
+    @Inject
+    protected Flags mFlags;
+
+    @Inject
+    protected DownloadRoleHolderContract mContract;
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
+        if (mFlags.isCosmicRayEnabled()) {
+            mContract.validate(this, getIntent());
+        }
 
         ProvisioningParams params = getIntent().getParcelableExtra(EXTRA_PROVISIONING_PARAMS);
         if (params.roleHolderDownloadInfo == null) {
@@ -72,7 +88,7 @@ public class DownloadRoleHolderActivity extends SetupGlifLayoutActivity {
     }
 
     private void onStateChanged(Integer state) {
-        switch(state) {
+        switch (state) {
             case DownloadRoleHolderViewModel.STATE_IDLE:
                 break;
             case DownloadRoleHolderViewModel.STATE_DOWNLOADING:
diff --git a/src/com/android/managedprovisioning/preprovisioning/PreProvisioningActivity.java b/src/com/android/managedprovisioning/preprovisioning/PreProvisioningActivity.java
index 9cbc2001e..bde89623c 100644
--- a/src/com/android/managedprovisioning/preprovisioning/PreProvisioningActivity.java
+++ b/src/com/android/managedprovisioning/preprovisioning/PreProvisioningActivity.java
@@ -72,6 +72,7 @@ import com.android.managedprovisioning.common.DefaultFeatureFlagChecker;
 import com.android.managedprovisioning.common.DefaultIntentResolverChecker;
 import com.android.managedprovisioning.common.DefaultPackageInstallChecker;
 import com.android.managedprovisioning.common.DeviceManagementRoleHolderUpdaterHelper;
+import com.android.managedprovisioning.common.Flags;
 import com.android.managedprovisioning.common.GetProvisioningModeUtils;
 import com.android.managedprovisioning.common.ManagedProvisioningSharedPreferences;
 import com.android.managedprovisioning.common.ProvisionLogger;
@@ -85,6 +86,8 @@ import com.android.managedprovisioning.common.ThemeHelper;
 import com.android.managedprovisioning.common.ThemeHelper.DefaultNightModeChecker;
 import com.android.managedprovisioning.common.ThemeHelper.DefaultSetupWizardBridge;
 import com.android.managedprovisioning.common.Utils;
+import com.android.managedprovisioning.contracts.DownloadRoleHolderArguments;
+import com.android.managedprovisioning.contracts.DownloadRoleHolderContract;
 import com.android.managedprovisioning.model.ProvisioningParams;
 import com.android.managedprovisioning.preprovisioning.PreProvisioningActivityController.UiParams;
 import com.android.managedprovisioning.provisioning.AdminIntegratedFlowPrepareActivity;
@@ -97,7 +100,12 @@ import com.google.android.setupcompat.logging.SetupMetricsLogger;
 import com.google.android.setupcompat.util.WizardManagerHelper;
 import com.google.android.setupdesign.transition.TransitionHelper;
 
-public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
+import dagger.hilt.android.AndroidEntryPoint;
+
+import javax.inject.Inject;
+
+@AndroidEntryPoint(SetupGlifLayoutActivity.class)
+public class PreProvisioningActivity extends Hilt_PreProvisioningActivity implements
         SimpleDialog.SimpleDialogListener, PreProvisioningActivityController.Ui {
 
     private static final int ENCRYPT_DEVICE_REQUEST_CODE = 1;
@@ -126,7 +134,7 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
             "ShowPreProvisioningScreen";
 
     private PreProvisioningActivityController mController;
-    private ControllerProvider mControllerProvider;
+    private final ControllerProvider mControllerProvider;
     private final AccessibilityContextMenuMaker mContextMenuMaker;
     private PreProvisioningActivityBridge mBridge;
     private boolean mShouldForwardTransition;
@@ -140,15 +148,20 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
     protected ScreenKey mScreenKey;
     protected String setupMetricScreenName;
 
+    @Inject
+    protected Flags mFlags;
+    @Inject
+    protected DownloadRoleHolderContract mDownloadRoleHolderContract;
+
     public PreProvisioningActivity() {
         this(activity ->
-                new PreProvisioningActivityController(activity, activity),
+                        new PreProvisioningActivityController(activity, activity),
                 null,
                 new Utils(),
                 new SettingsFacade(),
                 new ThemeHelper(
-                    new DefaultNightModeChecker(),
-                    new DefaultSetupWizardBridge()),
+                        new DefaultNightModeChecker(),
+                        new DefaultSetupWizardBridge()),
                 RoleHolderProvider.DEFAULT,
                 RoleHolderUpdaterProvider.DEFAULT);
     }
@@ -325,7 +338,7 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
             case GET_PROVISIONING_MODE_REQUEST_CODE:
                 mShouldForwardTransition = true;
                 if (resultCode == RESULT_OK) {
-                    if(data != null && mController.updateProvisioningParamsFromIntent(data)) {
+                    if (data != null && mController.updateProvisioningParamsFromIntent(data)) {
                         mController.showUserConsentScreen();
                     } else {
                         ProvisionLogger.loge(
@@ -456,7 +469,7 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
         if (TextUtils.isEmpty(
                 mRoleHolderUpdaterProvider.getPackageName(this))) {
             ProvisionLogger.logw("Role holder requested update, but there is no role "
-                            + "holder updater present. Restarting the role holder.");
+                    + "holder updater present. Restarting the role holder.");
             boolean isProvisioningStarted = mController.startAppropriateProvisioning(
                     getIntent(),
                     createRoleHolderAdditionalExtras(resultCode),
@@ -511,9 +524,9 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
                     mAnalyticsTracker.logRoleHolderUpdaterUpdateFailed();
                     failRoleHolderUpdate();
                     ProvisionLogger.loge("Failed to start provisioning after a "
-                        + "platform-requested role holder update. Result is " + resultCode
-                        + " and allow offline provisioning is "
-                        + mController.getParams().allowOffline);
+                            + "platform-requested role holder update. Result is " + resultCode
+                            + " and allow offline provisioning is "
+                            + mController.getParams().allowOffline);
                     SetupMetricsLogger.logMetrics(this, mScreenKey,
                             SetupMetric.ofError(setupMetricScreenName, resultCode));
                 }
@@ -801,7 +814,7 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
                         isRoleHolderRequestedUpdate));
         mAnalyticsTracker.logRoleHolderUpdaterUpdateStart();
         getTransitionHelper().startActivityForResultWithTransition(
-                 this,
+                this,
                 intent,
                 isRoleHolderRequestedUpdate
                         ? START_ROLE_HOLDER_REQUESTED_UPDATE_REQUEST_CODE
@@ -829,11 +842,22 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
     @Override
     public void startPlatformDrivenRoleHolderDownload() {
         mAnalyticsTracker.logPlatformRoleHolderUpdateStart();
-        Intent intent = new Intent(this,
-                getActivityForScreen(ManagedProvisioningScreens.DOWNLOAD_ROLE_HOLDER));
-        WizardManagerHelper.copyWizardManagerExtras(getIntent(), intent);
-        intent.putExtra(ProvisioningParams.EXTRA_PROVISIONING_PARAMS,
-                mController.getParams());
+
+        Intent intent;
+        if (mFlags.isCosmicRayEnabled()) {
+            intent = mDownloadRoleHolderContract.createIntent(
+                    this,
+                    new DownloadRoleHolderArguments(
+                            mDownloadRoleHolderContract.getSuwArgumentsSerializer().read(
+                                    getIntent()),
+                            requireNonNull(mController.getParams())));
+        } else {
+            intent = new Intent(this,
+                    getActivityForScreen(ManagedProvisioningScreens.DOWNLOAD_ROLE_HOLDER));
+            WizardManagerHelper.copyWizardManagerExtras(getIntent(), intent);
+            intent.putExtra(ProvisioningParams.EXTRA_PROVISIONING_PARAMS,
+                    mController.getParams());
+        }
         getTransitionHelper().startActivityForResultWithTransition(
                 this, intent, DOWNLOAD_DEVICE_MANAGEMENT_ROLE_HOLDER_FROM_PLATFORM_REQUEST_CODE);
     }
@@ -860,6 +884,7 @@ public class PreProvisioningActivity extends SetupGlifLayoutActivity implements
 
     // TODO: The below group of methods do not belong in the activity.
     // Move them to the controller instead.
+
     /**
      * Starts either the admin-integrated or the legacy flow, depending on the device state and
      * DPC capabilities.
diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index 4a433d5df..185ca9827 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -26,6 +26,8 @@ android_robolectric_test {
         "androidx.test.rules",
         "androidx.core_core",
         "testng", // used only for assertThrows
+        "android_onboarding.contracts.testing",
+        "hilt_android_testing",
     ],
 
     static_libs: [
diff --git a/tests/robotests/src/com/android/managedprovisioning/contracts/DownloadRoleHolderContractTest.kt b/tests/robotests/src/com/android/managedprovisioning/contracts/DownloadRoleHolderContractTest.kt
new file mode 100644
index 000000000..7269efd34
--- /dev/null
+++ b/tests/robotests/src/com/android/managedprovisioning/contracts/DownloadRoleHolderContractTest.kt
@@ -0,0 +1,45 @@
+package com.android.managedprovisioning.contracts
+
+import android.os.Build
+import com.android.onboarding.contracts.testing.assertArgumentEncodesCorrectly
+import dagger.hilt.android.testing.HiltAndroidRule
+import dagger.hilt.android.testing.HiltAndroidTest
+import dagger.hilt.android.testing.HiltTestApplication
+import javax.inject.Inject
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.annotation.Config
+
+@RunWith(RobolectricTestRunner::class)
+@Config(minSdk = Build.VERSION_CODES.VANILLA_ICE_CREAM, application = HiltTestApplication::class)
+@HiltAndroidTest
+class DownloadRoleHolderContractTest {
+    @get:Rule
+    val hiltRule = HiltAndroidRule(this)
+
+    @Inject
+    lateinit var contract: DownloadRoleHolderContract
+
+    @Before
+    fun setUp() {
+        hiltRule.inject()
+    }
+
+    @Test
+    fun completeArgument_encodesCorrectly() {
+        val rhDownloadInfo = requireNotNull(aProvisioningParams.roleHolderDownloadInfo)
+        val args =
+            DownloadRoleHolderArguments(
+                suwArguments = aSuwArguments,
+                provisioningParams = aProvisioningParams,
+                location = rhDownloadInfo.location,
+                checksum = Checksum.SignatureChecksum(rhDownloadInfo.signatureChecksum),
+                cookieHeader = rhDownloadInfo.cookieHeader,
+                minVersion = rhDownloadInfo.minVersion,
+            )
+        assertArgumentEncodesCorrectly(contract, args)
+    }
+}
diff --git a/tests/robotests/src/com/android/managedprovisioning/contracts/data.kt b/tests/robotests/src/com/android/managedprovisioning/contracts/data.kt
new file mode 100644
index 000000000..ccea229e5
--- /dev/null
+++ b/tests/robotests/src/com/android/managedprovisioning/contracts/data.kt
@@ -0,0 +1,89 @@
+package com.android.managedprovisioning.contracts
+
+import android.content.ComponentName
+import android.os.Bundle
+import com.android.managedprovisioning.model.DisclaimersParam
+import com.android.managedprovisioning.model.PackageDownloadInfo
+import com.android.managedprovisioning.model.ProvisioningParams
+import com.android.managedprovisioning.model.WifiInfo
+import com.android.onboarding.contracts.provisioning.ACTIONS
+import com.android.onboarding.contracts.provisioning.FLAGS
+import com.android.onboarding.contracts.setupwizard.SuwArguments
+
+/**
+ * Stub [SuwArguments] prebuilt for testing. Further customisation is available via
+ * [SuwArguments.copy] method
+ */
+val aSuwArguments =
+    SuwArguments(
+        isSubactivityFirstLaunched = true,
+        isSuwSuggestedActionFlow = true,
+        isSetupFlow = true,
+        preDeferredSetup = true,
+        deferredSetup = false,
+        firstRun = true,
+        portalSetup = false,
+        hasMultipleUsers = false,
+        theme = "dark",
+        wizardBundle = Bundle.EMPTY
+    )
+
+val aProvisioningParams: ProvisioningParams = run {
+    val flowType = FlowType.AdminIntegrated
+    val deviceAdminDownloadInfo =
+        PackageDownloadInfo.Builder()
+                .setLocation("some/body")
+                .setMinVersion(69)
+                .setPackageChecksum("checksum".toByteArray())
+                .build()
+    val deviceAdminComponentName = ComponentName("org.test", "Test")
+    val wifiInfo = WifiInfo.Builder().setSsid("MyFi").setPassword("secret").build()
+    val packageDownloadInfo = PackageDownloadInfo.Builder()
+            .setLocation("http://test.local")
+            .setMinVersion(420)
+            .setSignatureChecksum("69".toByteArray())
+            .setCookieHeader("I'm a cookie monster, gimmie all your cookies or else!")
+            .build()
+    ProvisioningParams.Builder()
+            .setProvisioningAction(ACTIONS.ACTION_PROVISION_MANAGED_PROFILE)
+            .setDeviceAdminComponentName(ComponentName("org.test", "local"))
+            .setFlowType(flowType.ordinal)
+            .setDeviceAdminDownloadInfo(deviceAdminDownloadInfo)
+            .setDeviceAdminComponentName(deviceAdminComponentName)
+            .setWifiInfo(wifiInfo)
+            .setUseMobileData(true)
+            .setIsNfc(true)
+            .setIsQrProvisioning(true)
+            .setSupportUrl("https://help.test")
+            .setInitiatorRequestedProvisioningModes(
+                FLAGS.FLAG_SUPPORTED_MODES_ORGANIZATION_OWNED or FLAGS.FLAG_SUPPORTED_MODES_DEVICE_OWNER
+            )
+            .setDisclaimersParam(
+                DisclaimersParam.Builder()
+                        .setDisclaimers(
+                            arrayOf(
+                                DisclaimersParam.Disclaimer("a", "a.html"),
+                                DisclaimersParam.Disclaimer("b", "b.html"),
+                            )
+                        )
+                        .build()
+            )
+            .setRoleHolderDownloadInfo(packageDownloadInfo)
+            .setIsOrganizationOwnedProvisioning(true)
+            .setSkipEducationScreens(true)
+            .setReturnBeforePolicyCompliance(true)
+            .setDeviceOwnerPermissionGrantOptOut(true)
+            .setAllowProvisioningAfterUserSetupComplete(true)
+            .build()
+}
+
+val aProvisioningArguments: BaseProvisioningArguments = BaseProvisioningArguments(
+    provisioningParams = aProvisioningParams,
+    flowType = FlowType.entries[aProvisioningParams.flowType],
+    deviceAdminDownloadInfo = aProvisioningParams.deviceAdminDownloadInfo,
+    deviceAdminComponentName = aProvisioningParams.deviceAdminComponentName,
+    wifiInfo = aProvisioningParams.wifiInfo,
+    useMobileData = aProvisioningParams.useMobileData,
+    isNfc = aProvisioningParams.isNfc,
+    isQr = aProvisioningParams.isQrProvisioning,
+)