Merge pull request #1236 from micahflee/929_download_errors

Prevent incomplete downloads in share mode, close after first download

14 files changed, 141 insertions, 18 deletions

diff --git a/cli/onionshare_cli/__init__.py b/cli/onionshare_cli/__init__.py
index 741db54f..8ba0aac2 100644
--- a/cli/onionshare_cli/__init__.py
+++ b/cli/onionshare_cli/__init__.py
@@ -180,11 +180,11 @@ def main(cwd=None):
     )
     # Share args
     parser.add_argument(
-        "--autostop-sharing",
+        "--no-autostop-sharing",
         action="store_true",
-        dest="autostop_sharing",
-        default=True,
-        help="Share files: Stop sharing after files have been sent",
+        dest="no_autostop_sharing",
+        default=False,
+        help="Share files: Continue sharing after files have been sent (default is to stop sharing)",
     )
     # Receive args
     parser.add_argument(
@@ -233,7 +233,7 @@ def main(cwd=None):
     autostop_timer = int(args.autostop_timer)
     legacy = bool(args.legacy)
     client_auth = bool(args.client_auth)
-    autostop_sharing = bool(args.autostop_sharing)
+    autostop_sharing = not bool(args.no_autostop_sharing)
     data_dir = args.data_dir
     disable_csp = bool(args.disable_csp)
     verbose = bool(args.verbose)
@@ -361,7 +361,7 @@ def main(cwd=None):
                 )
                 sys.exit()
 
-            app.start_onion_service(mode_settings, False, True)
+            app.start_onion_service(mode, mode_settings, False, True)
             url = build_url(mode_settings, app, web)
             schedule = datetime.now() + timedelta(seconds=autostart_timer)
             if mode == "receive":
@@ -397,9 +397,9 @@ def main(cwd=None):
             print("Waiting for the scheduled time before starting...")
             app.onion.cleanup(False)
             time.sleep(autostart_timer)
-            app.start_onion_service(mode_settings)
+            app.start_onion_service(mode, mode_settings)
         else:
-            app.start_onion_service(mode_settings)
+            app.start_onion_service(mode, mode_settings)
     except KeyboardInterrupt:
         print("")
         sys.exit()
diff --git a/cli/onionshare_cli/onion.py b/cli/onionshare_cli/onion.py
index aa057b12..7c9d0a67 100644
--- a/cli/onionshare_cli/onion.py
+++ b/cli/onionshare_cli/onion.py
@@ -169,6 +169,9 @@ class Onion(object):
         # Assigned later if we are using stealth mode
         self.auth_string = None
 
