diff --git a/hardware.cc b/hardware.cc
index 26435cf..9f3fe00 100644
--- a/hardware.cc
+++ b/hardware.cc
@@ -59,12 +59,14 @@
   CgptAddParams params;
   memset(&params, '\0', sizeof(params));
 
-  string root_dev = utils::RootDevice(kernel_device);
-  string partition_number_str = utils::PartitionNumber(kernel_device);
-  uint32_t partition_number = atoi(partition_number_str.c_str());
+  std::string disk_name;
+  int partition_num = 0;
 
-  params.drive_name = const_cast<char *>(root_dev.c_str());
-  params.partition = partition_number;
+  if (!utils::SplitPartitionName(kernel_device, &disk_name, &partition_num))
+    return false;
+
+  params.drive_name = const_cast<char *>(disk_name.c_str());
+  params.partition = partition_num;
 
   int retval = CgptGetPartitionDetails(&params);
   if (retval != CGPT_OK)
@@ -82,15 +84,17 @@
     return false;
   }
 
-  string root_dev = utils::RootDevice(kernel_device);
-  string partition_number_str = utils::PartitionNumber(kernel_device);
-  uint32_t partition_number = atoi(partition_number_str.c_str());
+  std::string disk_name;
+  int partition_num = 0;
+
+  if (!utils::SplitPartitionName(kernel_device, &disk_name, &partition_num))
+    return false;
 
   CgptAddParams params;
   memset(&params, 0, sizeof(params));
 
-  params.drive_name = const_cast<char *>(root_dev.c_str());
-  params.partition = partition_number;
+  params.drive_name = const_cast<char *>(disk_name.c_str());
+  params.partition = partition_num;
 
   params.successful = false;
   params.set_successful = true;
