diff --git a/components/drivers/usb/usbhost/class/udisk.c b/components/drivers/usb/usbhost/class/udisk.c
index 1796450348..deda263b78 100644
--- a/components/drivers/usb/usbhost/class/udisk.c
+++ b/components/drivers/usb/usbhost/class/udisk.c
@@ -13,6 +13,10 @@
 #include <drivers/usb_host.h>
 #include "mass.h"
 
+#define DBG_TAG    "usbhost.udisk"
+#define DBG_LVL    DBG_INFO
+#include <rtdbg.h>
+
 #ifdef RT_USBH_MSTORAGE
 
 #define UDISK_MAX_COUNT        8
@@ -311,6 +315,11 @@ rt_err_t rt_udisk_run(struct uhintf* intf)
         if (ret == RT_EOK)
         {
             struct ustor_data* data = rt_malloc(sizeof(struct ustor_data));
+            if (data == RT_NULL)
+            {
+                LOG_E("Allocate partition data buffer failed.");
+                continue;
+            }
             rt_memset(data, 0, sizeof(struct ustor_data));
             data->intf = intf;
             data->udisk_id = udisk_get_id();
@@ -349,6 +358,11 @@ rt_err_t rt_udisk_run(struct uhintf* intf)
             if(i == 0)
             {
                 struct ustor_data* data = rt_malloc(sizeof(struct ustor_data));
+                if (data == RT_NULL)
+                {
+                    LOG_E("Allocate partition data buffer failed.");
+                    break;
+                }
                 rt_memset(data, 0, sizeof(struct ustor_data));
                 data->udisk_id = udisk_get_id();