From b1e1450bfb637caf23506dee464ed3241d370bb6 Mon Sep 17 00:00:00 2001
From: Me No Dev <me-no-dev@vdnsbg.com>
Date: Fri, 8 Apr 2016 17:43:16 +0300
Subject: [PATCH] stability improvements for low memory

---
 src/ESPAsyncWebServer.h |  1 +
 src/WebRequest.cpp      | 17 +++++++++++++----
 src/WebResponseImpl.h   |  3 ++-
 src/WebResponses.cpp    |  6 +-----
 src/WebServer.cpp       | 14 ++++++++++++--
 5 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/ESPAsyncWebServer.h b/src/ESPAsyncWebServer.h
index f5513a8..1e67ecf 100644
--- a/src/ESPAsyncWebServer.h
+++ b/src/ESPAsyncWebServer.h
@@ -267,6 +267,7 @@ class AsyncWebServerResponse {
     virtual bool _started();
     virtual bool _finished();
     virtual bool _failed();
+    virtual bool _sourceValid();
     virtual void _respond(AsyncWebServerRequest *request);
     virtual size_t _ack(AsyncWebServerRequest *request, size_t len, uint32_t time);
 };
diff --git a/src/WebRequest.cpp b/src/WebRequest.cpp
index 37119a5..2cd0707 100644
--- a/src/WebRequest.cpp
+++ b/src/WebRequest.cpp
@@ -181,6 +181,8 @@ void AsyncWebServerRequest::_onDisconnect(){
 }
 
 void AsyncWebServerRequest::_addParam(AsyncWebParameter *p){
+  if(p == NULL)
+    return;
   if(_params == NULL)
     _params = p;
   else {
@@ -252,6 +254,8 @@ bool AsyncWebServerRequest::_parseReqHead(){
 bool AsyncWebServerRequest::_parseReqHeader(){
   if(_temp.indexOf(':')){
     AsyncWebHeader *h = new AsyncWebHeader(_temp);
+    if(h == NULL)
+      return false;
     if(h->name() == "Host"){
       _host = h->value();
       delete h;
@@ -610,14 +614,19 @@ void AsyncWebServerRequest::addInterestingHeader(String name){
   if(!_interestingHeaders->contains(name)) _interestingHeaders->add(name);
 }
 
-
 void AsyncWebServerRequest::send(AsyncWebServerResponse *response){
   _response = response;
-  _response->_respond(this);
+  if(_response == NULL){
+    _client->close(true);
+    _onDisconnect();
+    return;
+  }
+  if(!_response->_sourceValid())
+    send(500);
+  else
+    _response->_respond(this);
 }
 
-
-
 AsyncWebServerResponse * AsyncWebServerRequest::beginResponse(int code, String contentType, String content){
   return new AsyncBasicResponse(code, contentType, content);
 }
diff --git a/src/WebResponseImpl.h b/src/WebResponseImpl.h
index 1f73920..59400f6 100644
--- a/src/WebResponseImpl.h
+++ b/src/WebResponseImpl.h
@@ -28,6 +28,7 @@ class AsyncBasicResponse: public AsyncWebServerResponse {
     AsyncBasicResponse(int code, String contentType=String(), String content=String());
     void _respond(AsyncWebServerRequest *request);
     size_t _ack(AsyncWebServerRequest *request, size_t len, uint32_t time);
+    bool _sourceValid(){ return true; }
 };
 
 class AsyncAbstractResponse: public AsyncWebServerResponse {
@@ -36,8 +37,8 @@ class AsyncAbstractResponse: public AsyncWebServerResponse {
   public:
     void _respond(AsyncWebServerRequest *request);
     size_t _ack(AsyncWebServerRequest *request, size_t len, uint32_t time);
+    bool _sourceValid(){ return false; }
     virtual size_t _fillBuffer(uint8_t *buf, size_t maxLen){ return 0; }
-    virtual bool _sourceValid(){ return false; }
 };
 
 class AsyncFileResponse: public AsyncAbstractResponse {
diff --git a/src/WebResponses.cpp b/src/WebResponses.cpp
index f9e6103..e8380cb 100644
--- a/src/WebResponses.cpp
+++ b/src/WebResponses.cpp
@@ -144,6 +144,7 @@ String AsyncWebServerResponse::_assembleHead(uint8_t version){
 bool AsyncWebServerResponse::_started(){ return _state > RESPONSE_SETUP; }
 bool AsyncWebServerResponse::_finished(){ return _state > RESPONSE_WAIT_ACK; }
 bool AsyncWebServerResponse::_failed(){ return _state == RESPONSE_FAILED; }
+bool AsyncWebServerResponse::_sourceValid(){ return false; }
 void AsyncWebServerResponse::_respond(AsyncWebServerRequest *request){ _state = RESPONSE_END; request->client()->close(); }
 size_t AsyncWebServerResponse::_ack(AsyncWebServerRequest *request, size_t len, uint32_t time){ return 0; }
 
@@ -227,11 +228,6 @@ size_t AsyncBasicResponse::_ack(AsyncWebServerRequest *request, size_t len, uint
  * */
 
 void AsyncAbstractResponse::_respond(AsyncWebServerRequest *request){
-  if(!_sourceValid()){
-    _state = RESPONSE_FAILED;
-    request->send(500);
-    return;
-  }
   _head = _assembleHead(request->version());
   _state = RESPONSE_HEADERS;
   _ack(request, 0, 0);
diff --git a/src/WebServer.cpp b/src/WebServer.cpp
index 42b678d..300027b 100644
--- a/src/WebServer.cpp
+++ b/src/WebServer.cpp
@@ -22,9 +22,19 @@
 #include "WebHandlerImpl.h"
 
 
-AsyncWebServer::AsyncWebServer(uint16_t port):_server(port), _handlers(0), _catchAllHandler(new AsyncCallbackWebHandler()){
+AsyncWebServer::AsyncWebServer(uint16_t port):_server(port), _handlers(0){
+  _catchAllHandler = new AsyncCallbackWebHandler();
+  if(_catchAllHandler == NULL)
+    return;
   _server.onClient([](void *s, AsyncClient* c){
-    new AsyncWebServerRequest((AsyncWebServer*)s, c);
+    if(c == NULL)
+      return;
+    AsyncWebServerRequest *r = new AsyncWebServerRequest((AsyncWebServer*)s, c);
+    if(r == NULL){
+      c->close(true);
+      c->free();
+      delete c;
+    }
   }, this);
 }