[ffx] Abstract out flashing logic between versions.

Since unlocking logic is different from v1 and v2, and it's internal to
the product, some of the flashing logic needs to be abstracted out in
order create space for unlocking logic.

Bug: 85901
Change-Id: Ie7aa37c046f9a48694f7170a4b21026f44af4cd7
Reviewed-on: https://fuchsia-review.googlesource.com/c/fuchsia/+/590388
Commit-Queue: Matthew Boetger <boetger@google.com>
Reviewed-by: Andrew Davies <awdavies@google.com>
diff --git a/src/developer/ffx/plugins/target/flash/src/manifest/mod.rs b/src/developer/ffx/plugins/target/flash/src/manifest/mod.rs
index fc2bfee..f44325a 100644
--- a/src/developer/ffx/plugins/target/flash/src/manifest/mod.rs
+++ b/src/developer/ffx/plugins/target/flash/src/manifest/mod.rs
@@ -14,7 +14,7 @@
     async_trait::async_trait,
     chrono::{DateTime, Duration, Utc},
     errors::{ffx_bail, ffx_error},
-    ffx_flash_args::FlashCommand,
+    ffx_flash_args::{FlashCommand, OemFile},
     fidl::{
         endpoints::{create_endpoints, ServerEnd},
         Error as FidlError,
@@ -39,10 +39,20 @@
 pub(crate) mod v2;
 pub(crate) mod v3;
 
+pub(crate) const MISSING_PRODUCT: &str = "Manifest does not contain product";
+
+const REBOOT_ERR: &str = "Failed to reboot your device. \
+                          Power cycle the device manually and try again.";
+
 pub(crate) const UNKNOWN_VERSION: &str = "Unknown flash manifest version";
 const LARGE_FILE: &str = "large file, please wait... ";
 const REVISION_VAR: &str = "hw-revision";
 
+const LOCKED_VAR: &str = "vx-locked";
+
+pub(crate) const UNLOCK_ERR: &str = "The product requires the target is unlocked. \
+                                     Please unlock target and try again.";
+
 #[derive(Default, Deserialize)]
 pub struct Images(Vec<Image>);
 
@@ -53,6 +63,19 @@
     // Ignore the rest of the fields
 }
 
+pub(crate) trait Partition {
+    fn name(&self) -> &str;
+    fn file(&self) -> &str;
+    fn variable(&self) -> Option<&str>;
+    fn variable_value(&self) -> Option<&str>;
+}
+
+pub(crate) trait Product<P> {
+    fn bootloader_partitions(&self) -> &Vec<P>;
+    fn partitions(&self) -> &Vec<P>;
+    fn oem_files(&self) -> &Vec<OemFile>;
+}
+
 #[async_trait(?Send)]
 pub(crate) trait Flash {
     async fn flash<W, F>(
@@ -477,6 +500,98 @@
     })
 }
 
