From 7927ef83b235e1219f7825e740db9ce216147aee Mon Sep 17 00:00:00 2001
From: Micah Lee <micah@micahflee.com>
Date: Fri, 27 Nov 2020 11:22:29 -0800
Subject: [PATCH] Only wait for share mode rendezvous circuits to close, ignore
 the rest

---
 cli/onionshare_cli/__init__.py   |  6 +++---
 cli/onionshare_cli/onion.py      | 15 ++++++++++++---
 cli/onionshare_cli/onionshare.py |  4 ++--
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/cli/onionshare_cli/__init__.py b/cli/onionshare_cli/__init__.py
index dac42514..12808375 100644
--- a/cli/onionshare_cli/__init__.py
+++ b/cli/onionshare_cli/__init__.py
@@ -353,7 +353,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":
@@ -389,9 +389,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 f10401da..b801997d 100644
--- a/cli/onionshare_cli/onion.py
+++ b/cli/onionshare_cli/onion.py
@@ -188,6 +188,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,
@@ -611,7 +614,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.
@@ -691,6 +694,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)
 
@@ -754,9 +761,11 @@ class Onion(object):
                 try:
                     rendevouz_circuit_ids = []
                     for c in self.c.get_circuits():
-                        if c.purpose == "HS_SERVICE_REND":
+                        if (
+                            c.purpose == "HS_SERVICE_REND"
+                            and c.rend_query in self.graceful_close_onions
+                        ):
                             rendevouz_circuit_ids.append(c.id)
-                        # print(f"id={c.id} purpose={c.purpose} rend_query={c.rend_query} type={c.type}")
 
                     while True:
                         num_rend_circuits = 0
diff --git a/cli/onionshare_cli/onionshare.py b/cli/onionshare_cli/onionshare.py
index f74672ce..6fead913 100644
--- a/cli/onionshare_cli/onionshare.py
+++ b/cli/onionshare_cli/onionshare.py
@@ -63,7 +63,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.
         """
@@ -80,7 +80,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"):