Merge pull request #874 from micahflee/866_receive_mode_crash

Stopping the server during a receive mode transfer immediately stops the transfer

9 files changed, 122 insertions, 26 deletions

diff --git a/onionshare/web/receive_mode.py b/onionshare/web/receive_mode.py
index 30414b77..f035271a 100644
--- a/onionshare/web/receive_mode.py
+++ b/onionshare/web/receive_mode.py
@@ -60,7 +60,7 @@ class ReceiveModeWeb(object):
             """
             Upload files.
             """
-            # Make sure the receive mode dir exists
+            # Figure out what the receive mode dir should be
             now = datetime.now()
             date_dir = now.strftime("%Y-%m-%d")
             time_dir = now.strftime("%H.%M.%S")
@@ -134,6 +134,23 @@ class ReceiveModeWeb(object):
                         'dir': receive_mode_dir
                     })
 
+                    # Make sure receive mode dir exists before writing file
+                    valid = True
+                    try:
+                        os.makedirs(receive_mode_dir, 0o700, exist_ok=True)
+                    except PermissionError:
+                        self.web.add_request(self.web.REQUEST_ERROR_DATA_DIR_CANNOT_CREATE, request.path, {
+                            "receive_mode_dir": receive_mode_dir
+                        })
+                        print(strings._('error_cannot_create_data_dir').format(receive_mode_dir))
+                        valid = False
+                    if not valid:
+                        flash('Error uploading, please inform the OnionShare user', 'error')
+                        if self.common.settings.get('public_mode'):
+                            return redirect('/')
+                        else:
+                            return redirect('/{}'.format(slug_candidate))
+
                     self.common.log('ReceiveModeWeb', 'define_routes', '/upload, uploaded {}, saving to {}'.format(f.filename, local_path))
                     print(strings._('receive_mode_received_file').format(local_path))
                     f.save(local_path)
@@ -193,6 +210,7 @@ class ReceiveModeWSGIMiddleware(object):
 
     def __call__(self, environ, start_response):
         environ['web'] = self.web
+        environ['stop_q'] = self.web.stop_q
         return self.app(environ, start_response)
 
 
