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 "";
 }
 
diff --git a/hardware_unittest.cc b/hardware_unittest.cc
index 3e49fd5..2b52b91 100644
--- a/hardware_unittest.cc
+++ b/hardware_unittest.cc
@@ -44,6 +44,14 @@
 
   EXPECT_EQ("/dev/ubi2", hwut_.KernelDeviceOfBootDevice("/dev/ubi3"));
   EXPECT_EQ("", hwut_.KernelDeviceOfBootDevice("/dev/ubi4"));
+
+  EXPECT_EQ("/dev/mtdblock2",
+            hwut_.KernelDeviceOfBootDevice("/dev/ubiblock3_0"));
+  EXPECT_EQ("/dev/mtdblock4",
+            hwut_.KernelDeviceOfBootDevice("/dev/ubiblock5_0"));
+  EXPECT_EQ("/dev/mtdblock6",
+            hwut_.KernelDeviceOfBootDevice("/dev/ubiblock7_0"));
+  EXPECT_EQ("", hwut_.KernelDeviceOfBootDevice("/dev/ubiblock4_0"));
 }
 
 }  // namespace chromeos_update_engine
diff --git a/utils.cc b/utils.cc
index e3579bf..05e5758 100644
--- a/utils.cc
+++ b/utils.cc
@@ -1004,12 +1004,23 @@
 
 bool GetInstallDev(const std::string& boot_dev, std::string* install_dev) {
   TEST_AND_RETURN_FALSE(StringHasPrefix(boot_dev, "/dev/"));
-  string ret(boot_dev);
-  string::reverse_iterator it = ret.rbegin();  // last character in string
+  string::iterator it;
+  string ubiblock_prefix("/dev/ubiblock");
+
+  install_dev->assign(boot_dev);
+
+  if(StringHasPrefix(boot_dev, ubiblock_prefix)) {
+    // UBI-based device
+    it = install_dev->begin() + ubiblock_prefix.length();
+  } else {
+    // non-UBI device
+    it = install_dev->end() - 1;  // last character in string
+  }
+
   // Right now, we just switch '3' and '5' partition numbers.
-  TEST_AND_RETURN_FALSE((*it == '3') || (*it == '5'));
-  *it = (*it == '3') ? '5' : '3';
-  *install_dev = ret;
+  TEST_AND_RETURN_FALSE(*it == '3' || *it == '5');
+  *it = (*it == '3' ? '5' : '3');
+
   return true;
 }
 
diff --git a/utils_unittest.cc b/utils_unittest.cc
index 09b6667..5098912 100644
--- a/utils_unittest.cc
+++ b/utils_unittest.cc
@@ -283,6 +283,17 @@
 
   boot_dev = "/dev/sda12";
   EXPECT_FALSE(utils::GetInstallDev(boot_dev, &install_dev));
+
+  boot_dev = "/dev/ubiblock3_0";
+  EXPECT_TRUE(utils::GetInstallDev(boot_dev, &install_dev));
+  EXPECT_EQ(install_dev, "/dev/ubiblock5_0");
+
+  boot_dev = "/dev/ubiblock5_0";
+  EXPECT_TRUE(utils::GetInstallDev(boot_dev, &install_dev));
+  EXPECT_EQ(install_dev, "/dev/ubiblock3_0");
+
+  boot_dev = "/dev/ubiblock12_0";
+  EXPECT_FALSE(utils::GetInstallDev(boot_dev, &install_dev));
 }
 
 namespace {