Merge pull request #1236 from micahflee/929_download_errors

Prevent incomplete downloads in share mode, close after first download

4 files changed, 62 insertions, 15 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):