@@ -201,7 +219,8 @@ class ReceiveModeTemporaryFile(object):
     A custom TemporaryFile that tells ReceiveModeRequest every time data gets
     written to it, in order to track the progress of uploads.
     """
-    def __init__(self, filename, write_func, close_func):
+    def __init__(self, request, filename, write_func, close_func):
+        self.onionshare_request = request
         self.onionshare_filename = filename
         self.onionshare_write_func = write_func
         self.onionshare_close_func = close_func
@@ -222,6 +241,11 @@ class ReceiveModeTemporaryFile(object):
         """
         Custom write method that calls out to onionshare_write_func
         """
+        if not self.onionshare_request.stop_q.empty():
+            self.close()
+            self.onionshare_request.close()
+            return
+
         bytes_written = self.f.write(b)
         self.onionshare_write_func(self.onionshare_filename, bytes_written)
 
@@ -241,6 +265,12 @@ class ReceiveModeRequest(Request):
     def __init__(self, environ, populate_request=True, shallow=False):
         super(ReceiveModeRequest, self).__init__(environ, populate_request, shallow)
         self.web = environ['web']
+        self.stop_q = environ['stop_q']
+
+        self.web.common.log('ReceiveModeRequest', '__init__')
+
+        # Prevent running the close() method more than once
+        self.closed = False
 
         # Is this a valid upload request?
         self.upload_request = False
@@ -301,21 +331,37 @@ class ReceiveModeRequest(Request):
                 'complete': False
             }
 
-        return ReceiveModeTemporaryFile(filename, self.file_write_func, self.file_close_func)
+        return ReceiveModeTemporaryFile(self, filename, self.file_write_func, self.file_close_func)
 
     def close(self):
         """
         Closing the request.
         """
         super(ReceiveModeRequest, self).close()
+
+        # Prevent calling this method more than once per request
+        if self.closed:
+            return
+        self.closed = True
+
+        self.web.common.log('ReceiveModeRequest', 'close')
+
         try:
             if self.told_gui_about_request:
                 upload_id = self.upload_id
-                # Inform the GUI that the upload has finished
-                self.web.add_request(self.web.REQUEST_UPLOAD_FINISHED, self.path, {
-                    'id': upload_id
-                })
+
+                if not self.web.stop_q.empty():
+                    # Inform the GUI that the upload has canceled
+                    self.web.add_request(self.web.REQUEST_UPLOAD_CANCELED, self.path, {
+                        'id': upload_id
+                    })
+                else:
+                    # Inform the GUI that the upload has finished
+                    self.web.add_request(self.web.REQUEST_UPLOAD_FINISHED, self.path, {
+                        'id': upload_id
+                    })
                 self.web.receive_mode.uploads_in_progress.remove(upload_id)
+
         except AttributeError:
             pass
 
@@ -323,6 +369,9 @@ class ReceiveModeRequest(Request):
         """
         This function gets called when a specific file is written to.
         """
+        if self.closed:
+            return
+
         if self.upload_request:
             self.progress[filename]['uploaded_bytes'] += length
 
diff --git a/onionshare/web/share_mode.py b/onionshare/web/share_mode.py
index a57d0a39..eb487c42 100644
--- a/onionshare/web/share_mode.py
+++ b/onionshare/web/share_mode.py
@@ -34,11 +34,6 @@ class ShareModeWeb(object):
         # one download at a time.
         self.download_in_progress = False
 
-        # If the client closes the OnionShare window while a download is in progress,
-        # it should immediately stop serving the file. The client_cancel global is
-        # used to tell the download function that the client is canceling the download.
-        self.client_cancel = False
-
         self.define_routes()
 
     def define_routes(self):
@@ -146,9 +141,6 @@ class ShareModeWeb(object):
             basename = os.path.basename(self.download_filename)
 
             def generate():
-                # The user hasn't canceled the download
-                self.client_cancel = False
-
                 # Starting a new download
                 if not self.web.stay_open:
                     self.download_in_progress = True
@@ -160,7 +152,7 @@ class ShareModeWeb(object):
                 canceled = False
                 while not self.web.done:
                     # The user has canceled the download, so stop serving the file
-                    if self.client_cancel:
+                    if not self.web.stop_q.empty():
                         self.web.add_request(self.web.REQUEST_CANCELED, path, {
                             'id': download_id
                         })
diff --git a/onionshare/web/web.py b/onionshare/web/web.py
index c0e9d6b6..e2b22c4d 100644
--- a/onionshare/web/web.py
+++ b/onionshare/web/web.py
@@ -38,7 +38,8 @@ class Web(object):
     REQUEST_UPLOAD_FILE_RENAMED = 6
     REQUEST_UPLOAD_SET_DIR = 7
     REQUEST_UPLOAD_FINISHED = 8
-    REQUEST_ERROR_DATA_DIR_CANNOT_CREATE = 9
+    REQUEST_UPLOAD_CANCELED = 9
+    REQUEST_ERROR_DATA_DIR_CANNOT_CREATE = 10
 
     def __init__(self, common, is_gui, mode='share'):
         self.common = common
@@ -57,6 +58,11 @@ class Web(object):
         # Are we running in GUI mode?
         self.is_gui = is_gui
 
+        # If the user stops the server while a transfer is in progress, it should
+        # immediately stop the transfer. In order to make it thread-safe, stop_q
+        # is a queue. If anything is in it, then the user stopped the server
+        self.stop_q = queue.Queue()
+
         # Are we using receive mode?
         self.mode = mode
         if self.mode == 'receive':
@@ -224,6 +230,13 @@ class Web(object):
 
         self.stay_open = stay_open
 
+        # Make sure the stop_q is empty when starting a new server
+        while not self.stop_q.empty():
+            try:
+                self.stop_q.get(block=False)
+            except queue.Empty:
+                pass
+
         # In Whonix, listen on 0.0.0.0 instead of 127.0.0.1 (#220)
         if os.path.exists('/usr/share/anon-ws-base-files/workstation'):
             host = '0.0.0.0'
@@ -237,11 +250,10 @@ class Web(object):
         """
         Stop the flask web server by loading /shutdown.
         """
+        self.common.log('Web', 'stop', 'stopping server')
 
-        if self.mode == 'share':
-            # If the user cancels the download, let the download function know to stop
-            # serving the file
-            self.share_mode.client_cancel = True
+        # Let the mode know that the user stopped the server
+        self.stop_q.put(True)
 
         # To stop flask, load http://127.0.0.1:<port>/<shutdown_slug>/shutdown
         if self.running:
diff --git a/onionshare_gui/mode/__init__.py b/onionshare_gui/mode/__init__.py
index edc1777d..4fe335e7 100644
--- a/onionshare_gui/mode/__init__.py
+++ b/onionshare_gui/mode/__init__.py
@@ -329,3 +329,9 @@ class Mode(QtWidgets.QWidget):
         Handle REQUEST_UPLOAD_FINISHED event.
         """
         pass
