UdpContext: check that pbuf_alloc doesn't return nullptr (esp8266#3354)

peterzjin · Sep 22, 2017 · 13c1e8b · 13c1e8b
1 parent 369edb6
commit 13c1e8b
Showing 1 changed file with 27 additions and 6 deletions.
diff --git a/libraries/ESP8266WiFi/src/include/UdpContext.h b/libraries/ESP8266WiFi/src/include/UdpContext.h
@@ -246,6 +246,11 @@ class UdpContext
         {
             _reserve(_tx_buf_offset + size);
         }
+        if (!_tx_buf_head || _tx_buf_head->tot_len < _tx_buf_offset + size)
+        {
+            DEBUGV("failed _reserve");
+            return 0;
+        }
 
         size_t left_to_copy = size;
         while(left_to_copy)
@@ -271,17 +276,25 @@ class UdpContext
     {
         size_t data_size = _tx_buf_offset;
         pbuf* tx_copy = pbuf_alloc(PBUF_TRANSPORT, data_size, PBUF_RAM);
-        uint8_t* dst = reinterpret_cast<uint8_t*>(tx_copy->payload);
-        for (pbuf* p = _tx_buf_head; p; p = p->next) {
-            size_t will_copy = (data_size < p->len) ? data_size : p->len;
-            memcpy(dst, p->payload, will_copy);
-            dst += will_copy;
-            data_size -= will_copy;
+        if(!tx_copy){
+            DEBUGV("failed pbuf_alloc");
+        }
+        else{
+            uint8_t* dst = reinterpret_cast<uint8_t*>(tx_copy->payload);
+            for (pbuf* p = _tx_buf_head; p; p = p->next) {
+                size_t will_copy = (data_size < p->len) ? data_size : p->len;
+                memcpy(dst, p->payload, will_copy);
+                dst += will_copy;
+                data_size -= will_copy;
+            }
         }
         pbuf_free(_tx_buf_head);
         _tx_buf_head = 0;
         _tx_buf_cur = 0;
         _tx_buf_offset = 0;
+        if(!tx_copy){
+            return false;
+        }
 
 
         if (!addr) {
@@ -313,6 +326,10 @@ class UdpContext
         if (!_tx_buf_head)
         {
             _tx_buf_head = pbuf_alloc(PBUF_TRANSPORT, pbuf_unit_size, PBUF_RAM);
+            if (!_tx_buf_head)
+            {
+                return;
+            }
             _tx_buf_cur = _tx_buf_head;
             _tx_buf_offset = 0;
         }
@@ -326,6 +343,10 @@ class UdpContext
         while(grow_size)
         {
             pbuf* pb = pbuf_alloc(PBUF_TRANSPORT, pbuf_unit_size, PBUF_RAM);
+            if (!pb)
+            {
+                return;
+            }
             pbuf_cat(_tx_buf_head, pb);
             if (grow_size < pbuf_unit_size)
                 return;