From 1d80b5fff72cc31cba49f54fe7dd56357842031a Mon Sep 17 00:00:00 2001
From: Piotr Szulc <piotr.szulc@gmail.com>
Date: Sat, 6 Jan 2024 19:41:01 +0100
Subject: [PATCH] [beken-72xx] Free list returned by wlan_sta_scan_result()
 (#226)

* Free list returned by wlan_sta_scan_result()

* scanAlloc improvements

There were a few things I didn't like about this function:
1) realloc() was called a bit too often.
2) if realloc() failed, the previous memory was not freed.
3) scanAlloc returned previous count or 255 on error. But there was no real check for error and 255 could've been used as index to null. I think it's better to simple return boolean.
4) scanAlloc was clearing memory only up to (and excluding) the new entries.

* Corrected clearing new entries in scanAlloc

* scanAlloc() now returns number of allocated items

* Fixed compilation issues related to goto.
---
 .../arduino/libraries/WiFi/WiFiScan.cpp       | 16 ++++++++++---
 .../arduino/libraries/api/WiFi/WiFiScan.cpp   | 24 +++++++++++++------
 .../arduino/libraries/WiFi/WiFiScan.cpp       |  5 +++-
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/cores/beken-72xx/arduino/libraries/WiFi/WiFiScan.cpp b/cores/beken-72xx/arduino/libraries/WiFi/WiFiScan.cpp
index ca585fa33..fe739658d 100644
--- a/cores/beken-72xx/arduino/libraries/WiFi/WiFiScan.cpp
+++ b/cores/beken-72xx/arduino/libraries/WiFi/WiFiScan.cpp
@@ -15,20 +15,27 @@ static void scanHandler(void *ctx, uint8_t param) {
 		return;
 	}
 
+	uint8_t apNum = 0;
 	ScanResult_adv result;
+	result.ApNum  = 0;
+	result.ApList = NULL;
 	if (wlan_sta_scan_result(&result)) {
 		LT_EM(WIFI, "Failed to get scan result");
 		goto end;
 	}
 	LT_IM(WIFI, "Found %d APs", result.ApNum);
 
-	cls->scanAlloc(result.ApNum);
-	if (!scan->ap) {
+	apNum = cls->scanAlloc(result.ApNum);
+	if (0 == apNum) {
 		LT_WM(WIFI, "scan->ap alloc failed");
 		goto end;
 	}
 
-	for (uint8_t i = 0; i < result.ApNum; i++) {
+	if (apNum < result.ApNum) {
+		LT_WM(WIFI, "alloc failed, only %d APs will be copied");
+	}
+
+	for (uint8_t i = 0; i < apNum; i++) {
 		scan->ap[i].ssid	= strdup(result.ApList[i].ssid);
 		scan->ap[i].auth	= securityTypeToAuthMode(result.ApList[i].security);
 		scan->ap[i].rssi	= result.ApList[i].ApPower;
@@ -47,6 +54,9 @@ static void scanHandler(void *ctx, uint8_t param) {
 		scan->running = false;
 		xSemaphoreGive(cDATA->scanSem);
 	}
+	if (result.ApList) {
+		free(result.ApList);
+	}
 	LT_HEAP_I();
 	return;
 }
diff --git a/cores/common/arduino/libraries/api/WiFi/WiFiScan.cpp b/cores/common/arduino/libraries/api/WiFi/WiFiScan.cpp
index 6cd295217..20bb83c7c 100644
--- a/cores/common/arduino/libraries/api/WiFi/WiFiScan.cpp
+++ b/cores/common/arduino/libraries/api/WiFi/WiFiScan.cpp
@@ -39,13 +39,23 @@ void WiFiClass::scanDelete() {
 }
 
 uint8_t WiFiClass::scanAlloc(uint8_t count) {
-	uint8_t last = scan->count;
-	scan->count	 = count;
-	scan->ap	 = (WiFiScanAP *)realloc(scan->ap, count * sizeof(WiFiScanAP));
-	if (!scan->ap)
-		return 255;
-	memset(scan->ap + last, 0, sizeof(WiFiScanAP));
-	return last;
+	if ((!scan->ap) || (count > scan->count)) {
+		auto newMem = (WiFiScanAP *)realloc(scan->ap, count * sizeof(WiFiScanAP));
+		if (!newMem) {
+			return scan->count;
+		}
+		scan->ap = newMem;
+	}
+	if (!scan->ap) {
+		scan->count = 0;
+		return 0;
+	}
+	if (count > scan->count) {
+		// clear only new entries
+		memset(scan->ap + scan->count, 0, sizeof(WiFiScanAP) * (count - scan->count));
+	}
+	scan->count = count;
+	return count;
 }
 
 String WiFiClass::SSID(uint8_t networkItem) {
diff --git a/cores/realtek-amb/arduino/libraries/WiFi/WiFiScan.cpp b/cores/realtek-amb/arduino/libraries/WiFi/WiFiScan.cpp
index 75731ef28..95408815d 100644
--- a/cores/realtek-amb/arduino/libraries/WiFi/WiFiScan.cpp
+++ b/cores/realtek-amb/arduino/libraries/WiFi/WiFiScan.cpp
@@ -20,7 +20,10 @@ static rtw_result_t scanHandler(rtw_scan_handler_result_t *result) {
 	if (!net->SSID.len)
 		return RTW_SUCCESS;
 
-	uint8_t last = cls->scanAlloc(scan->count + 1);
+	uint8_t last = scan->count + 1;
+	if (cls->scanAlloc(last) < last) {
+		return RTW_SUCCESS;
+	}
 
 	scan->ap[last].ssid	   = strdup((char *)net->SSID.val);
 	scan->ap[last].auth	   = securityTypeToAuthMode(net->security);