factory_ufs: support UFS 3.1 provisioning

In this CL, we
(1) support UFS 3.1 (ufs_spec/ufs_3_1.rs). Before we only support
    UFS 2.1.
(2) parse the header of current UFS config to get the version info.
(3) provision the UFS according to the version info.
(4) split config.rs to ufs_spec/ufs_2_1.rs, ufs_spec/ufs_common.rs and
    ufs_spec/ufs_config_descriptor.rs.
(5) define trait CommonDescriptorTrait that the DeviceConfigDescriptor,
    UnitConfigDescriptor and ConfigDescriptor must implement.
(6) define trait ProvisionLun that UnitConfigDescriptor must implement.
(7) define trait ProvisionLun0 that ConfigDescriptor must implement.

BUG=b:228796554
TEST=run on RVP

Change-Id: I57d3100f1e709d451a5798b434f5ba11e85f55ac
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/factory_installer/+/3692623
Tested-by: Phoebe Wang <phoebewang@chromium.org>
Reviewed-by: Yu-An Wang <wyuang@google.com>
Commit-Queue: Phoebe Wang <phoebewang@chromium.org>
Reviewed-by: Daniil Lunev <dlunev@chromium.org>
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 9e17542..1bd5c66 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -6,6 +6,7 @@
 [dependencies]
 anyhow = "1.0.0"
 bincode = "1.0.1"
+byteorder = "1.4.3"
 clap = { version = "3.1.0", features = ["derive"] }
 glob = "0.3.0"
 libc = "0.2.0"
diff --git a/rust/src/bin/factory_ufs.rs b/rust/src/bin/factory_ufs.rs
index 857045f..b180334 100644
--- a/rust/src/bin/factory_ufs.rs
+++ b/rust/src/bin/factory_ufs.rs
@@ -1,14 +1,17 @@
-// Copyright 2022 The Chromium OS Authors. All rights reserved.
+// Copyright 2022 The ChromiumOS Authors.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 use anyhow::{self, Context, Result};
-use factory_installer::factory_ufs::ufs::{self, PathMatchError};
+use factory_installer::factory_ufs::ufs_provision::{self, PathMatchError};
+use factory_installer::factory_ufs::ufs_spec::ufs_config_descriptor::{
+    UFS2ConfigDescriptor, UFS3ConfigDescriptor,
+};
 
 const UFS_BSG_ERROR_MSG: &str = "Unexpected error occurs when finding";
 
 fn main() -> Result<()> {
-    let ufs_bsg_node_path = ufs::get_ufs_bsg_node();
+    let ufs_bsg_node_path = ufs_provision::get_ufs_bsg_node();
     let ufs_bsg_node_path = match ufs_bsg_node_path {
         Ok(path) => path,
         Err(err) => {
@@ -18,19 +21,39 @@
                     return Ok(());
                 }
             }
-            return Err(anyhow::anyhow!(
-                "{} UFS BSG node: {}",
-                UFS_BSG_ERROR_MSG,
-                err
-            ));
+            anyhow::bail!("{} UFS BSG node: {}", UFS_BSG_ERROR_MSG, err)
         }
     };
-    let ufs_bsg_sys_path =
-        ufs::get_ufs_bsg_sys().context(format!("{} UFS BSG sys path", UFS_BSG_ERROR_MSG))?;
-    println!("UFS BSG node path: {}", ufs_bsg_node_path);
-    println!("UFS BSG sys path: {}", ufs_bsg_sys_path);
+    let ufs_bsg_sys_path = ufs_provision::get_ufs_bsg_sys()
+        .context(format!("{} UFS BSG sys path", UFS_BSG_ERROR_MSG))?;
+    println!("UFS BSG node path: {}", ufs_bsg_node_path.display());
+    println!("UFS BSG sys path: {}", ufs_bsg_sys_path.display());
 
-    ufs::do_provisioning_if_needed(&ufs_bsg_node_path, &ufs_bsg_sys_path)?;
+    // We need to parse the header of current UFS config to decide which UFS
+    // spec to load.
+    let cur_config_binary = ufs_provision::get_current_config_binary(&ufs_bsg_node_path)?;
+    let header_value = ufs_provision::read_header(&cur_config_binary)?;
 
-    Ok(())
+    match header_value {
+        0x90 => {
+            println!("Provision UFS 2.1...");
+            ufs_provision::do_provisioning_if_needed::<UFS2ConfigDescriptor>(
+                &ufs_bsg_node_path,
+                &ufs_bsg_sys_path,
+                &cur_config_binary,
+            )
+        }
+        0xe6 => {
+            println!("Provision UFS 3.1...");
+            ufs_provision::do_provisioning_if_needed::<UFS3ConfigDescriptor>(
+                &ufs_bsg_node_path,
+                &ufs_bsg_sys_path,
+                &cur_config_binary,
+            )
+        }
+        value => anyhow::bail!(
+            "Cannot parse current UFS config! Unknown header value: {}",
+            value
+        ),
+    }
 }
