From f5ddd23b70144a388e0ce3e20bb3339e5fd756e9 Mon Sep 17 00:00:00 2001
From: Micah Lee <micah@micahflee.com>
Date: Mon, 22 Sep 2014 20:22:30 +0000
Subject: [PATCH] handles canceled downloads properly (fixes #81)

---
 locale/en.json                   |  1 +
 onionshare/web.py                | 26 ++++++++++++++++++--------
 onionshare_gui/downloads.py      |  3 +++
 onionshare_gui/onionshare_gui.py |  3 +++
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/locale/en.json b/locale/en.json
index b0a0c3b8..87e1eb27 100644
--- a/locale/en.json
+++ b/locale/en.json
@@ -36,6 +36,7 @@
     "gui_stop_server": "Stop Server",
     "gui_copy_url": "Copy URL",
     "gui_downloads": "Downloads:",
+    "gui_canceled": "Canceled",
     "gui_copied_url": "Copied URL to clipboard",
     "gui_starting_server1": "Starting Tor hidden service...",
     "gui_starting_server2": "Crunching files...",
diff --git a/onionshare/web.py b/onionshare/web.py
index 31676264..fabf6f12 100644
--- a/onionshare/web.py
+++ b/onionshare/web.py
@@ -63,6 +63,7 @@ REQUEST_LOAD = 0
 REQUEST_DOWNLOAD = 1
 REQUEST_PROGRESS = 2
 REQUEST_OTHER = 3
+REQUEST_CANCELED = 4
 q = Queue.Queue()
 
 def add_request(type, path, data=None):
@@ -137,25 +138,34 @@ def download(slug_candidate):
 
         fp = open(zip_filename, 'rb')
         done = False
+        canceled = False
         while not done:
             chunk = fp.read(102400)
             if chunk == '':
                 done = True
             else:
-                yield chunk
+                try:
+                    yield chunk
 
-                # tell GUI the progress
-                downloaded_bytes = fp.tell()
-                percent = round((1.0 * downloaded_bytes / zip_filesize) * 100, 2);
-                sys.stdout.write("\r{0}, {1}%          ".format(helpers.human_readable_filesize(downloaded_bytes), percent))
-                sys.stdout.flush()
-                add_request(REQUEST_PROGRESS, path, { 'id':download_id, 'bytes':downloaded_bytes })
+                    # tell GUI the progress
+                    downloaded_bytes = fp.tell()
+                    percent = round((1.0 * downloaded_bytes / zip_filesize) * 100, 2);
+                    sys.stdout.write("\r{0}, {1}%          ".format(helpers.human_readable_filesize(downloaded_bytes), percent))
+                    sys.stdout.flush()
+                    add_request(REQUEST_PROGRESS, path, { 'id':download_id, 'bytes':downloaded_bytes })
+                except:
+                    # looks like the download was canceled
+                    done = True
+                    canceled = True
+
+                    # tell the GUI the download has canceled
+                    add_request(REQUEST_CANCELED, path, { 'id':download_id })
 
         fp.close()
         sys.stdout.write("\n")
 
         # download is finished, close the server
-        if not stay_open:
+        if not stay_open and not canceled:
             print strings._("closing_automatically")
             if shutdown_func is None:
                 raise RuntimeError('Not running with the Werkzeug Server')
diff --git a/onionshare_gui/downloads.py b/onionshare_gui/downloads.py
index 7f8b7c10..90ef6f0d 100644
--- a/onionshare_gui/downloads.py
+++ b/onionshare_gui/downloads.py
@@ -66,3 +66,6 @@ class Downloads(QtGui.QVBoxLayout):
         else:
             pb.setFormat("{0}, %p%".format(helpers.human_readable_filesize(downloaded_bytes)))
 
+    def cancel_download(self, download_id):
+        pb = self.progress_bars[download_id]
+        pb.setFormat(strings._('gui_canceled'))
diff --git a/onionshare_gui/onionshare_gui.py b/onionshare_gui/onionshare_gui.py
index 68727ee4..41f01852 100644
--- a/onionshare_gui/onionshare_gui.py
+++ b/onionshare_gui/onionshare_gui.py
@@ -194,6 +194,9 @@ class OnionShareGui(QtGui.QWidget):
                     if not web.get_stay_open():
                         self.server_status.stop_server()
 
+            elif event["type"] == web.REQUEST_CANCELED:
+                self.downloads.cancel_download(event["data"]["id"])
+
             elif event["path"] != '/favicon.ico':
                 self.status_bar.showMessage('{0}: {1}'.format(strings._('other_page_loaded', True), event["path"]))