Browse Source

[HUST CSE][components][udisk.c] add two checkout procedures after malloc to make sure allocating successfully (#7475)

kunpHust 2 years ago
parent
commit
7a2bccd1c2
1 changed files with 14 additions and 0 deletions
  1. 14 0
      components/drivers/usb/usbhost/class/udisk.c

+ 14 - 0
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();