diff --git a/rust/src/factory_ufs/config.rs b/rust/src/factory_ufs/config.rs
deleted file mode 100644
index 24bf64d..0000000
--- a/rust/src/factory_ufs/config.rs
+++ /dev/null
@@ -1,130 +0,0 @@
-// Copyright 2022 The Chromium OS Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-use serde::{Deserialize, Serialize};
-
-/// Defines the logical block size of each LUN.
-///
-/// The size of logical block size is calculated by 2^(LOGICAL_BLOCK_SIZE),
-/// which equals to 4k here.
-pub const LOGICAL_BLOCK_SIZE: u8 = 0x0C;
-/// Defines the privisioning type for the UFS device.
-///
-/// 0x03 means enabling thin provisioning and setting TPRZ bit to one.
-pub const PROVISIONING_TYPE: u8 = 0x03;
-
-pub const UFS_NODE_PATTERN: &str = "/dev/bsg/ufs-bsg*";
-pub const SYS_DEV_PATTERN: &str = "/sys/devices/*/*/host*/ufs-bsg*";
-pub const UFS_CONFIG_NAME: &str = "config_desc_data_ind_0";
-
-/// Defines the config descriptor for device.
-///
-/// `DeviceConfigDescriptor` and `UnitConfigDescriptor` specifies the
-/// configuration for UFS provisioning. Both of them need to be serialized and
-/// deserialized in big-endian.
-/// For spec, please refer to `UFS 2.2 table 14-10`.
-#[derive(Serialize, Deserialize, Debug, PartialEq)]
-#[repr(C)]
-pub struct DeviceConfigDescriptor {
-    pub b_length: u8,
-    pub b_descriptor_type: u8,
-    pub b_conf_desc_continue: u8,
-    pub b_boot_enable: u8,
-    pub b_descr_access_en: u8,
-    pub b_init_power_mode: u8,
-    pub b_high_priority_lun: u8,
-    pub b_secure_removal_type: u8,
-    pub b_init_active_icc_level: u8,
-    pub w_periodic_rtc_update: u16,
-    pub b_hpb_control: u8,
-    pub b_rpmb_region_enable: u8,
-    pub b_rpmb_region_1_size: u8,
-    pub b_rpmb_region_2_size: u8,
-    pub b_rpmb_region_3_size: u8,
-}
-
-impl Default for DeviceConfigDescriptor {
-    fn default() -> Self {
-        DeviceConfigDescriptor {
-            b_length: 0x90,
-            b_descriptor_type: 0x01,
-            b_conf_desc_continue: 0x00,
-            b_boot_enable: 0x01,
-            b_descr_access_en: 0x00,
-            b_init_power_mode: 0x01,
-            b_high_priority_lun: 0x7F,
-            b_secure_removal_type: 0x00,
-            b_init_active_icc_level: 0x00,
-            w_periodic_rtc_update: 0x00,
-            b_hpb_control: 0x00,
-            b_rpmb_region_enable: 0x00,
-            b_rpmb_region_1_size: 0x00,
-            b_rpmb_region_2_size: 0x00,
-            b_rpmb_region_3_size: 0x00,
-        }
-    }
-}
-
-/// Defines the config descriptor for each LUN unit.
-///
-/// For spec, please refer to `UFS 2.2 table 13-3`.
-#[derive(Serialize, Deserialize, Debug, PartialEq)]
-#[repr(C)]
-pub struct UnitConfigDescriptor {
-    pub b_lu_enable: u8,
-    pub b_boot_lun_id: u8,
-    pub b_lu_write_protect: u8,
-    pub b_memory_type: u8,
-    pub d_num_alloc_units: u32,
-    pub b_data_reliability: u8,
-    pub b_logical_block_size: u8,
-    pub b_provisioning_type: u8,
-    pub w_context_capabilities: u32,
-    pub w_lu_max_active_hpb_regions: u8,
-}
-
-impl Default for UnitConfigDescriptor {
-    fn default() -> Self {
-        UnitConfigDescriptor {
-            b_lu_enable: 0x0,
-            b_boot_lun_id: 0x0,
-            b_lu_write_protect: 0x0,
-            b_memory_type: 0x0,
-            d_num_alloc_units: 0x0,
-            b_data_reliability: 0x0,
-            b_logical_block_size: LOGICAL_BLOCK_SIZE,
-            b_provisioning_type: 0x0,
-            w_context_capabilities: 0x0,
-            w_lu_max_active_hpb_regions: 0x0,
-        }
-    }
-}
-
-/// Defines the the config descriptor for UFS provisioning.
-///
-/// This contains 1 [DeviceConfigDescriptor] and 8 [UnitConfigDescriptor]s.
-#[derive(Serialize, Deserialize, Debug, PartialEq)]
-#[repr(C)]
-pub struct ConfigDescriptor {
-    pub device_config: DeviceConfigDescriptor,
-    pub units_config: [UnitConfigDescriptor; 8],
-}
-
-impl Default for ConfigDescriptor {
-    fn default() -> ConfigDescriptor {
-        ConfigDescriptor {
-            device_config: DeviceConfigDescriptor::default(),
-            units_config: [
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-                UnitConfigDescriptor::default(),
-            ],
-        }
-    }
-}
diff --git a/rust/src/factory_ufs/mod.rs b/rust/src/factory_ufs/mod.rs
index 4d8f8ec..f7d342e 100644
--- a/rust/src/factory_ufs/mod.rs
+++ b/rust/src/factory_ufs/mod.rs
@@ -1,6 +1,6 @@
-// Copyright 2022 The Chromium OS Authors. All rights reserved.
+// Copyright 2022 The ChromiumOS Authors.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-pub mod config;
-pub mod ufs;
+pub mod ufs_provision;
+pub mod ufs_spec;
diff --git a/rust/src/factory_ufs/ufs.rs b/rust/src/factory_ufs/ufs.rs
deleted file mode 100644
index 6ac99af..0000000
--- a/rust/src/factory_ufs/ufs.rs
+++ /dev/null
@@ -1,188 +0,0 @@
-// Copyright 2022 The Chromium OS Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-use crate::factory_ufs::config;
-use crate::utils::process_utils::Command;
-
-use anyhow::{bail, Result};
-use bincode::config as bincode_config;
-use glob::glob;
-use std::env::temp_dir;
-use std::error::Error;
-use std::fmt::{Display, Formatter, Result as fmt_Result};
-use std::fs::{read, read_to_string, write};
-use tempfile::NamedTempFile;
-
-#[derive(Debug)]
-pub enum PathMatchError {
-    PathNotFound,
-    MultiplePathFound,
-}
-
-impl Error for PathMatchError {}
-
-impl Display for PathMatchError {
-    fn fmt(&self, f: &mut Formatter) -> fmt_Result {
-        match self {
-            PathMatchError::PathNotFound => write!(f, "Path not found."),
-            PathMatchError::MultiplePathFound => write!(f, "Multiple paths found."),
-        }
-    }
-}
-
-/// Gets one match path given a path regex.
-///
-/// # Arguments
-///
-/// * `path_regex` - A string which defines the pattern to match.
-///
-/// # Return value
-///
-/// Return a result type. Raise error if path is not found or multiple paths
-/// are found.
-fn match_path(path_regex: &str) -> Result<String> {
-    let mut paths = Vec::<String>::new();
-    for path in glob(path_regex)? {
-        paths.push(path?.display().to_string());
-    }
-    if paths.is_empty() {
-        bail!(PathMatchError::PathNotFound);
-    } else if paths.len() != 1 {
-        bail!(PathMatchError::MultiplePathFound);
-    }
-
-    Ok(paths[0].to_string())
-}
-
-/// Returns the path to pattern [config::UFS_NODE_PATTERN].
-pub fn get_ufs_bsg_node() -> Result<String> {
-    match_path(config::UFS_NODE_PATTERN)
-}
-
-/// Returns the path to pattern [config::SYS_DEV_PATTERN].
-pub fn get_ufs_bsg_sys() -> Result<String> {
-    match_path(config::SYS_DEV_PATTERN)
-}
-
-/// Rescans to bring the device up.
-fn rescan_dev_and_sync(ufs_dev_path: &str) -> Result<()> {
-    let pattern = format!("{}/host*/scsi_host/host*/scan", ufs_dev_path);
-    let scsi_path = match_path(&pattern)?;
-    Command::new("echo")
-        .args(["- - -", ">", &scsi_path])
-        .output()?;
-
-    Command::new("sync").output()?;
-    Ok(())
-}
-
-/// Writes new UFS config using ufs-utils.
-fn write_ufs_config(ufs_bsg_node_path: &str, ufs_config: &config::ConfigDescriptor) -> Result<()> {
-    let config_path = NamedTempFile::new()?.path().display().to_string();
-    let config_bin_vec = bincode_config().big_endian().serialize(&ufs_config)?;
-    write(&config_path, config_bin_vec)?;
-
-    Command::new("ufs-utils")
-        .args([
-            "desc",
-            "-t",
-            "1",
-            "-p",
-            &ufs_bsg_node_path,
-            "-w",
-            &config_path,
-        ])
-        .current_dir(temp_dir())
-        .output()?;
-    Ok(())
-}
-
-/// Gets the current UFS config using ufs-utils.
-fn get_current_config(ufs_bsg_node_path: &str) -> Result<config::ConfigDescriptor> {
-    // Needs to switch to /tmp since ufs-utils dumps the config to cwd
-    // with name `config::UFS_CONFIG_NAME`.
-    Command::new("ufs-utils")
-        .args(["desc", "-t", "1", "-p", ufs_bsg_node_path])
-        .current_dir(temp_dir())
-        .output()
-        .ok();
-
-    let cur_config_path = format!(
-        "{}/{}",
-        temp_dir().display().to_string(),
-        config::UFS_CONFIG_NAME
-    );
-
-    let config_binary = read(cur_config_path)?;
-
-    Ok(bincode_config().big_endian().deserialize(&config_binary)?)
-}
-
-/// Generates a new UFS config given the allocation units.
-fn generate_new_config(alloc_units: u32) -> Result<config::ConfigDescriptor> {
-    let mut unit0 = config::UnitConfigDescriptor::default();
-    unit0.b_lu_enable = 0x01;
-    unit0.d_num_alloc_units = alloc_units;
-    unit0.b_provisioning_type = config::PROVISIONING_TYPE;
-
-    let mut ufs_config = config::ConfigDescriptor::default();
-    ufs_config.units_config[0] = unit0;
-
-    Ok(ufs_config)
-}
-
-/// Reads the geometry_descriptor from a given path.
-fn read_geometry_descriptor(ufs_dev_path: &str, f_name: &str) -> Result<u32> {
-    // The content of the geometry_descriptor is a hex string.
-    // We need to transform it into u32.
-    let path_name = format!("{}/geometry_descriptor/{}", ufs_dev_path, f_name);
-    let raw_string = read_to_string(path_name)?
-        .trim()
-        .trim_start_matches("0x")
-        .to_string();
-
-    Ok(u32::from_str_radix(&raw_string, 16)?)
-}
-
-/// Gets the total number of allocation units from the geometry_descriptor.
-fn get_alloc_units(ufs_dev_path: &str) -> Result<u32> {
-    let total_capacity = read_geometry_descriptor(ufs_dev_path, "raw_device_capacity")?;
-    let segment_size = read_geometry_descriptor(ufs_dev_path, "segment_size")?;
-    let segments_per_allocation_unit =
-        read_geometry_descriptor(ufs_dev_path, "allocation_unit_size")?;
-
-    println!("UFS capacity: {} GB", total_capacity / 0x200000);
-    let units = total_capacity / segment_size / segments_per_allocation_unit;
-
-    Ok(units)
-}
-
-/// Do UFS provisioning if needed.
-///
-/// This function compares the current UFS config and the newly generated
-/// config, and skips provisioning if these two configs are the same.
-pub fn do_provisioning_if_needed(ufs_bsg_node_path: &str, ufs_bsg_sys_path: &str) -> Result<()> {
-    // `ufs_bsg_sys_path` looks like: /sys/devices/pcixxx/xxx/hostx/ufs-bsgx,
-    // but we actually want: /sys/devices/pcixxx/xxx
-    let mut string_vec: Vec<&str> = ufs_bsg_sys_path.split("/").collect();
-    string_vec.truncate(string_vec.len() - 2);
-    let ufs_dev_path = string_vec.join("/");
-    let alloc_units = get_alloc_units(&ufs_dev_path)?;
-    let cur_config = get_current_config(&ufs_bsg_node_path)?;
-    let new_config = generate_new_config(alloc_units)?;
-
-    println!("Current UFS config on DUT: {:#?}", cur_config);
-    println!("New UFS config: {:#?}", new_config);
-
-    if cur_config == new_config {
-        println!("Current UFS config is the same as the new config. Skip UFS provisioning.");
-        return Ok(());
-    }
-
-    write_ufs_config(&ufs_bsg_node_path, &new_config)?;
-    rescan_dev_and_sync(&ufs_dev_path)?;
-    println!("Finish UFS provisioning.");
-
-    Ok(())
-}
diff --git a/rust/src/factory_ufs/ufs_provision.rs b/rust/src/factory_ufs/ufs_provision.rs
new file mode 100644
index 0000000..e4b0e8c
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_provision.rs
@@ -0,0 +1,227 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use crate::factory_ufs::ufs_spec::ufs_common::{self, CommonDescriptorTrait};
+use crate::factory_ufs::ufs_spec::ufs_config_descriptor::ProvisionLun0;
+use crate::utils::process_utils::{Command, StringOutput};
+
+use anyhow::{self, Result};
+use bincode;
+use byteorder::ReadBytesExt;
+use glob;
+use serde::{Deserialize, Serialize};
+use std::env;
+use std::error::Error;
+use std::fmt::{Debug, Display, Formatter, Result as fmt_Result};
+use std::fs;
+use std::io::Cursor;
+use std::path::{Path, PathBuf};
+use tempfile::NamedTempFile;
+
+#[derive(Debug)]
+pub enum PathMatchError {
+    PathNotFound,
+    MultiplePathFound,
+}
+
+impl Error for PathMatchError {}
+
+impl Display for PathMatchError {
+    fn fmt(&self, f: &mut Formatter) -> fmt_Result {
+        match self {
+            PathMatchError::PathNotFound => write!(f, "Path not found."),
+            PathMatchError::MultiplePathFound => write!(f, "Multiple paths found."),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct NoParentDir;
+
+impl Error for NoParentDir {}
+
+impl Display for NoParentDir {
+    fn fmt(&self, f: &mut Formatter) -> fmt_Result {
+        write!(f, "No parent directory can be found.")
+    }
+}
+
+/// Gets one match path given a path regex.
+///
+/// # Arguments
+///
+/// * `path_regex` - A string which defines the pattern to match.
+///
+/// # Return value
+///
+/// Return a result type. Raise error if path is not found or multiple paths
+/// are found.
+fn match_path(path_regex: &str) -> Result<PathBuf> {
+    let mut paths = Vec::<PathBuf>::new();
+    for path in glob::glob(path_regex)? {
+        paths.push(path?);
+    }
+    if paths.is_empty() {
+        anyhow::bail!(PathMatchError::PathNotFound);
+    } else if paths.len() != 1 {
+        anyhow::bail!(PathMatchError::MultiplePathFound);
+    }
+
+    Ok(paths[0].clone())
+}
+
+/// Returns the path to pattern [ufs_common::UFS_NODE_PATTERN].
+pub fn get_ufs_bsg_node() -> Result<PathBuf> {
+    match_path(ufs_common::UFS_NODE_PATTERN)
+}
+
+/// Returns the path to pattern [ufs_common::SYS_DEV_PATTERN].
+pub fn get_ufs_bsg_sys() -> Result<PathBuf> {
+    match_path(ufs_common::SYS_DEV_PATTERN)
+}
+
+/// Rescans to bring the device up.
+fn rescan_dev_and_sync(ufs_dev_path: &Path) -> Result<()> {
+    let pattern = Path::new(ufs_dev_path)
+        .join("host*")
+        .join("scsi_host")
+        .join("host*")
+        .join("scan")
+        .display()
+        .to_string();
+    let scsi_path = match_path(&pattern)?;
+
+    println!("Perform scsi rescan: {}", scsi_path.display());
+    fs::write(&scsi_path, "- - -")?;
+
+    println!("Sync write...");
+    Command::new("sync").output()?;
+
+    Ok(())
+}
+
+/// Writes new UFS config using ufs-utils.
+fn write_ufs_config<T: Serialize>(ufs_bsg_node_path: &Path, ufs_config: &T) -> Result<()> {
+    let config_path = NamedTempFile::new()?.path().display().to_string();
+    let config_bin_vec = bincode::config().big_endian().serialize::<T>(&ufs_config)?;
+    fs::write(&config_path, config_bin_vec)?;
+
+    println!("Write config to {}...", ufs_bsg_node_path.display());
+    let output = Command::new("ufs-utils")
+        .args([
+            "desc",
+            "-t",
+            "1",
+            "-p",
+            &ufs_bsg_node_path.display().to_string(),
+            "-w",
+            &config_path,
+        ])
+        .current_dir(env::temp_dir())
+        .output()?;
+    println!("{}", output.stdout());
+    Ok(())
+}
+
+/// Deserialize a binary config to a given type.
+fn deserialize_config<'de, T: Deserialize<'de>>(config_binary: &'de Vec<u8>) -> Result<T> {
+    Ok(bincode::config().big_endian().deserialize(&config_binary)?)
+}
+
+/// Generates a new UFS config given the allocation units.
+fn generate_new_config<T: Default + ProvisionLun0>(alloc_units: u32) -> Result<T> {
+    let mut ufs_config = T::default();
+    ufs_config.provision_lun_0(alloc_units);
+
+    Ok(ufs_config)
+}
+
+/// Reads the geometry_descriptor from a given path.
+fn read_geometry_descriptor(ufs_dev_path: &Path, f_name: &str) -> Result<u32> {
+    // The content of the geometry_descriptor is a hex string.
+    // We need to transform it into u32.
+    let path = Path::new(ufs_dev_path)
+        .join("geometry_descriptor")
+        .join(f_name);
+    let raw_string = fs::read_to_string(path)?
+        .trim()
+        .trim_start_matches("0x")
+        .to_string();
+
+    Ok(u32::from_str_radix(&raw_string, 16)?)
+}
+
+/// Gets the total number of allocation units from the geometry_descriptor.
+fn get_alloc_units(ufs_dev_path: &Path) -> Result<u32> {
+    let total_capacity = read_geometry_descriptor(ufs_dev_path, "raw_device_capacity")?;
+    let segment_size = read_geometry_descriptor(ufs_dev_path, "segment_size")?;
+    let segments_per_allocation_unit =
+        read_geometry_descriptor(ufs_dev_path, "allocation_unit_size")?;
+
+    println!("UFS capacity: {} GB", total_capacity / 0x200000);
+    let units = total_capacity / segment_size / segments_per_allocation_unit;
+
+    Ok(units)
+}
+
+/// Gets the current UFS config using ufs-utils.
+pub fn get_current_config_binary(ufs_bsg_node_path: &Path) -> Result<Vec<u8>> {
+    // Needs to switch to /tmp since ufs-utils dumps the config to cwd
+    // with name `ufs_common::UFS_CONFIG_NAME`.
+    Command::new("ufs-utils")
+        .args([
+            "desc",
+            "-t",
+            "1",
+            "-p",
+            &ufs_bsg_node_path.display().to_string(),
+        ])
+        .current_dir(env::temp_dir())
+        .output()
+        .ok();
+    let cur_config_path = env::temp_dir().join(ufs_common::UFS_CONFIG_NAME);
+
+    Ok(fs::read(cur_config_path)?)
+}
+
+/// Read the first byte from a config binary.
+pub fn read_header(ufs_config_binary: &Vec<u8>) -> Result<u8> {
+    Ok(Cursor::new(ufs_config_binary).read_u8()?)
+}
+
+/// Do UFS provisioning if needed.
+///
+/// This function compares the current UFS config and the newly generated
+/// config, and skips provisioning if these two configs are the same.
+pub fn do_provisioning_if_needed<'de, T: CommonDescriptorTrait<'de> + ProvisionLun0>(
+    ufs_bsg_node_path: &Path,
+    ufs_bsg_sys_path: &Path,
+    cur_config_binary: &'de Vec<u8>,
+) -> Result<()> {
+    // `ufs_bsg_sys_path` looks like: /sys/devices/pcixxx/xxx/hostx/ufs-bsgx,
+    // but we actually want: /sys/devices/pcixxx/xxx.
+    let ufs_dev_path = Path::new(ufs_bsg_sys_path)
+        .parent()
+        .ok_or(NoParentDir)?
+        .parent()
+        .ok_or(NoParentDir)?;
+    let alloc_units = get_alloc_units(&ufs_dev_path)?;
+
+    let cur_config = deserialize_config::<'de, T>(cur_config_binary)?;
+    let new_config = generate_new_config::<T>(alloc_units)?;
+
+    println!("Current UFS config on DUT: {:#?}", cur_config);
+    println!("New UFS config: {:#?}", new_config);
+
+    if cur_config == new_config {
+        println!("Current UFS config is the same as the new config. Skip UFS provisioning.");
+        return Ok(());
+    }
+
+    write_ufs_config(&ufs_bsg_node_path, &new_config)?;
+    rescan_dev_and_sync(&ufs_dev_path)?;
+    println!("Finish UFS provisioning.");
+
+    Ok(())
+}
diff --git a/rust/src/factory_ufs/ufs_spec/mod.rs b/rust/src/factory_ufs/ufs_spec/mod.rs
new file mode 100644
index 0000000..0e0071a
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_spec/mod.rs
@@ -0,0 +1,9 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+pub mod ufs_2_1;
+pub mod ufs_3_1;
+pub mod ufs_common;
+pub mod ufs_config_descriptor;
+pub mod ufs_unit_descriptor;
diff --git a/rust/src/factory_ufs/ufs_spec/ufs_2_1.rs b/rust/src/factory_ufs/ufs_spec/ufs_2_1.rs
new file mode 100644
index 0000000..748fcfb
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_spec/ufs_2_1.rs
@@ -0,0 +1,102 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use crate::factory_ufs::ufs_spec::ufs_common::{self, CommonDescriptorTrait};
+use crate::factory_ufs::ufs_spec::ufs_unit_descriptor::ProvisionLun;
+use serde::{Deserialize, Serialize};
+
+/// Defines the config descriptor for device.
+///
+/// `DeviceConfigDescriptor` and `UnitConfigDescriptor` specifies the
+/// configuration for UFS provisioning. Both of them need to be serialized and
+/// deserialized in big-endian.
+/// For spec, please refer to `UFS 2.1`.
+#[derive(Serialize, Deserialize, Debug, PartialEq)]
+#[repr(C)]
+pub struct DeviceConfigDescriptor {
+    pub b_length: u8,
+    pub b_descriptor_type: u8,
+    pub b_conf_desc_continue: u8,
+    pub b_boot_enable: u8,
+    pub b_descr_access_en: u8,
+    pub b_init_power_mode: u8,
+    pub b_high_priority_lun: u8,
+    pub b_secure_removal_type: u8,
+    pub b_init_active_icc_level: u8,
+    pub w_periodic_rtc_update: u16,
+    pub b_hpb_control: u8,
+    pub b_rpmb_region_enable: u8,
+    pub b_rpmb_region_1_size: u8,
+    pub b_rpmb_region_2_size: u8,
+    pub b_rpmb_region_3_size: u8,
+}
+
+impl Default for DeviceConfigDescriptor {
+    fn default() -> Self {
+        DeviceConfigDescriptor {
+            b_length: 0x90,
+            b_descriptor_type: 0x01,
+            b_conf_desc_continue: 0x00,
+            b_boot_enable: 0x01,
+            b_descr_access_en: 0x00,
+            b_init_power_mode: 0x01,
+            b_high_priority_lun: 0x7F,
+            b_secure_removal_type: 0x00,
+            b_init_active_icc_level: 0x00,
+            w_periodic_rtc_update: 0x00,
+            b_hpb_control: 0x00,
+            b_rpmb_region_enable: 0x00,
+            b_rpmb_region_1_size: 0x00,
+            b_rpmb_region_2_size: 0x00,
+            b_rpmb_region_3_size: 0x00,
+        }
+    }
+}
+
+impl CommonDescriptorTrait<'_> for DeviceConfigDescriptor {}
+
+/// Defines the config descriptor for each LUN unit.
+///
+/// For spec, please refer to `UFS 2.1`.
+#[derive(Serialize, Deserialize, Debug, PartialEq)]
+#[repr(C)]
+pub struct UnitConfigDescriptor {
+    pub b_lu_enable: u8,
+    pub b_boot_lun_id: u8,
+    pub b_lu_write_protect: u8,
+    pub b_memory_type: u8,
+    pub d_num_alloc_units: u32,
+    pub b_data_reliability: u8,
+    pub b_logical_block_size: u8,
+    pub b_provisioning_type: u8,
+    pub w_context_capabilities: u32,
+    pub w_lu_max_active_hpb_regions: u8,
+}
+
+impl Default for UnitConfigDescriptor {
+    fn default() -> Self {
+        UnitConfigDescriptor {
+            b_lu_enable: 0x0,
+            b_boot_lun_id: 0x0,
+            b_lu_write_protect: 0x0,
+            b_memory_type: 0x0,
+            d_num_alloc_units: 0x0,
+            b_data_reliability: 0x0,
+            b_logical_block_size: ufs_common::LOGICAL_BLOCK_SIZE,
+            b_provisioning_type: 0x0,
+            w_context_capabilities: 0x0,
+            w_lu_max_active_hpb_regions: 0x0,
+        }
+    }
+}
+
+impl ProvisionLun for UnitConfigDescriptor {
+    fn provision_lun(&mut self, alloc_units: u32) {
+        self.b_lu_enable = 0x01;
+        self.d_num_alloc_units = alloc_units;
+        self.b_provisioning_type = ufs_common::PROVISIONING_TYPE;
+    }
+}
+
+impl CommonDescriptorTrait<'_> for UnitConfigDescriptor {}
diff --git a/rust/src/factory_ufs/ufs_spec/ufs_3_1.rs b/rust/src/factory_ufs/ufs_spec/ufs_3_1.rs
new file mode 100644
index 0000000..64fc05f
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_spec/ufs_3_1.rs
@@ -0,0 +1,118 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use crate::factory_ufs::ufs_spec::ufs_common::{self, CommonDescriptorTrait};
+use crate::factory_ufs::ufs_spec::ufs_unit_descriptor::ProvisionLun;
+use serde::{Deserialize, Serialize};
+
+/// Defines the config descriptor for device.
+///
+/// `DeviceConfigDescriptor` and `UnitConfigDescriptor` specifies the
+/// configuration for UFS provisioning. Both of them need to be serialized and
+/// deserialized in big-endian.
+/// For spec, please refer to `UFS 3.1 table 14.10`.
+#[derive(Serialize, Deserialize, Debug, PartialEq)]
+#[repr(C)]
+pub struct DeviceConfigDescriptor {
+    pub b_length: u8,
+    pub b_descriptor_type: u8,
+    pub b_conf_desc_continue: u8,
+    pub b_boot_enable: u8,
+    pub b_descr_access_en: u8,
+    pub b_init_power_mode: u8,
+    pub b_high_priority_lun: u8,
+    pub b_secure_removal_type: u8,
+    pub b_init_active_icc_level: u8,
+    pub w_periodic_rtc_update: u16,
+    pub reserved: u8,
+    pub b_rpmb_region_enable: u8,
+    pub b_rpmb_region_1_size: u8,
+    pub b_rpmb_region_2_size: u8,
+    pub b_rpmb_region_3_size: u8,
+    pub b_write_booster_buffer_preserve_user_space_en: u8,
+    pub b_write_booster_buffer_type: u8,
+    pub d_num_shared_write_booster_buffer_alloc_units: u32,
+}
+
+impl Default for DeviceConfigDescriptor {
+    fn default() -> Self {
+        DeviceConfigDescriptor {
+            b_length: 0xe6,
+            b_descriptor_type: 0x01,
+            b_conf_desc_continue: 0x00,
+            b_boot_enable: 0x01,
+            b_descr_access_en: 0x00,
+            b_init_power_mode: 0x01,
+            b_high_priority_lun: 0x7F,
+            b_secure_removal_type: 0x00,
+            b_init_active_icc_level: 0x00,
+            w_periodic_rtc_update: 0x00,
+            reserved: 0x00,
+            b_rpmb_region_enable: 0x00,
+            b_rpmb_region_1_size: 0x00,
+            b_rpmb_region_2_size: 0x00,
+            b_rpmb_region_3_size: 0x00,
+            b_write_booster_buffer_preserve_user_space_en: 0x00,
+            b_write_booster_buffer_type: 0x00,
+            d_num_shared_write_booster_buffer_alloc_units: 0x00,
+        }
+    }
+}
+
+impl CommonDescriptorTrait<'_> for DeviceConfigDescriptor {}
+
+/// Defines the config descriptor for each LUN unit.
+///
+/// For spec, please refer to `UFS 3.1 table 14.12`.
+#[derive(Serialize, Deserialize, Debug, PartialEq)]
+#[repr(C)]
+pub struct UnitConfigDescriptor {
+    pub b_lu_enable: u8,
+    pub b_boot_lun_id: u8,
+    pub b_lu_write_protect: u8,
+    pub b_memory_type: u8,
+    pub d_num_alloc_units: u32,
+    pub b_data_reliability: u8,
+    pub b_logical_block_size: u8,
+    pub b_provisioning_type: u8,
+    pub w_context_capabilities: u16,
+    pub reserved_empty_1: u16,
+    pub reserved_empty_2: u8,
+    pub reserved_hpb: u16,
+    pub w_hpb_pinned_region_start_idx: u16,
+    pub w_num_hpb_pinned_regions: u16,
+    pub w_lu_max_active_hpb_regions: u32,
+}
+
+impl Default for UnitConfigDescriptor {
+    fn default() -> Self {
+        UnitConfigDescriptor {
+            b_lu_enable: 0x0,
+            b_boot_lun_id: 0x0,
+            b_lu_write_protect: 0x0,
+            b_memory_type: 0x0,
+            d_num_alloc_units: 0x0,
+            b_data_reliability: 0x0,
+            b_logical_block_size: ufs_common::LOGICAL_BLOCK_SIZE,
+            b_provisioning_type: 0x0,
+            w_context_capabilities: 0x0,
+            reserved_empty_1: 0x0,
+            reserved_empty_2: 0x0,
+            reserved_hpb: 0x0,
+            w_hpb_pinned_region_start_idx: 0x0,
+            w_num_hpb_pinned_regions: 0x0,
+            w_lu_max_active_hpb_regions: 0x0,
+        }
+    }
+}
+
+impl ProvisionLun for UnitConfigDescriptor {
+    fn provision_lun(&mut self, alloc_units: u32) {
+        self.b_lu_enable = 0x01;
+        self.d_num_alloc_units = alloc_units;
+        self.b_provisioning_type = ufs_common::PROVISIONING_TYPE;
+    }
+}
+
+impl CommonDescriptorTrait<'_> for UnitConfigDescriptor {}
diff --git a/rust/src/factory_ufs/ufs_spec/ufs_common.rs b/rust/src/factory_ufs/ufs_spec/ufs_common.rs
new file mode 100644
index 0000000..7e98027
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_spec/ufs_common.rs
@@ -0,0 +1,25 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use serde::{Deserialize, Serialize};
+use std::fmt::Debug;
+
+/// Defines the logical block size of each LUN.
+///
+/// The size of logical block size is calculated by 2^(LOGICAL_BLOCK_SIZE),
+/// which equals to 4k here.
+pub const LOGICAL_BLOCK_SIZE: u8 = 0x0C;
+/// Defines the privisioning type for the UFS device.
+///
+/// 0x03 means enabling thin provisioning and setting TPRZ bit to one.
+pub const PROVISIONING_TYPE: u8 = 0x03;
+
+pub const UFS_NODE_PATTERN: &str = "/dev/bsg/ufs-bsg*";
+pub const SYS_DEV_PATTERN: &str = "/sys/devices/*/*/host*/ufs-bsg*";
+pub const UFS_CONFIG_NAME: &str = "config_desc_data_ind_0";
+
+pub trait CommonDescriptorTrait<'de>:
+    Serialize + Deserialize<'de> + Default + Debug + PartialEq
+{
+}
diff --git a/rust/src/factory_ufs/ufs_spec/ufs_config_descriptor.rs b/rust/src/factory_ufs/ufs_spec/ufs_config_descriptor.rs
new file mode 100644
index 0000000..55dc72f
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_spec/ufs_config_descriptor.rs
@@ -0,0 +1,64 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use crate::factory_ufs::ufs_spec::ufs_common::CommonDescriptorTrait;
+use crate::factory_ufs::ufs_spec::{ufs_2_1, ufs_3_1, ufs_unit_descriptor::ProvisionLun};
+use serde::{Deserialize, Serialize};
+
+/// Defines the the config descriptor for UFS provisioning.
+///
+/// This contains 1 DeviceConfigDescriptor and 8 UnitConfigDescriptors.
+#[derive(Serialize, Deserialize, Debug, PartialEq)]
+#[repr(C)]
+pub struct ConfigDescriptor<D, U> {
+    pub device_config: D,
+    pub units_config: [U; 8],
+}
+
+impl<D, U> Default for ConfigDescriptor<D, U>
+where
+    D: Default,
+    U: Default,
+{
+    fn default() -> ConfigDescriptor<D, U> {
+        ConfigDescriptor {
+            device_config: D::default(),
+            units_config: [
+                U::default(),
+                U::default(),
+                U::default(),
+                U::default(),
+                U::default(),
+                U::default(),
+                U::default(),
+                U::default(),
+            ],
+        }
+    }
+}
+
+impl<'de, D, U> CommonDescriptorTrait<'de> for ConfigDescriptor<D, U>
+where
+    D: CommonDescriptorTrait<'de>,
+    U: CommonDescriptorTrait<'de>,
+{
+}
+
+pub trait ProvisionLun0 {
+    fn provision_lun_0(&mut self, alloc_units: u32);
+}
+
+impl<D, U> ProvisionLun0 for ConfigDescriptor<D, U>
+where
+    U: ProvisionLun,
+{
+    fn provision_lun_0(&mut self, alloc_units: u32) {
+        self.units_config[0].provision_lun(alloc_units);
+    }
+}
+
+pub type UFS2ConfigDescriptor =
+    ConfigDescriptor<ufs_2_1::DeviceConfigDescriptor, ufs_2_1::UnitConfigDescriptor>;
+pub type UFS3ConfigDescriptor =
+    ConfigDescriptor<ufs_3_1::DeviceConfigDescriptor, ufs_3_1::UnitConfigDescriptor>;
diff --git a/rust/src/factory_ufs/ufs_spec/ufs_unit_descriptor.rs b/rust/src/factory_ufs/ufs_spec/ufs_unit_descriptor.rs
new file mode 100644
index 0000000..d2f8cb0
--- /dev/null
+++ b/rust/src/factory_ufs/ufs_spec/ufs_unit_descriptor.rs
@@ -0,0 +1,7 @@
+// Copyright 2022 The ChromiumOS Authors.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+pub trait ProvisionLun {
+    fn provision_lun(&mut self, alloc_units: u32);
+}