+
+    def handle_request_upload_canceled(self, event):
+        """
+        Handle REQUEST_UPLOAD_CANCELED event.
+        """
+        pass
diff --git a/onionshare_gui/mode/history.py b/onionshare_gui/mode/history.py
index b54b6f5f..6af804b2 100644
--- a/onionshare_gui/mode/history.py
+++ b/onionshare_gui/mode/history.py
@@ -45,19 +45,32 @@ class HistoryItem(QtWidgets.QWidget):
         When an item finishes, returns a string displaying the start/end datetime range.
         started is a datetime object.
         """
+        return self._get_label_text('gui_all_modes_transfer_finished', 'gui_all_modes_transfer_finished_range', started)
+
+    def get_canceled_label_text(self, started):
+        """
+        When an item is canceled, returns a string displaying the start/end datetime range.
+        started is a datetime object.
+        """
+        return self._get_label_text('gui_all_modes_transfer_canceled', 'gui_all_modes_transfer_canceled_range', started)
+
+    def _get_label_text(self, string_name, string_range_name, started):
+        """
+        Return a string that contains a date, or date range.
+        """
         ended = datetime.now()
         if started.year == ended.year and started.month == ended.month and started.day == ended.day:
             if started.hour == ended.hour and started.minute == ended.minute:
-                text = strings._('gui_all_modes_transfer_finished').format(
+                text = strings._(string_name).format(
                     started.strftime("%b %d, %I:%M%p")
                 )
             else:
-                text = strings._('gui_all_modes_transfer_finished_range').format(
+                text = strings._(string_range_name).format(
                     started.strftime("%b %d, %I:%M%p"),
                     ended.strftime("%I:%M%p")
                 )
         else:
-            text = strings._('gui_all_modes_transfer_finished_range').format(
+            text = strings._(string_range_name).format(
                 started.strftime("%b %d, %I:%M%p"),
                 ended.strftime("%b %d, %I:%M%p")
             )
@@ -306,6 +319,13 @@ class ReceiveHistoryItem(HistoryItem):
             # Change the label
             self.label.setText(self.get_finished_label_text(self.started))
 
+        elif data['action'] == 'canceled':
+            # Hide the progress bar
+            self.progress_bar.hide()
+
+            # Change the label
+            self.label.setText(self.get_canceled_label_text(self.started))
+
 
 class HistoryItemList(QtWidgets.QScrollArea):
     """
diff --git a/onionshare_gui/mode/receive_mode/__init__.py b/onionshare_gui/mode/receive_mode/__init__.py
index 90a1f731..3a90f2f4 100644
--- a/onionshare_gui/mode/receive_mode/__init__.py
+++ b/onionshare_gui/mode/receive_mode/__init__.py
@@ -83,7 +83,7 @@ class ReceiveMode(Mode):
         # Wrapper layout
         self.wrapper_layout = QtWidgets.QHBoxLayout()
         self.wrapper_layout.addLayout(self.main_layout)
-        self.wrapper_layout.addWidget(self.history)
+        self.wrapper_layout.addWidget(self.history, stretch=1)
         self.setLayout(self.wrapper_layout)
 
     def get_stop_server_shutdown_timeout_text(self):
@@ -191,6 +191,18 @@ class ReceiveMode(Mode):
         self.history.update_completed()
         self.history.update_in_progress()
 
+    def handle_request_upload_canceled(self, event):
+        """
+        Handle REQUEST_UPLOAD_CANCELED event.
+        """
+        self.history.update(event["data"]["id"], {
+            'action': 'canceled'
+        })
+        self.history.completed_count += 1
+        self.history.in_progress_count -= 1
+        self.history.update_completed()
+        self.history.update_in_progress()
+
     def on_reload_settings(self):
         """
         We should be ok to re-enable the 'Start Receive Mode' button now.
diff --git a/onionshare_gui/mode/share_mode/__init__.py b/onionshare_gui/mode/share_mode/__init__.py
index d0806cfa..1f5ad00b 100644
--- a/onionshare_gui/mode/share_mode/__init__.py
+++ b/onionshare_gui/mode/share_mode/__init__.py
@@ -115,7 +115,7 @@ class ShareMode(Mode):
         # Wrapper layout
         self.wrapper_layout = QtWidgets.QHBoxLayout()
         self.wrapper_layout.addLayout(self.main_layout)
-        self.wrapper_layout.addWidget(self.history)
+        self.wrapper_layout.addWidget(self.history, stretch=1)
         self.setLayout(self.wrapper_layout)
 
         # Always start with focus on file selection
diff --git a/onionshare_gui/onionshare_gui.py b/onionshare_gui/onionshare_gui.py
index 710a1d84..27abf5e5 100644
--- a/onionshare_gui/onionshare_gui.py
+++ b/onionshare_gui/onionshare_gui.py
@@ -393,6 +393,9 @@ class OnionShareGui(QtWidgets.QMainWindow):
             elif event["type"] == Web.REQUEST_UPLOAD_FINISHED:
                 mode.handle_request_upload_finished(event)
 
+            elif event["type"] == Web.REQUEST_UPLOAD_CANCELED:
+                mode.handle_request_upload_canceled(event)
+
             if event["type"] == Web.REQUEST_ERROR_DATA_DIR_CANNOT_CREATE:
                 Alert(self.common, strings._('error_cannot_create_data_dir').format(event["data"]["receive_mode_dir"]))
 
diff --git a/share/locale/en.json b/share/locale/en.json
index a3307e49..3ad2efda 100644
--- a/share/locale/en.json
+++ b/share/locale/en.json
@@ -171,6 +171,8 @@
     "gui_all_modes_transfer_started": "Started {}",
     "gui_all_modes_transfer_finished_range": "Transferred {} - {}",
     "gui_all_modes_transfer_finished": "Transferred {}",
+    "gui_all_modes_transfer_canceled_range": "Canceled {} - {}",
+    "gui_all_modes_transfer_canceled": "Canceled {}",
     "gui_all_modes_progress_complete": "%p%, {0:s} elapsed.",
     "gui_all_modes_progress_starting": "{0:s}, %p% (calculating)",
     "gui_all_modes_progress_eta": "{0:s}, ETA: {1:s}, %p%",