Explicitly cleanup temp files and dirs

6 files changed, 38 insertions, 26 deletions

diff --git a/cli/onionshare_cli/onionshare.py b/cli/onionshare_cli/onionshare.py
index c2711b89..2bb22296 100644
--- a/cli/onionshare_cli/onionshare.py
+++ b/cli/onionshare_cli/onionshare.py
@@ -40,9 +40,6 @@ class OnionShare(object):
         self.onion_host = None
         self.port = None
 
-        # files and dirs to delete on shutdown
-        self.cleanup_filenames = []
-
         # do not use tor -- for development
         self.local_only = local_only
 
@@ -75,7 +72,9 @@ class OnionShare(object):
         if self.local_only:
             self.onion_host = f"127.0.0.1:{self.port}"
             if not mode_settings.get("general", "public"):
-                self.auth_string = "E2GOT5LTUTP3OAMRCRXO4GSH6VKJEUOXZQUC336SRKAHTTT5OVSA"
+                self.auth_string = (
+                    "E2GOT5LTUTP3OAMRCRXO4GSH6VKJEUOXZQUC336SRKAHTTT5OVSA"
+                )
             return
 
         self.onion_host = self.onion.start_onion_service(
diff --git a/cli/onionshare_cli/web/send_base_mode.py b/cli/onionshare_cli/web/send_base_mode.py
index 7b587182..d786b8b5 100644
--- a/cli/onionshare_cli/web/send_base_mode.py
+++ b/cli/onionshare_cli/web/send_base_mode.py
@@ -202,6 +202,9 @@ class SendBaseModeWeb:
 
             file_to_download = self.gzip_individual_files[filesystem_path]
             filesize = os.path.getsize(self.gzip_individual_files[filesystem_path])
+
+            # Cleanup this temp file
+            self.web.cleanup_tempfiles.append(gzip_file)
         else:
             file_to_download = filesystem_path
             filesize = os.path.getsize(filesystem_path)
diff --git a/cli/onionshare_cli/web/share_mode.py b/cli/onionshare_cli/web/share_mode.py
index 75c11a83..d655255e 100644
--- a/cli/onionshare_cli/web/share_mode.py
+++ b/cli/onionshare_cli/web/share_mode.py
@@ -504,10 +504,13 @@ class ShareModeWeb(SendBaseModeWeb):
 
             self.is_zipped = False
 
+            # Cleanup this tempfile
+            self.web.cleanup_tempfiles.append(self.gzip_file)
+
         else:
             # Zip up the files and folders
             self.zip_writer = ZipWriter(
-                self.common, processed_size_callback=processed_size_callback
+                self.common, self.web, processed_size_callback=processed_size_callback
             )
             self.download_filename = self.zip_writer.zip_filename
             for info in self.file_info["files"]:
@@ -538,8 +541,9 @@ class ZipWriter(object):
     filename.
     """
 
-    def __init__(self, common, zip_filename=None, processed_size_callback=None):
+    def __init__(self, common, web, zip_filename=None, processed_size_callback=None):
         self.common = common
+        self.web = web
         self.cancel_compression = False
 
         if zip_filename:
@@ -550,6 +554,9 @@ class ZipWriter(object):
             )
             self.zip_filename = f"{self.zip_temp_dir.name}/onionshare_{self.common.random_string(4, 6)}.zip"
 
+            # Cleanup this temp dir
+            self.web.cleanup_tempdirs.append(self.zip_temp_dir)
+
         self.z = zipfile.ZipFile(self.zip_filename, "w", allowZip64=True)
         self.processed_size_callback = processed_size_callback
         if self.processed_size_callback is None:
diff --git a/cli/onionshare_cli/web/web.py b/cli/onionshare_cli/web/web.py
index e12fccc7..fe2dee87 100644
--- a/cli/onionshare_cli/web/web.py
+++ b/cli/onionshare_cli/web/web.py
@@ -155,7 +155,8 @@ class Web:
             self.socketio.init_app(self.app)
             self.chat_mode = ChatModeWeb(self.common, self)
 
-        self.cleanup_filenames = []
+        self.cleanup_tempfiles = []
+        self.cleanup_tempdirs = []
 
     def get_mode(self):
         if self.mode == "share":
@@ -199,7 +200,10 @@ class Web:
             for header, value in self.security_headers:
                 r.headers.set(header, value)
             # Set a CSP header unless in website mode and the user has disabled it
-            if not self.settings.get("website", "disable_csp") or self.mode != "website":
+            if (
+                not self.settings.get("website", "disable_csp")
+                or self.mode != "website"
+            ):
                 r.headers.set(
                     "Content-Security-Policy",
                     "default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;",
@@ -380,14 +384,13 @@ class Web:
         """
         self.common.log("Web", "cleanup")
 
-        # Cleanup files
-        try:
-            for filename in self.cleanup_filenames:
-                if os.path.isfile(filename):
-                    os.remove(filename)
-                elif os.path.isdir(filename):
-                    shutil.rmtree(filename)
-        except Exception:
-            # Don't crash if file is still in use
-            pass
-        self.cleanup_filenames = []
+        # Close all of the tempfile.NamedTemporaryFile
+        for file in self.cleanup_tempfiles:
+            file.close()
+
+        # Clean up the tempfile.NamedTemporaryDirectory objects
+        for dir in self.cleanup_tempdirs:
+            dir.cleanup()
+
+        self.cleanup_tempfiles = []
+        self.cleanup_tempdirs = []
diff --git a/desktop/src/onionshare/tab/tab.py b/desktop/src/onionshare/tab/tab.py
index 2e592771..b0aa25df 100644
--- a/desktop/src/onionshare/tab/tab.py
+++ b/desktop/src/onionshare/tab/tab.py
@@ -660,9 +660,6 @@ class Tab(QtWidgets.QWidget):
 
         # Close
         if self.close_dialog.clickedButton() == self.close_dialog.accept_button:
-            self.common.log("Tab", "close_tab", "close, closing tab")
-            self.get_mode().stop_server()
-            self.get_mode().web.cleanup()
             return True
         # Cancel
         else:
@@ -671,8 +668,10 @@ class Tab(QtWidgets.QWidget):
 
     def cleanup(self):
         self.common.log("Tab", "cleanup", f"tab_id={self.tab_id}")
-        if self.get_mode() and self.get_mode().web_thread:
-            self.get_mode().web.stop(self.get_mode().app.port)
-            self.get_mode().web_thread.quit()
-            self.get_mode().web_thread.wait()
+        if self.get_mode():
+            if self.get_mode().web_thread:
+                self.get_mode().web.stop(self.get_mode().app.port)
+                self.get_mode().web_thread.quit()
+                self.get_mode().web_thread.wait()
+
             self.get_mode().web.cleanup()
diff --git a/desktop/src/onionshare/tab_widget.py b/desktop/src/onionshare/tab_widget.py
index 7162fcc4..c43a9b9a 100644
--- a/desktop/src/onionshare/tab_widget.py
+++ b/desktop/src/onionshare/tab_widget.py
@@ -316,6 +316,7 @@ class TabWidget(QtWidgets.QTabWidget):
             self.common.log("TabWidget", "closing a service tab")
             if tab.close_tab():
                 self.common.log("TabWidget", "user is okay with closing the tab")
+                tab.cleanup()
 
                 # If the tab is persistent, delete the settings file from disk
                 if tab.settings.get("persistent", "enabled"):