diff --git a/hardware.cc b/hardware.cc
index 07901f5..b330f64 100644
--- a/hardware.cc
+++ b/hardware.cc
@@ -3,6 +3,7 @@
 // found in the LICENSE file.
 
 #include "update_engine/hardware.h"
+#include "update_engine/utils.h"
 
 #include <base/logging.h>
 #include <rootdev/rootdev.h>
@@ -30,17 +31,31 @@
 
 const string Hardware::KernelDeviceOfBootDevice(
     const std::string& boot_device) {
-  // Currently this assumes the last digit of the boot device is
+  if (boot_device.empty())
+    return boot_device;
+
+  string ubiblock_prefix("/dev/ubiblock");
+  string ret;
+  char partition_num;
+  if(utils::StringHasPrefix(boot_device, ubiblock_prefix)) {
+    // eg: /dev/ubiblock3_0 becomes /dev/mtdblock2
+    ret = "/dev/mtdblock";
+    partition_num = boot_device[ubiblock_prefix.size()];
+  } else {
+    // eg: /dev/sda3 becomes /dev/sda2
+    // eg: /dev/mmcblk0p3 becomes /dev/mmcblk0p2
+    ret = boot_device.substr(0, boot_device.size() - 1);
+    partition_num = boot_device[boot_device.size() - 1];
+  }
+
+  // Currently this assumes the partition number of the boot device is
   // 3, 5, or 7, and changes it to 2, 4, or 6, respectively, to
   // get the kernel device.
-  string ret = boot_device;
-  if (ret.empty())
-    return ret;
-  char last_char = ret[ret.size() - 1];
-  if (last_char == '3' || last_char == '5' || last_char == '7') {
-    ret[ret.size() - 1] = last_char - 1;
+  if (partition_num == '3' || partition_num == '5' || partition_num == '7') {
+    ret.append(1, partition_num - 1);
     return ret;
   }
+
   return "";
 }