From 7e8bfe314d6e441c27ab90c3c1ba73c4de857bb6 Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Mon, 25 Apr 2016 17:13:40 -0700
Subject: [PATCH] Call w.close() exactly once, in both success and error cases.

One downside is that we keep the wormhole channel allocated longer (we
have to finish the file transfer before we can deallocate it, which
could take a while for large files). Maybe we can fix this in the
future.
---
 src/wormhole/cli/cmd_receive.py | 25 +++++++++----------------
 src/wormhole/cli/cmd_send.py    | 12 ++++++++++--
 2 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/src/wormhole/cli/cmd_receive.py b/src/wormhole/cli/cmd_receive.py
index 4e80567..cbe3037 100644
--- a/src/wormhole/cli/cmd_receive.py
+++ b/src/wormhole/cli/cmd_receive.py
@@ -1,7 +1,7 @@
 from __future__ import print_function
 import os, sys, json, binascii, six, tempfile, zipfile
 from tqdm import tqdm
-from twisted.internet import reactor, defer
+from twisted.internet import reactor
 from twisted.internet.defer import inlineCallbacks, returnValue
 from ..twisted.transcribe import Wormhole
 from ..twisted.transit import TransitReceiver
@@ -33,9 +33,8 @@ class TwistedReceiver:
     def msg(self, *args, **kwargs):
         print(*args, file=self.args.stdout, **kwargs)
 
-    # TODO: @handle_server_error
+    @inlineCallbacks
     def go(self):
-        d = defer.succeed(None)
         tor_manager = None
         if self.args.tor:
             _start = self.args.timing.add_event("import TorManager")
@@ -46,18 +45,10 @@ class TwistedReceiver:
             # tor in parallel with everything else, make sure the TorManager
             # can lazy-provide an endpoint, and overlap the startup process
             # with the user handing off the wormhole code
-            d.addCallback(lambda _: tor_manager.start())
-        def _make_wormhole(_):
-            self._w = Wormhole(APPID, self.args.relay_url, tor_manager,
-                               timing=self.args.timing,
-                               reactor=self._reactor)
-        d.addCallback(_make_wormhole)
-        d.addCallback(lambda _: self._go(self._w, tor_manager))
-        def _always_close(res):
-            d2 = self._w.close()
-            d2.addBoth(lambda _: res)
-            return d2
-        d.addBoth(_always_close)
+            yield tor_manager.start()
+        w = Wormhole(APPID, self.args.relay_url, tor_manager,
+                     timing=self.args.timing,
+                     reactor=self._reactor)
         # I wanted to do this instead:
         #
         #    try:
@@ -68,7 +59,9 @@ class TwistedReceiver:
         # but when _go had a UsageError, the stacktrace was always displayed
         # as coming from the "yield self._go" line, which wasn't very useful
         # for tracking it down.
-        return d
+        d = self._go(w, tor_manager)
+        d.addBoth(w.close)
+        yield d
 
     @inlineCallbacks
     def _go(self, w, tor_manager):
diff --git a/src/wormhole/cli/cmd_send.py b/src/wormhole/cli/cmd_send.py
index d5af074..ce58372 100644
--- a/src/wormhole/cli/cmd_send.py
+++ b/src/wormhole/cli/cmd_send.py
@@ -51,6 +51,13 @@ def send(args, reactor=reactor):
     w = Wormhole(APPID, args.relay_url, tor_manager, timing=args.timing,
                  reactor=reactor)
 
+    d = _send(reactor, w, args, phase1, fd_to_send, tor_manager)
+    d.addBoth(w.close)
+    yield d
+
+@inlineCallbacks
+def _send(reactor, w, args, phase1, fd_to_send, tor_manager):
+    transit_sender = None
     if fd_to_send:
         transit_sender = TransitSender(args.transit_helper,
                                        no_listen=args.no_listen,
@@ -103,7 +110,6 @@ def send(args, reactor=reactor):
     if fd_to_send is None:
         if them_phase1["message_ack"] == "ok":
             print(u"text message sent", file=args.stdout)
-            yield w.close()
             returnValue(None) # terminates this function
         raise TransferError("error sending text: %r" % (them_phase1,))
 
@@ -114,7 +120,9 @@ def send(args, reactor=reactor):
         raise TransferError("ambiguous response from remote, "
                             "transfer abandoned: %s" % (them_phase1,))
     tdata = them_phase1["transit"]
-    yield w.close()
+    # XXX the downside of closing above, rather than here, is that it leaves
+    # the channel claimed for a longer time
+    #yield w.close()
     yield _send_file_twisted(tdata, transit_sender, fd_to_send,
                              args.stdout, args.hide_progress, args.timing)
     returnValue(None)