From 0f0d0b1b34ee9acbdc200c8673f8040400afa729 Mon Sep 17 00:00:00 2001
From: Max Prokhorov <prokhorov.max@outlook.com>
Date: Mon, 24 Jun 2019 00:15:27 +0300
Subject: [PATCH] Close connection in notFound callback (#1768)

* Close connection in notFound callback

* Add early check of AP mode
---
 code/espurna/web.ino | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/code/espurna/web.ino b/code/espurna/web.ino
index 7c4b6f15..42b8bf6c 100644
--- a/code/espurna/web.ino
+++ b/code/espurna/web.ino
@@ -336,27 +336,62 @@ void _onUpgradeData(AsyncWebServerRequest *request, String filename, size_t inde
     }
 }
 
+bool _onAPModeRequest(AsyncWebServerRequest *request) {
+
+    if ((WiFi.getMode() & WIFI_AP) > 0) {
+        const String domain = getSetting("hostname") + ".";
+        const String host = request->header("Host");
+        const String ip = WiFi.softAPIP().toString();
+
+        // Only allow requests that use our hostname or ip
+        if (host.equals(ip)) return true;
+        if (host.startsWith(domain)) return true;
+
+        // Immediatly close the connection, ref: https://github.com/xoseperez/espurna/issues/1660
+        // Not doing so will cause memory exhaustion, because the connection will linger
+        request->send(404);
+        request->client()->close();
+
+        return false;
+    }
+
+    return true;
+
+}
+
 void _onRequest(AsyncWebServerRequest *request){
 
+    if (!_onAPModeRequest(request)) return;
+
     // Send request to subscribers
     for (unsigned char i = 0; i < _web_request_callbacks.size(); i++) {
         bool response = (_web_request_callbacks[i])(request);
         if (response) return;
     }
 
-    // No subscriber handled the request, return a 404
+    // No subscriber handled the request, return a 404 with implicit "Connection: close"
     request->send(404);
 
+    // And immediatly close the connection, ref: https://github.com/xoseperez/espurna/issues/1660
+    // Not doing so will cause memory exhaustion, because the connection will linger
+    request->client()->close();
+
 }
 
 void _onBody(AsyncWebServerRequest *request, uint8_t *data, size_t len, size_t index, size_t total) {
 
+    if (!_onAPModeRequest(request)) return;
+
     // Send request to subscribers
     for (unsigned char i = 0; i < _web_body_callbacks.size(); i++) {
         bool response = (_web_body_callbacks[i])(request, data, len, index, total);
         if (response) return;
     }
 
+    // Same as _onAPModeRequest(...)
+    request->send(404);
+    request->client()->close();
+
 }