From 98c8e7e77cddfa6300680e4a6c7a2b219337b867 Mon Sep 17 00:00:00 2001
From: Dave Allie <dave@daveallie.com>
Date: Sat, 6 Dec 2025 02:49:10 +1100
Subject: [PATCH] Fix memory leak with Epub object getting orphaned

---
 src/main.cpp                        |  4 +++-
 src/screens/EpubReaderScreen.cpp    | 16 +++++++++++++---
 src/screens/EpubReaderScreen.h      |  4 ----
 src/screens/FileSelectionScreen.cpp |  8 ++++++--
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/src/main.cpp b/src/main.cpp
index 37d5eec..c283408 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -51,7 +51,7 @@ Epub* loadEpub(const std::string& path) {
   }
 
   Serial.println("Failed to load epub");
-  free(epub);
+  delete epub;
   return nullptr;
 }
 
@@ -122,6 +122,8 @@ void onSelectEpubFile(const std::string& path) {
   } else {
     exitScreen();
     enterNewScreen(new FullScreenMessageScreen(renderer, "Failed to load epub", REGULAR, false, false));
+    delay(2000);
+    onGoHome();
   }
 }
 
diff --git a/src/screens/EpubReaderScreen.cpp b/src/screens/EpubReaderScreen.cpp
index 50639ce..a6015b7 100644
--- a/src/screens/EpubReaderScreen.cpp
+++ b/src/screens/EpubReaderScreen.cpp
@@ -14,6 +14,10 @@ void EpubReaderScreen::taskTrampoline(void* param) {
 }
 
 void EpubReaderScreen::onEnter() {
+  if (!epub) {
+    return;
+  }
+
   sectionMutex = xSemaphoreCreateMutex();
 
   epub->setupCacheDir();
@@ -42,10 +46,16 @@ void EpubReaderScreen::onEnter() {
 
 void EpubReaderScreen::onExit() {
   xSemaphoreTake(sectionMutex, portMAX_DELAY);
-  vTaskDelete(displayTaskHandle);
+  if (displayTaskHandle) {
+    vTaskDelete(displayTaskHandle);
+    displayTaskHandle = nullptr;
+  }
   vSemaphoreDelete(sectionMutex);
-  displayTaskHandle = nullptr;
   sectionMutex = nullptr;
+  delete section;
+  section = nullptr;
+  delete epub;
+  epub = nullptr;
 }
 
 void EpubReaderScreen::handleInput(const Input input) {
@@ -141,7 +151,7 @@ void EpubReaderScreen::renderPage() {
       section->setupCacheDir();
       if (!section->persistPageDataToSD()) {
         Serial.println("Failed to persist page data to SD");
-        free(section);
+        delete section;
         section = nullptr;
         xSemaphoreGive(sectionMutex);
         return;
diff --git a/src/screens/EpubReaderScreen.h b/src/screens/EpubReaderScreen.h
index ff93f23..d88ad49 100644
--- a/src/screens/EpubReaderScreen.h
+++ b/src/screens/EpubReaderScreen.h
@@ -26,10 +26,6 @@ class EpubReaderScreen final : public Screen {
  public:
   explicit EpubReaderScreen(EpdRenderer* renderer, Epub* epub, const std::function<void()>& onGoHome)
       : Screen(renderer), epub(epub), onGoHome(onGoHome) {}
-  ~EpubReaderScreen() override {
-    free(section);
-    free(epub);
-  }
   void onEnter() override;
   void onExit() override;
   void handleInput(Input input) override;
diff --git a/src/screens/FileSelectionScreen.cpp b/src/screens/FileSelectionScreen.cpp
index 4fb2d98..0869819 100644
--- a/src/screens/FileSelectionScreen.cpp
+++ b/src/screens/FileSelectionScreen.cpp
@@ -32,6 +32,7 @@ void FileSelectionScreen::loadFiles() {
 void FileSelectionScreen::onEnter() {
   basepath = "/";
   loadFiles();
+  selectorIndex = 0;
 
   // Trigger first update
   updateRequired = true;
@@ -45,8 +46,11 @@ void FileSelectionScreen::onEnter() {
 }
 
 void FileSelectionScreen::onExit() {
-  vTaskDelete(displayTaskHandle);
-  displayTaskHandle = nullptr;
+  if (displayTaskHandle) {
+    vTaskDelete(displayTaskHandle);
+    displayTaskHandle = nullptr;
+  }
+  files.clear();
 }
 
 void FileSelectionScreen::handleInput(const Input input) {