diff --git a/payload_state.cc b/payload_state.cc
index a23b545..7505b96 100644
--- a/payload_state.cc
+++ b/payload_state.cc
@@ -1166,13 +1166,13 @@
   // payload was marked as ready immediately before the last reboot, and we
   // need to check if such payload successfully rebooted or not.
   if (prefs_->Exists(kPrefsTargetVersionInstalledFrom)) {
-    string installed_from;
-    if (!prefs_->GetString(kPrefsTargetVersionInstalledFrom, &installed_from)) {
+    int64_t installed_from = 0;
+    if (!prefs_->GetInt64(kPrefsTargetVersionInstalledFrom, &installed_from)) {
       LOG(ERROR) << "Error reading TargetVersionInstalledFrom on reboot.";
       return;
     }
-    if (installed_from ==
-        utils::PartitionNumber(system_state_->hardware()->BootDevice())) {
+    if (int(installed_from) ==
+        utils::GetPartitionNumber(system_state_->hardware()->BootDevice())) {
       // A reboot was pending, but the chromebook is again in the same
       // BootDevice where the update was installed from.
       int64_t target_attempt;
@@ -1221,8 +1221,8 @@
   }
   prefs_->SetInt64(kPrefsTargetVersionAttempt, target_attempt + 1);
 
-  prefs_->SetString(kPrefsTargetVersionInstalledFrom,
-                    utils::PartitionNumber(
+  prefs_->SetInt64(kPrefsTargetVersionInstalledFrom,
+                    utils::GetPartitionNumber(
                         system_state_->hardware()->BootDevice()));
 }
 
diff --git a/update_check_scheduler.cc b/update_check_scheduler.cc
index 1835de9..682bae7 100644
--- a/update_check_scheduler.cc
+++ b/update_check_scheduler.cc
@@ -59,7 +59,7 @@
 }
 
 bool UpdateCheckScheduler::IsBootDeviceRemovable() {
-  return utils::IsRemovableDevice(utils::RootDevice(
+  return utils::IsRemovableDevice(utils::GetDiskName(
       system_state_->hardware()->BootDevice()));
 }
 
diff --git a/utils.cc b/utils.cc
index 8a73cd7..7558f35 100644
--- a/utils.cc
+++ b/utils.cc
@@ -368,31 +368,78 @@
   return true;
 }
 
-string RootDevice(const string& partition_device) {
-  FilePath device_path(partition_device);
-  if (device_path.DirName().value() != "/dev") {
-    return "";
-  }
-  string::const_iterator it = --partition_device.end();
-  for (; it >= partition_device.begin(); --it) {
-    if (!isdigit(*it))
-      break;
-  }
-  // Some devices contain a p before the partitions. For example:
-  // /dev/mmc0p4 should be shortened to /dev/mmc0.
-  if (*it == 'p')
-    --it;
-  return string(partition_device.begin(), it + 1);
+std::string GetDiskName(const string& partition_name) {
+  std::string disk_name;
+  return SplitPartitionName(partition_name, &disk_name, nullptr) ?
+         disk_name : std::string();
 }
 
-string PartitionNumber(const string& partition_device) {
-  CHECK(!partition_device.empty());
-  string::const_iterator it = --partition_device.end();
-  for (; it >= partition_device.begin(); --it) {
-    if (!isdigit(*it))
-      break;
+int GetPartitionNumber(const std::string& partition_name) {
+  int partition_num = 0;
+  return SplitPartitionName(partition_name, nullptr, &partition_num) ?
+      partition_num : 0;
+}
+
+bool SplitPartitionName(const std::string& partition_name,
+                        std::string* out_disk_name,
+                        int* out_partition_num) {
+  if (!StringHasPrefix(partition_name, "/dev/")) {
+    LOG(ERROR) << "Invalid partition device name: " << partition_name;
+    return false;
   }
-  return string(it + 1, partition_device.end());
+
+  if (StringHasPrefix(partition_name, "/dev/ubiblock")) {
+    // NAND block devices have weird naming which could be something
+    // like "/dev/ubiblock2_0". Since update engine doesn't have proper
+    // support for NAND devices yet, don't bother parsing their names for now.
+    LOG(ERROR) << "UBI block devices are not supported: " << partition_name;
+    return false;
+  }
+
+  size_t pos = partition_name.find_last_not_of("0123456789");
+  if (pos == string::npos || (pos + 1) == partition_name.size()) {
+    LOG(ERROR) << "Unable to parse partition device name: "
+               << partition_name;
+    return false;
+  }
+
+  if (out_disk_name) {
+    // Special case for MMC devices which have the following naming scheme:
+    // mmcblk0p2
+    bool valid_mmc_device = StringHasPrefix(partition_name, "/dev/mmcblk") &&
+                            partition_name[pos] == 'p';
+    *out_disk_name =
+        partition_name.substr(0, valid_mmc_device ? pos : (pos + 1));
+  }
+
+  if (out_partition_num) {
+    std::string partition_str = partition_name.substr(pos + 1);
+    *out_partition_num = atoi(partition_str.c_str());
+  }
+  return true;
+}
+
+std::string MakePartitionName(const std::string& disk_name,
+                              int partition_num) {
+  if (!StringHasPrefix(disk_name, "/dev/")) {
+    LOG(ERROR) << "Invalid disk name: " << disk_name;
+    return std::string();
+  }
+
+  if (partition_num < 1) {
+    LOG(ERROR) << "Invalid partition number: " << partition_num;
+    return std::string();
+  }
+
+  std::string partition_name = disk_name;
+  if (!StringHasPrefix(disk_name, "/dev/mmcblk")) {
+    // Special case for MMC devices. Add "p" to separate the disk name
+    // from partition number
+    partition_name += 'p';
+  }
+
+  base::StringAppendF(&partition_name, "%d", partition_num);
+  return partition_name;
 }
 
 string SysfsBlockDevice(const string& device) {
diff --git a/utils.h b/utils.h
index 3e1fd16..a84f512 100644
--- a/utils.h
+++ b/utils.h
@@ -136,14 +136,35 @@
 // This WILL cross filesystem boundaries.
 bool RecursiveUnlinkDir(const std::string& path);
 
-// Returns the root device for a partition. For example,
-// RootDevice("/dev/sda3") returns "/dev/sda". Returns an empty string
-// if the input device is not of the "/dev/xyz" form.
-std::string RootDevice(const std::string& partition_device);
+// Returns the disk device name for a partition. For example,
+// GetDiskName("/dev/sda3") returns "/dev/sda". Returns an empty string
+// if the input device is not of the "/dev/xyz#" form.
+std::string GetDiskName(const std::string& partition_name);
 
-// Returns the partition number, as a string, of partition_device. For example,
-// PartitionNumber("/dev/sda3") returns "3".
-std::string PartitionNumber(const std::string& partition_device);
+// Returns the partition number, of partition device name. For example,
+// GetPartitionNumber("/dev/sda3") returns 3.
+// Returns 0 on failure
+int GetPartitionNumber(const std::string& partition_name);
+
+// Splits the partition device name into the block device name and partition
+// number. For example, "/dev/sda3" will be split into {"/dev/sda", 3} and
+// "/dev/mmcblk0p2" into {"/dev/mmcblk0", 2}
+// Returns false when malformed device name is passed in.
+// If both output parameters are omitted (nullptr), can be used
+// just to test the validity of the device name. Note that the function
+// simply checks if the device name looks like a valid device, no other
+// checks are performed (i.e. it doesn't check if the device actually exists).
+bool SplitPartitionName(const std::string& partition_name,
+                        std::string* out_disk_name,
+                        int* out_partition_num);
+
+// Builds a partition device name from the block device name and partition
+// number. For example:
+// {"/dev/sda", 1} => "/dev/sda1"
+// {"/dev/mmcblk2", 12} => "/dev/mmcblk2p12"
+// Returns empty string when invalid parameters are passed in
+std::string MakePartitionName(const std::string& disk_name,
+                              int partition_num);
 
 // Returns the sysfs block device for a root block device. For
 // example, SysfsBlockDevice("/dev/sda") returns
diff --git a/utils_unittest.cc b/utils_unittest.cc
index 8602337..d75a6b0 100644
--- a/utils_unittest.cc
+++ b/utils_unittest.cc
@@ -202,12 +202,16 @@
   EXPECT_FALSE(utils::StringHasSuffix(result, "XXXXXX"));
 }
 
-TEST(UtilsTest, RootDeviceTest) {
-  EXPECT_EQ("/dev/sda", utils::RootDevice("/dev/sda3"));
-  EXPECT_EQ("/dev/mmc0", utils::RootDevice("/dev/mmc0p3"));
-  EXPECT_EQ("", utils::RootDevice("/dev/foo/bar"));
-  EXPECT_EQ("", utils::RootDevice("/"));
-  EXPECT_EQ("", utils::RootDevice(""));
+TEST(UtilsTest, GetDiskNameTest) {
+  EXPECT_EQ("/dev/sda", utils::GetDiskName("/dev/sda3"));
+  EXPECT_EQ("/dev/sda", utils::GetDiskName("/dev/sda1234"));
+  EXPECT_EQ("/dev/mmcblk0", utils::GetDiskName("/dev/mmcblk0p3"));
+  EXPECT_EQ("", utils::GetDiskName("/dev/mmcblk0p"));
+  EXPECT_EQ("", utils::GetDiskName("/dev/sda"));
+  EXPECT_EQ("", utils::GetDiskName("/dev/ubiblock3_2"));
+  EXPECT_EQ("", utils::GetDiskName("/dev/foo/bar"));
+  EXPECT_EQ("", utils::GetDiskName("/"));
+  EXPECT_EQ("", utils::GetDiskName(""));
 }
 
 TEST(UtilsTest, SysfsBlockDeviceTest) {
@@ -224,9 +228,16 @@
   EXPECT_FALSE(utils::IsRemovableDevice("/dev/non-existent-device"));
 }
 
-TEST(UtilsTest, PartitionNumberTest) {
-  EXPECT_EQ("3", utils::PartitionNumber("/dev/sda3"));
-  EXPECT_EQ("3", utils::PartitionNumber("/dev/mmc0p3"));
+TEST(UtilsTest, GetPartitionNumberTest) {
+  EXPECT_EQ(3, utils::GetPartitionNumber("/dev/sda3"));
+  EXPECT_EQ(123, utils::GetPartitionNumber("/dev/sda123"));
+  EXPECT_EQ(2, utils::GetPartitionNumber("/dev/mmcblk0p2"));
+  EXPECT_EQ(0, utils::GetPartitionNumber("/dev/mmcblk0p"));
+  EXPECT_EQ(0, utils::GetPartitionNumber("/dev/ubiblock3_2"));
+  EXPECT_EQ(0, utils::GetPartitionNumber(""));
+  EXPECT_EQ(0, utils::GetPartitionNumber("/"));
+  EXPECT_EQ(0, utils::GetPartitionNumber("/dev/"));
+  EXPECT_EQ(0, utils::GetPartitionNumber("/dev/sda"));
 }
 
 TEST(UtilsTest, CompareCpuSharesTest) {