From 50e466b581d6809664c8225659eac3c711947f65 Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Thu, 19 Feb 2015 16:51:59 -0800
Subject: [PATCH] don't complain about unconnectable sockets, improve error
 messages

---
 bin/receive_file.py              | 15 +++++++++++++--
 bin/send_file.py                 |  9 ++++++---
 src/wormhole/blocking/transit.py | 20 +++++++++++++-------
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/bin/receive_file.py b/bin/receive_file.py
index 2ee976f..82441a0 100644
--- a/bin/receive_file.py
+++ b/bin/receive_file.py
@@ -45,9 +45,20 @@ decrypted = SecretBox(xfer_key).decrypt(encrypted)
 # only write to the current directory, and never overwrite anything
 here = os.path.abspath(os.getcwd())
 target = os.path.abspath(os.path.join(here, filename))
-assert os.path.dirname(target) == here
-assert not os.path.exists(target)
+if os.path.dirname(target) != here:
+    print("Error: suggested filename (%s) would be outside current directory"
+          % (filename,))
+    skt.send("bad filename\n")
+    skt.close()
+    sys.exit(1)
+if os.path.exists(target):
+    print("Error: refusing to overwrite existing file %s" % (filename,))
+    skt.send("file already exists\n")
+    skt.close()
+    sys.exit(1)
 with open(target, "wb") as f:
     f.write(decrypted)
 print("Received file written to %s" % filename)
 skt.send("ok\n")
+skt.close()
+sys.exit(0)
diff --git a/bin/send_file.py b/bin/send_file.py
index 11ce780..9d2ed9d 100644
--- a/bin/send_file.py
+++ b/bin/send_file.py
@@ -53,9 +53,12 @@ print("Sending %d bytes.." % filesize)
 skt.send(encrypted)
 
 print("File sent.. waiting for confirmation")
-ack = skt.recv(3)
+# ack is a short newline-terminated string, followed by socket close. A long
+# read is probably good enough.
+ack = skt.recv(300)
 if ack == "ok\n":
     print("Confirmation received. Transfer complete.")
+    sys.exit(0)
 else:
-    print("Transfer failed (remote says: '%r')" % ack)
-skt.close()
+    print("Transfer failed (remote says: %r)" % ack)
+    sys.exit(1)
diff --git a/src/wormhole/blocking/transit.py b/src/wormhole/blocking/transit.py
index b636c58..3b39153 100644
--- a/src/wormhole/blocking/transit.py
+++ b/src/wormhole/blocking/transit.py
@@ -120,10 +120,11 @@ def connector(owner, hint, send_handshake, expected_handshake):
     if isinstance(hint, type(u"")):
         hint = hint.encode("ascii")
     addr,port = hint.split(",")
-    skt = socket.create_connection((addr,port)) # timeout here
-    skt.settimeout(TIMEOUT)
-    #print("socket(%s) connected" % (hint,))
+    skt = None
     try:
+        skt = socket.create_connection((addr,port)) # timeout here
+        skt.settimeout(TIMEOUT)
+        #print("socket(%s) connected" % (hint,))
         skt.send(send_handshake)
         got = b""
         while len(got) < len(expected_handshake):
@@ -132,13 +133,18 @@ def connector(owner, hint, send_handshake, expected_handshake):
                 raise BadHandshake("got '%r' want '%r' on %s" %
                                    (got, expected_handshake, hint))
         #print("connector ready %r" % (hint,))
-    except:
+    except Exception as e:
         try:
-            skt.shutdown(socket.SHUT_WR)
+            if skt:
+                skt.shutdown(socket.SHUT_WR)
         except socket.error:
             pass
-        skt.close()
-        raise
+        if skt:
+            skt.close()
+        # ignore socket errors, warn about coding errors
+        if not isinstance(e, (socket.error, socket.timeout, BadHandshake)):
+            raise
+        return
     # owner is now responsible for the socket
     owner._negotiation_finished(skt) # note thread