+        # Keep track of onions where it's important to gracefully close to prevent truncated downloads
+        self.graceful_close_onions = []
+
     def connect(
         self,
         custom_settings=None,
@@ -600,7 +603,7 @@ class Onion(object):
         else:
             return False
 
-    def start_onion_service(self, mode_settings, port, await_publication):
+    def start_onion_service(self, mode, mode_settings, port, await_publication):
         """
         Start a onion service on port 80, pointing to the given port, and
         return the onion hostname.
@@ -678,6 +681,10 @@ class Onion(object):
 
         onion_host = res.service_id + ".onion"
 
+        # Gracefully close share mode rendezvous circuits
+        if mode == "share":
+            self.graceful_close_onions.append(res.service_id)
+
         # Save the service_id
         mode_settings.set("general", "service_id", res.service_id)
 
@@ -709,7 +716,7 @@ class Onion(object):
                     "Onion", "stop_onion_service", f"failed to remove {onion_host}"
                 )
 
-    def cleanup(self, stop_tor=True):
+    def cleanup(self, stop_tor=True, wait=True):
         """
         Stop onion services that were created earlier. If there's a tor subprocess running, kill it.
         """
@@ -736,6 +743,46 @@ class Onion(object):
         if stop_tor:
             # Stop tor process
             if self.tor_proc:
+                if wait:
+                    # Wait for Tor rendezvous circuits to close
+                    # Catch exceptions to prevent crash on Ctrl-C
+                    try:
+                        rendevouz_circuit_ids = []
+                        for c in self.c.get_circuits():
+                            if (
+                                c.purpose == "HS_SERVICE_REND"
+                                and c.rend_query in self.graceful_close_onions
+                            ):
+                                rendevouz_circuit_ids.append(c.id)
+
+                        symbols = [c for c in "\\|/-"]
+                        symbols_i = 0
+
+                        while True:
+                            num_rend_circuits = 0
+                            for c in self.c.get_circuits():
+                                if c.id in rendevouz_circuit_ids:
+                                    num_rend_circuits += 1
+
+                            if num_rend_circuits == 0:
+                                print(
+                                    "\rTor rendezvous circuits have closed" + " " * 20
+                                )
+                                break
+
+                            if num_rend_circuits == 1:
+                                circuits = "circuit"
+                            else:
+                                circuits = "circuits"
+                            print(
+                                f"\rWaiting for {num_rend_circuits} Tor rendezvous {circuits} to close {symbols[symbols_i]} ",
+                                end="",
+                            )
+                            symbols_i = (symbols_i + 1) % len(symbols)
+                            time.sleep(1)
+                    except:
+                        pass
+
                 self.tor_proc.terminate()
                 time.sleep(0.2)
                 if self.tor_proc.poll() is None:
diff --git a/cli/onionshare_cli/onionshare.py b/cli/onionshare_cli/onionshare.py
index 513dbc34..db15a03d 100644
--- a/cli/onionshare_cli/onionshare.py
+++ b/cli/onionshare_cli/onionshare.py
@@ -62,7 +62,7 @@ class OnionShare(object):
         except:
             raise OSError("Cannot find an available OnionShare port")
 
-    def start_onion_service(self, mode_settings, await_publication=True):
+    def start_onion_service(self, mode, mode_settings, await_publication=True):
         """
         Start the onionshare onion service.
         """
@@ -79,7 +79,7 @@ class OnionShare(object):
             return
 
         self.onion_host = self.onion.start_onion_service(
-            mode_settings, self.port, await_publication
+            mode, mode_settings, self.port, await_publication
         )
 
         if mode_settings.get("general", "client_auth"):
diff --git a/cli/tests/test_cli.py b/cli/tests/test_cli.py
index 99f26547..be71240b 100644
--- a/cli/tests/test_cli.py
+++ b/cli/tests/test_cli.py
@@ -15,7 +15,7 @@ class MyOnion:
 
     @staticmethod
     def start_onion_service(
-        self, mode_settings_obj, await_publication=True, save_scheduled_key=False
+        self, mode, mode_settings_obj, await_publication=True, save_scheduled_key=False
     ):
         return "test_service_id.onion"
 
@@ -40,13 +40,13 @@ class TestOnionShare:
         assert onionshare_obj.local_only is False
 
     def test_start_onion_service(self, onionshare_obj, mode_settings_obj):
-        onionshare_obj.start_onion_service(mode_settings_obj)
+        onionshare_obj.start_onion_service("share", mode_settings_obj)
         assert 17600 <= onionshare_obj.port <= 17650
         assert onionshare_obj.onion_host == "test_service_id.onion"
 
     def test_start_onion_service_local_only(self, onionshare_obj, mode_settings_obj):
         onionshare_obj.local_only = True
-        onionshare_obj.start_onion_service(mode_settings_obj)
+        onionshare_obj.start_onion_service("share", mode_settings_obj)
         assert onionshare_obj.onion_host == "127.0.0.1:{}".format(onionshare_obj.port)
 
     def test_cleanup(self, onionshare_obj, temp_dir_1024, temp_file_1024):
diff --git a/desktop/src/onionshare/main_window.py b/desktop/src/onionshare/main_window.py
index c1e0cd1a..a1b44032 100644
--- a/desktop/src/onionshare/main_window.py
+++ b/desktop/src/onionshare/main_window.py
@@ -30,6 +30,7 @@ from .widgets import Alert
 from .update_checker import UpdateThread
 from .tab_widget import TabWidget
 from .gui_common import GuiCommon
+from .threads import OnionCleanupThread
 
 
 class MainWindow(QtWidgets.QMainWindow):
@@ -285,8 +286,32 @@ class MainWindow(QtWidgets.QMainWindow):
         e.accept()
 
     def cleanup(self):
+        self.common.log("MainWindow", "cleanup")
         self.tabs.cleanup()
-        self.common.gui.onion.cleanup()
+
+        alert = Alert(
+            self.common,
+            strings._("gui_rendezvous_cleanup"),
+            QtWidgets.QMessageBox.Information,
+            buttons=QtWidgets.QMessageBox.NoButton,
+            autostart=False,
+        )
+        quit_early_button = QtWidgets.QPushButton(
+            strings._("gui_rendezvous_cleanup_quit_early")
+        )
+        alert.addButton(quit_early_button, QtWidgets.QMessageBox.RejectRole)
+
+        self.onion_cleanup_thread = OnionCleanupThread(self.common)
+        self.onion_cleanup_thread.finished.connect(alert.accept)
+        self.onion_cleanup_thread.start()
+
+        alert.exec_()
+        if alert.clickedButton() == quit_early_button:
+            self.common.log("MainWindow", "cleanup", "quitting early")
+            if self.onion_cleanup_thread.isRunning():
+                self.onion_cleanup_thread.terminate()
+                self.onion_cleanup_thread.wait()
+            self.common.gui.onion.cleanup(wait=False)
 
         # Wait 1 second for threads to close gracefully, so tests finally pass
         time.sleep(1)
diff --git a/desktop/src/onionshare/resources/locale/en.json b/desktop/src/onionshare/resources/locale/en.json
index 9617e9c4..502fe13e 100644
--- a/desktop/src/onionshare/resources/locale/en.json
+++ b/desktop/src/onionshare/resources/locale/en.json
@@ -188,5 +188,7 @@
     "settings_error_bundled_tor_not_supported": "Using the Tor version that comes with OnionShare does not work in developer mode on Windows or macOS.",
     "settings_error_bundled_tor_timeout": "Taking too long to connect to Tor. Maybe you aren't connected to the Internet, or have an inaccurate system clock?",
     "settings_error_bundled_tor_broken": "OnionShare could not connect to Tor:\n{}",
+    "gui_rendezvous_cleanup": "Waiting for Tor circuits to close to be sure your files have successfully transferred.\n\nThis might take a few minutes.",
+    "gui_rendezvous_cleanup_quit_early": "Quit Early",
     "error_port_not_available": "OnionShare port not available"
 }
\ No newline at end of file
diff --git a/desktop/src/onionshare/tab/mode/__init__.py b/desktop/src/onionshare/tab/mode/__init__.py
index 7b325704..0bef7628 100644
--- a/desktop/src/onionshare/tab/mode/__init__.py
+++ b/desktop/src/onionshare/tab/mode/__init__.py
@@ -107,6 +107,12 @@ class Mode(QtWidgets.QWidget):
         """
         pass
 
+    def get_type(self):
+        """
+        Returns the type of mode as a string (e.g. "share", "receive", etc.)
+        """
+        pass
+
     def human_friendly_time(self, secs):
         """
         Returns a human-friendly time delta from given seconds.
diff --git a/desktop/src/onionshare/tab/mode/chat_mode/__init__.py b/desktop/src/onionshare/tab/mode/chat_mode/__init__.py
index 25a02969..a7c2929b 100644
--- a/desktop/src/onionshare/tab/mode/chat_mode/__init__.py
+++ b/desktop/src/onionshare/tab/mode/chat_mode/__init__.py
@@ -101,6 +101,12 @@ class ChatMode(Mode):
         self.wrapper_layout.addLayout(self.column_layout)
         self.setLayout(self.wrapper_layout)
 
+    def get_type(self):
+        """
+        Returns the type of mode as a string (e.g. "share", "receive", etc.)
+        """
+        return "chat"
+
     def get_stop_server_autostop_timer_text(self):
         """
         Return the string to put on the stop server button, if there's an auto-stop timer
diff --git a/desktop/src/onionshare/tab/mode/receive_mode/__init__.py b/desktop/src/onionshare/tab/mode/receive_mode/__init__.py
index 95d1ecbe..95a68dcb 100644
--- a/desktop/src/onionshare/tab/mode/receive_mode/__init__.py
+++ b/desktop/src/onionshare/tab/mode/receive_mode/__init__.py
@@ -149,6 +149,12 @@ class ReceiveMode(Mode):
         self.wrapper_layout.addLayout(self.column_layout)
         self.setLayout(self.wrapper_layout)
 
+    def get_type(self):
+        """
+        Returns the type of mode as a string (e.g. "share", "receive", etc.)
+        """
+        return "receive"
+
     def data_dir_button_clicked(self):
         """
         Browse for a new OnionShare data directory, and save to tab settings
diff --git a/desktop/src/onionshare/tab/mode/share_mode/__init__.py b/desktop/src/onionshare/tab/mode/share_mode/__init__.py
index ccf85dbd..bf1498d5 100644
--- a/desktop/src/onionshare/tab/mode/share_mode/__init__.py
+++ b/desktop/src/onionshare/tab/mode/share_mode/__init__.py
@@ -173,6 +173,12 @@ class ShareMode(Mode):
         # Always start with focus on file selection
         self.file_selection.setFocus()
 
+    def get_type(self):
+        """
+        Returns the type of mode as a string (e.g. "share", "receive", etc.)
+        """
+        return "share"
+
     def autostop_sharing_checkbox_clicked(self):
         """
         Save autostop sharing setting to the tab settings
diff --git a/desktop/src/onionshare/tab/mode/website_mode/__init__.py b/desktop/src/onionshare/tab/mode/website_mode/__init__.py
index 325b22f1..6df6ff02 100644
--- a/desktop/src/onionshare/tab/mode/website_mode/__init__.py
+++ b/desktop/src/onionshare/tab/mode/website_mode/__init__.py
@@ -173,6 +173,12 @@ class WebsiteMode(Mode):
         # Always start with focus on file selection
         self.file_selection.setFocus()
 
+    def get_type(self):
+        """
+        Returns the type of mode as a string (e.g. "share", "receive", etc.)
+        """
+        return "website"
+
     def disable_csp_checkbox_clicked(self):
         """
         Save disable CSP setting to the tab settings
diff --git a/desktop/src/onionshare/tab/tab.py b/desktop/src/onionshare/tab/tab.py
index 8f5ffd08..8cbddfed 100644
--- a/desktop/src/onionshare/tab/tab.py
+++ b/desktop/src/onionshare/tab/tab.py
@@ -669,7 +669,9 @@ class Tab(QtWidgets.QWidget):
             return False
 
     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()
         self.app.cleanup()
diff --git a/desktop/src/onionshare/tab_widget.py b/desktop/src/onionshare/tab_widget.py
index 84d16e83..b78e67dd 100644
--- a/desktop/src/onionshare/tab_widget.py
+++ b/desktop/src/onionshare/tab_widget.py
@@ -81,6 +81,8 @@ class TabWidget(QtWidgets.QTabWidget):
         self.event_handler_t.start()
 
     def cleanup(self):
+        self.common.log("TabWidget", "cleanup")
+
         # Stop the event thread
         self.event_handler_t.should_quit = True
         self.event_handler_t.quit()
diff --git a/desktop/src/onionshare/threads.py b/desktop/src/onionshare/threads.py
index 285b51b9..0c31a838 100644
--- a/desktop/src/onionshare/threads.py
+++ b/desktop/src/onionshare/threads.py
@@ -77,7 +77,7 @@ class OnionThread(QtCore.QThread):
         try:
             if self.mode.obtain_onion_early:
                 self.mode.app.start_onion_service(
-                    self.mode.settings, await_publication=False
+                    self.mode.get_type(), self.mode.settings, await_publication=False
                 )
                 # wait for modules in thread to load, preventing a thread-related cx_Freeze crash
                 time.sleep(0.2)
@@ -86,7 +86,7 @@ class OnionThread(QtCore.QThread):
                 self.mode.app.stop_onion_service(self.mode.settings)
             else:
                 self.mode.app.start_onion_service(
-                    self.mode.settings, await_publication=True
+                    self.mode.get_type(), self.mode.settings, await_publication=True
                 )
                 # wait for modules in thread to load, preventing a thread-related cx_Freeze crash
                 time.sleep(0.2)
@@ -258,3 +258,18 @@ class EventHandlerThread(QtCore.QThread):
             if self.should_quit:
                 break
             time.sleep(0.2)
+
+
+class OnionCleanupThread(QtCore.QThread):
+    """
+    Wait for Tor rendezvous circuits to close in a separate thread
+    """
+
+    def __init__(self, common):
+        super(OnionCleanupThread, self).__init__()
+        self.common = common
+        self.common.log("OnionCleanupThread", "__init__")
+
+    def run(self):
+        self.common.log("OnionCleanupThread", "run")
+        self.common.gui.onion.cleanup()