Fix random crashing of ClientContext::write(Stream) and write_P(PGM_P buf, size_t size) (esp8266#2504)

Author: Harald Frostel

libraries/ESP8266WiFi/src/include/DataSource.h

@@ -31,14 +31,14 @@ class BufferDataSource : public DataSource {
 
     const uint8_t* get_buffer(size_t size) override
     {
-        (void) size;
+        (void)size;
         assert(_pos + size <= _size);
         return _data + _pos;
     }
 
     void release_buffer(const uint8_t* buffer, size_t size) override
     {
-        (void) buffer;
+        (void)buffer;
         assert(buffer == _data + _pos);
         _pos += size;
     }
@@ -66,28 +66,48 @@ class BufferedStreamDataSource : public DataSource {
     const uint8_t* get_buffer(size_t size) override
     {
         assert(_pos + size <= _size);
+
+        //Data that was already read from the stream but not released (e.g. if error occured). Otherwise this should be 0.
+        const size_t stream_read = _streamPos - _pos;
+
         if (_bufferSize < size) {
-            _buffer.reset(new uint8_t[size]);
+            uint8_t *new_buffer = new uint8_t[size];
+            //If stream reading is ahead, than some data is already in the old buffer and needs to be copied to new resized buffer
+            if (_buffer && stream_read > 0) {
+                memcpy(new_buffer, _buffer.get(), stream_read);
+            }
+            _buffer.reset(new_buffer);
             _bufferSize = size;
         }
-        size_t cb = _stream.readBytes(reinterpret_cast<char*>(_buffer.get()), size);
-        assert(cb == size);
-        (void) cb;
+
+        //If error in tcp_write in ClientContext::_write_some() occured earlier and therefore release_buffer was not called, than there might not even be data needed to be read from the stream
+        if (size > stream_read) {
+            //Remaining bytes to read from stream
+            const size_t stream_rem = size - stream_read;
+            const size_t cb = _stream.readBytes(reinterpret_cast<char*>(_buffer.get() + stream_read), stream_rem);
+            assert(cb == stream_rem);
+            (void)cb;
+            _streamPos += stream_rem;
+        }
         return _buffer.get();
+
     }
 
     void release_buffer(const uint8_t* buffer, size_t size) override
     {
-        (void) buffer;
+        (void)buffer;
         _pos += size;
+        //Release size needs to be the same as get_buffer(size)
+        assert(_pos == _streamPos);
     }
 
 protected:
-    TStream& _stream;
+    TStream & _stream;
     std::unique_ptr<uint8_t[]> _buffer;
     size_t _size;
     size_t _pos = 0;
     size_t _bufferSize = 0;
+    size_t _streamPos = 0;
 };
 
 class ProgmemStream
@@ -104,7 +124,7 @@ class ProgmemStream
         size_t will_read = (_left < size) ? _left : size;
         memcpy_P((void*)dst, (PGM_VOID_P)_buf, will_read);
         _left -= will_read;
-        _buf  += will_read;
+        _buf += will_read;
         return will_read;
     }