From 13c1e8b293fe00f81db83f3001594cd0ab472179 Mon Sep 17 00:00:00 2001
From: Matthew Blythe <mblythester+git@gmail.com>
Date: Fri, 22 Sep 2017 02:56:27 -0600
Subject: [PATCH] UdpContext: check that pbuf_alloc doesn't return nullptr
 (#3354)

---
 .../ESP8266WiFi/src/include/UdpContext.h      | 33 +++++++++++++++----
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/libraries/ESP8266WiFi/src/include/UdpContext.h b/libraries/ESP8266WiFi/src/include/UdpContext.h
index cb5c5628d..3fd224842 100644
--- a/libraries/ESP8266WiFi/src/include/UdpContext.h
+++ b/libraries/ESP8266WiFi/src/include/UdpContext.h
@@ -246,6 +246,11 @@ public:
         {
             _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 @@ public:
     {
         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 @@ private:
         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 @@ private:
         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;