+pub(crate) async fn stage_oem_files<W: Write, F: FileResolver + Sync>(
+    writer: &mut W,
+    file_resolver: &mut F,
+    resolve: bool,
+    oem_files: &Vec<OemFile>,
+    fastboot_proxy: &FastbootProxy,
+) -> Result<()> {
+    for oem_file in oem_files {
+        stage_file(writer, file_resolver, resolve, oem_file.file(), &fastboot_proxy).await?;
+        writeln!(writer, "Sending command \"{}\"", oem_file.command())?;
+        fastboot_proxy.oem(oem_file.command()).await?.map_err(|_| {
+            anyhow!("There was an error sending oem command \"{}\"", oem_file.command())
+        })?;
+    }
+    Ok(())
+}
+
+pub(crate) async fn flash_partitions<W: Write, F: FileResolver + Sync, P: Partition>(
+    writer: &mut W,
+    file_resolver: &mut F,
+    partitions: &Vec<P>,
+    fastboot_proxy: &FastbootProxy,
+) -> Result<()> {
+    for partition in partitions {
+        match (partition.variable(), partition.variable_value()) {
+            (Some(var), Some(value)) => {
+                if verify_variable_value(var, value, fastboot_proxy).await? {
+                    flash_partition(
+                        writer,
+                        file_resolver,
+                        partition.name(),
+                        partition.file(),
+                        fastboot_proxy,
+                    )
+                    .await?;
+                }
+            }
+            _ => {
+                flash_partition(
+                    writer,
+                    file_resolver,
+                    partition.name(),
+                    partition.file(),
+                    fastboot_proxy,
+                )
+                .await?
+            }
+        }
+    }
+    Ok(())
+}
+
+pub(crate) async fn flash_and_reboot<W, F, Part, P>(
+    writer: &mut W,
+    file_resolver: &mut F,
+    product: &P,
+    fastboot_proxy: &FastbootProxy,
+    cmd: FlashCommand,
+) -> Result<()>
+where
+    W: Write,
+    F: FileResolver + Sync,
+    Part: Partition,
+    P: Product<Part>,
+{
+    flash_partitions(writer, file_resolver, product.bootloader_partitions(), fastboot_proxy)
+        .await?;
+    if product.bootloader_partitions().len() > 0 && !cmd.no_bootloader_reboot {
+        reboot_bootloader(writer, &fastboot_proxy)
+            .await
+            .map_err(|_| ffx_error!("{}", REBOOT_ERR))?;
+    }
+    stage_oem_files(writer, file_resolver, false, &cmd.oem_stage, fastboot_proxy).await?;
+    flash_partitions(writer, file_resolver, product.partitions(), fastboot_proxy).await?;
+    stage_oem_files(writer, file_resolver, true, product.oem_files(), fastboot_proxy).await?;
+    finish(writer, fastboot_proxy).await
+}
+
+pub(crate) async fn finish<W: Write>(writer: &mut W, fastboot_proxy: &FastbootProxy) -> Result<()> {
+    if fastboot_proxy.erase("misc").await?.is_err() {
+        log::debug!("Could not erase misc partition");
+    }
+    fastboot_proxy.set_active("a").await?.map_err(|_| anyhow!("Could not set active slot"))?;
+    fastboot_proxy.continue_boot().await?.map_err(|_| anyhow!("Could not reboot device"))?;
+    writeln!(writer, "Continuing to boot - this could take awhile")?;
+    Ok(())
+}
+
+pub(crate) async fn is_locked(fastboot_proxy: &FastbootProxy) -> Result<bool> {
+    verify_variable_value(LOCKED_VAR, "no", &fastboot_proxy).await.map(|l| !l)
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // tests
 
diff --git a/src/developer/ffx/plugins/target/flash/src/manifest/v1.rs b/src/developer/ffx/plugins/target/flash/src/manifest/v1.rs
index 13bf61c..0f3d692 100644
--- a/src/developer/ffx/plugins/target/flash/src/manifest/v1.rs
+++ b/src/developer/ffx/plugins/target/flash/src/manifest/v1.rs
@@ -5,26 +5,20 @@
 use {
     crate::{
         file::FileResolver,
-        manifest::{flash_partition, reboot_bootloader, stage_file, verify_variable_value, Flash},
+        manifest::{
+            flash_and_reboot, is_locked, Flash, Partition as PartitionTrait,
+            Product as ProductTrait, MISSING_PRODUCT, UNLOCK_ERR,
+        },
     },
-    anyhow::{anyhow, Result},
+    anyhow::Result,
     async_trait::async_trait,
-    errors::{ffx_bail, ffx_error},
+    errors::ffx_bail,
     ffx_flash_args::{FlashCommand, OemFile},
     fidl_fuchsia_developer_bridge::FastbootProxy,
     serde::{Deserialize, Serialize},
     std::io::Write,
 };
 
-const REBOOT_ERR: &str = "Failed to reboot your device.  \
-                          Power cycle the device manually and try again.";
-pub(crate) const MISSING_PRODUCT: &str = "Manifest does not contain product";
-
-const UNLOCK_ERR: &str = "The product requires the target is unlocked. \
-                          Please unlock target and try again.";
-
-const LOCKED_VAR: &str = "vx-locked";
-
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub(crate) struct Product {
     pub(crate) name: String,
@@ -35,6 +29,20 @@
     pub(crate) requires_unlock: bool,
 }
 
+impl ProductTrait<Partition> for Product {
+    fn bootloader_partitions(&self) -> &Vec<Partition> {
+        &self.bootloader_partitions
+    }
+
+    fn partitions(&self) -> &Vec<Partition> {
+        &self.partitions
+    }
+
+    fn oem_files(&self) -> &Vec<OemFile> {
+        &self.oem_files
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub(crate) struct Partition(
     String,
@@ -50,22 +58,24 @@
         variable: Option<String>,
         variable_value: Option<String>,
     ) -> Self {
-        Partition(name, file, variable, variable_value)
+        Self(name, file, variable, variable_value)
     }
+}
 
-    pub(crate) fn name(&self) -> &str {
+impl PartitionTrait for Partition {
+    fn name(&self) -> &str {
         self.0.as_str()
     }
 
-    pub(crate) fn file(&self) -> &str {
+    fn file(&self) -> &str {
         self.1.as_str()
     }
 
-    pub(crate) fn variable(&self) -> Option<&str> {
+    fn variable(&self) -> Option<&str> {
         self.2.as_ref().map(|s| s.as_str())
     }
 
-    pub(crate) fn variable_value(&self) -> Option<&str> {
+    fn variable_value(&self) -> Option<&str> {
         self.3.as_ref().map(|s| s.as_str())
     }
 }
@@ -90,83 +100,13 @@
             Some(res) => res,
             None => ffx_bail!("{} {}", MISSING_PRODUCT, cmd.product),
         };
-        if product.requires_unlock
-            && !verify_variable_value(LOCKED_VAR, "no", &fastboot_proxy).await?
-        {
+        if product.requires_unlock && is_locked(&fastboot_proxy).await? {
             ffx_bail!("{}", UNLOCK_ERR);
         }
-        flash_partitions(writer, file_resolver, &product.bootloader_partitions, &fastboot_proxy)
-            .await?;
-        if product.bootloader_partitions.len() > 0 && !cmd.no_bootloader_reboot {
-            reboot_bootloader(writer, &fastboot_proxy)
-                .await
-                .map_err(|_| ffx_error!("{}", REBOOT_ERR))?;
-        }
-        stage_oem_files(writer, file_resolver, false, &cmd.oem_stage, &fastboot_proxy).await?;
-        flash_partitions(writer, file_resolver, &product.partitions, &fastboot_proxy).await?;
-        stage_oem_files(writer, file_resolver, true, &product.oem_files, &fastboot_proxy).await?;
-        if fastboot_proxy.erase("misc").await?.is_err() {
-            log::debug!("Could not erase misc partition");
-        }
-        fastboot_proxy.set_active("a").await?.map_err(|_| anyhow!("Could not set active slot"))?;
-        fastboot_proxy.continue_boot().await?.map_err(|_| anyhow!("Could not reboot device"))?;
-        writeln!(writer, "Continuing to boot - this could take awhile")?;
-        Ok(())
+        flash_and_reboot(writer, file_resolver, product, &fastboot_proxy, cmd).await
     }
 }
 
-pub(crate) async fn stage_oem_files<W: Write, F: FileResolver + Sync>(
-    writer: &mut W,
-    file_resolver: &mut F,
-    resolve: bool,
-    oem_files: &Vec<OemFile>,
-    fastboot_proxy: &FastbootProxy,
-) -> Result<()> {
-    for oem_file in oem_files {
-        stage_file(writer, file_resolver, resolve, oem_file.file(), &fastboot_proxy).await?;
-        writeln!(writer, "Sending command \"{}\"", oem_file.command())?;
-        fastboot_proxy.oem(oem_file.command()).await?.map_err(|_| {
-            anyhow!("There was an error sending oem command \"{}\"", oem_file.command())
-        })?;
-    }
-    Ok(())
-}
-
-pub(crate) async fn flash_partitions<W: Write, F: FileResolver + Sync>(
-    writer: &mut W,
-    file_resolver: &mut F,
-    partitions: &Vec<Partition>,
-    fastboot_proxy: &FastbootProxy,
-) -> Result<()> {
-    for partition in partitions {
-        match (partition.variable(), partition.variable_value()) {
-            (Some(var), Some(value)) => {
-                if verify_variable_value(var, value, fastboot_proxy).await? {
-                    flash_partition(
-                        writer,
-                        file_resolver,
-                        partition.name(),
-                        partition.file(),
-                        fastboot_proxy,
-                    )
-                    .await?;
-                }
-            }
-            _ => {
-                flash_partition(
-                    writer,
-                    file_resolver,
-                    partition.name(),
-                    partition.file(),
-                    fastboot_proxy,
-                )
-                .await?
-            }
-        }
-    }
-    Ok(())
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // tests
 
diff --git a/src/developer/ffx/plugins/target/flash/src/manifest/v2.rs b/src/developer/ffx/plugins/target/flash/src/manifest/v2.rs
index a3bf960..f0b5ae9 100644
--- a/src/developer/ffx/plugins/target/flash/src/manifest/v2.rs
+++ b/src/developer/ffx/plugins/target/flash/src/manifest/v2.rs
@@ -5,16 +5,24 @@
 use {
     crate::{
         file::FileResolver,
-        manifest::{v1::FlashManifest as FlashManifestV1, verify_hardware, Flash},
+        manifest::{
+            flash_and_reboot, is_locked, v1::FlashManifest as FlashManifestV1, verify_hardware,
+            Flash, MISSING_PRODUCT, UNLOCK_ERR,
+        },
     },
     anyhow::Result,
     async_trait::async_trait,
+    errors::ffx_bail,
     ffx_flash_args::FlashCommand,
     fidl_fuchsia_developer_bridge::FastbootProxy,
     serde::{Deserialize, Serialize},
     std::io::Write,
 };
 
+const MISSING_CREDENTIALS: &str =
+    "The flash manifest is missing the credential files to unlock this device.\n\
+     Please unlock the target and try again.";
+
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub(crate) struct FlashManifest {
     pub(crate) hw_revision: String,
@@ -40,7 +48,19 @@
         if !cmd.skip_verify {
             verify_hardware(&self.hw_revision, &fastboot_proxy).await?;
         }
-        self.v1.flash(writer, file_resolver, fastboot_proxy, cmd).await
+        let product = match self.v1.0.iter().find(|product| product.name == cmd.product) {
+            Some(res) => res,
+            None => ffx_bail!("{} {}", MISSING_PRODUCT, cmd.product),
+        };
+        if product.requires_unlock && is_locked(&fastboot_proxy).await? {
+            if self.credentials.len() == 0 {
+                ffx_bail!("{}", MISSING_CREDENTIALS);
+            } else {
+                //TODO: Try unlock the device.
+                ffx_bail!("{}", UNLOCK_ERR);
+            }
+        }
+        flash_and_reboot(writer, file_resolver, product, &fastboot_proxy, cmd).await
     }
 }
 
@@ -91,6 +111,25 @@
         ]
     }"#;
 
+    const NO_CREDS_MANIFEST: &'static str = r#"{
+        "hw_revision": "zedboot",
+        "products": [
+            {
+                "name": "zedboot",
+                "requires_unlock": false,
+                "bootloader_partitions": [],
+                "partitions": [
+                    ["test1", "path1"],
+                    ["test2", "path2"],
+                    ["test3", "path3"],
+                    ["test4", "path4"],
+                    ["test5", "path5"]
+                ],
+                "oem_files": []
+            }
+        ]
+    }"#;
+
     #[fuchsia_async::run_singlethreaded(test)]
     async fn test_matching_revision_should_work() -> Result<()> {
         let v: FlashManifest = from_str(MANIFEST)?;
@@ -134,4 +173,27 @@
             .is_err());
         Ok(())
     }
+
+    #[fuchsia_async::run_singlethreaded(test)]
+    async fn test_no_creds_and_requires_unlock_should_err() -> Result<()> {
+        let v: FlashManifest = from_str(NO_CREDS_MANIFEST)?;
+        let tmp_file = NamedTempFile::new().expect("tmp access failed");
+        let tmp_file_name = tmp_file.path().to_string_lossy().to_string();
+        let (_, proxy) = setup();
+        let mut writer = Vec::<u8>::new();
+        assert!(v
+            .flash(
+                &mut writer,
+                &mut TestResolver::new(),
+                proxy,
+                FlashCommand {
+                    manifest: Some(PathBuf::from(tmp_file_name)),
+                    product: "zedboot".to_string(),
+                    ..Default::default()
+                }
+            )
+            .await
+            .is_err());
+        Ok(())
+    }
 }