Merge PR91: reject transfers upon insufficient space

src/wormhole/cli/cmd_receive.py

@@ -8,7 +8,8 @@ from twisted.python import log
 from ..wormhole import wormhole
 from ..transit import TransitReceiver
 from ..errors import TransferError, WormholeClosedError
-from ..util import dict_to_bytes, bytes_to_dict, bytes_to_hexstr
+from ..util import (dict_to_bytes, bytes_to_dict, bytes_to_hexstr,
+                    estimate_free_space)
 
 APPID = u"lothar.com/wormhole/text-or-file-xfer"
 
@@ -198,6 +199,11 @@ class TwistedReceiver:
         self.abs_destname = self._decide_destname("file",
                                                   file_data["filename"])
         self.xfersize = file_data["filesize"]
+        free = estimate_free_space(self.abs_destname)
+        if free is not None and free < self.xfersize:
+            self._msg(u"Error: insufficient free space (%sB) for file (%sB)"
+                      % (free, self.xfersize))
+            raise TransferRejectedError()
 
         self._msg(u"Receiving file (%s) into: %s" %
                   (naturalsize(self.xfersize), os.path.basename(self.abs_destname)))
@@ -214,6 +220,11 @@ class TwistedReceiver:
         self.abs_destname = self._decide_destname("directory",
                                                   file_data["dirname"])
         self.xfersize = file_data["zipsize"]
+        free = estimate_free_space(self.abs_destname)
+        if free is not None and free < file_data["numbytes"]:
+            self._msg(u"Error: insufficient free space (%sB) for directory (%sB)"
+                      % (free, file_data["numbytes"]))
+            raise TransferRejectedError()
 
         self._msg(u"Receiving directory (%s) into: %s/" %
                   (naturalsize(self.xfersize), os.path.basename(self.abs_destname)))
@@ -232,8 +243,7 @@ class TwistedReceiver:
 
         # get confirmation from the user before writing to the local directory
         if os.path.exists(abs_destname):
-            self._msg(u"Error: refusing to overwrite existing %s %s" %
+            self._msg(u"Error: refusing to overwrite existing '%s'" % destname)
-                      (mode, destname))
             raise TransferRejectedError()
         return abs_destname
 

src/wormhole/test/test_scripts.py

@@ -449,7 +449,9 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase):
         return self._do_test(mode="directory", override_filename=True)
 
     @inlineCallbacks
-    def test_file_noclobber(self):
+    def _do_test_fail(self, mode, failmode):
+        assert mode in ("file", "directory")
+        assert failmode in ("noclobber", "toobig")
         send_cfg = config("send")
         recv_cfg = config("receive")
 
@@ -458,27 +460,44 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase):
             cfg.relay_url = self.relayurl
             cfg.transit_helper = ""
             cfg.listen = False
-            cfg.code = code = "1-abc"
+            cfg.code = "1-abc"
             cfg.stdout = io.StringIO()
             cfg.stderr = io.StringIO()
 
-        message = "test message"
-
-        recv_cfg.accept_file = True
-
         send_dir = self.mktemp()
         os.mkdir(send_dir)
         receive_dir = self.mktemp()
         os.mkdir(receive_dir)
+        recv_cfg.accept_file = True # don't ask for permission
 
-        send_filename = "testfile"
+        if mode == "file":
-        with open(os.path.join(send_dir, send_filename), "w") as f:
+            message = "test message\n"
+            send_cfg.what = receive_name = send_filename = "testfile"
+            fn = os.path.join(send_dir, send_filename)
+            with open(fn, "w") as f:
                 f.write(message)
-        send_cfg.what = receive_filename = send_filename
+            size = os.stat(fn).st_size
-        recv_cfg.what = receive_filename
 
+        elif mode == "directory":
+            # $send_dir/
+            # $send_dir/$dirname/
+            # $send_dir/$dirname/[12345]
+            # cd $send_dir && wormhole send $dirname
+            # cd $receive_dir && wormhole receive
+            # expect: $receive_dir/$dirname/[12345]
+
+            size = 0
+            send_cfg.what = receive_name = send_dirname = "testdir"
+            os.mkdir(os.path.join(send_dir, send_dirname))
+            for i in range(5):
+                path = os.path.join(send_dir, send_dirname, str(i))
+                with open(path, "w") as f:
+                    f.write("test message %d\n" % i)
+                size += os.stat(path).st_size
+
+        if failmode == "noclobber":
             PRESERVE = "don't clobber me\n"
-        clobberable = os.path.join(receive_dir, receive_filename)
+            clobberable = os.path.join(receive_dir, receive_name)
             with open(clobberable, "w") as f:
                 f.write(PRESERVE)
 
@@ -488,11 +507,15 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase):
         recv_cfg.cwd = receive_dir
         receive_d = cmd_receive.receive(recv_cfg)
 
-        # both sides will fail because of the pre-existing file
+        # both sides will fail
-
+        if failmode == "noclobber":
+            free_space = 10000000
+        else:
+            free_space = 0
+        with mock.patch("wormhole.cli.cmd_receive.estimate_free_space",
+                        return_value=free_space):
             f = yield self.assertFailure(send_d, TransferError)
             self.assertEqual(str(f), "remote error, transfer abandoned: transfer rejected")
-
             f = yield self.assertFailure(receive_d, TransferError)
             self.assertEqual(str(f), "transfer rejected")
 
@@ -513,29 +536,75 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase):
                              (receive_stdout, receive_stderr))
 
         # check sender
+        if mode == "file":
             self.failUnlessIn("Sending {size:s} file named '{name}'{NL}"
-                          .format(size=naturalsize(len(message)),
+                              .format(size=naturalsize(size),
                                       name=send_filename,
                                       NL=NL), send_stdout)
             self.failUnlessIn("On the other computer, please run: "
                               "wormhole receive{NL}"
                               "Wormhole code is: {code}{NL}{NL}"
-                          .format(code=code, NL=NL),
+                              .format(code=send_cfg.code, NL=NL),
                               send_stdout)
             self.failIfIn("File sent.. waiting for confirmation{NL}"
                           "Confirmation received. Transfer complete.{NL}"
                           .format(NL=NL), send_stdout)
+        elif mode == "directory":
+            self.failUnlessIn("Sending directory", send_stdout)
+            self.failUnlessIn("named 'testdir'", send_stdout)
+            self.failUnlessIn("On the other computer, please run: "
+                              "wormhole receive{NL}"
+                              "Wormhole code is: {code}{NL}{NL}"
+                              .format(code=send_cfg.code, NL=NL), send_stdout)
+            self.failIfIn("File sent.. waiting for confirmation{NL}"
+                          "Confirmation received. Transfer complete.{NL}"
+                          .format(NL=NL), send_stdout)
 
         # check receiver
-        self.failUnlessIn("Error: "
+        if mode == "file":
-                          "refusing to overwrite existing file testfile{NL}"
-                          .format(NL=NL), receive_stdout)
             self.failIfIn("Received file written to ", receive_stdout)
-        fn = os.path.join(receive_dir, receive_filename)
+            if failmode == "noclobber":
+                self.failUnlessIn("Error: "
+                                  "refusing to overwrite existing 'testfile'{NL}"
+                                  .format(NL=NL), receive_stdout)
+            else:
+                self.failUnlessIn("Error: "
+                                  "insufficient free space (0B) for file ({size:d}B){NL}"
+                                  .format(NL=NL, size=size), receive_stdout)
+        elif mode == "directory":
+            self.failIfIn("Received files written to {name}"
+                          .format(name=receive_name), receive_stdout)
+            #want = (r"Receiving directory \(\d+ \w+\) into: {name}/"
+            #        .format(name=receive_name))
+            #self.failUnless(re.search(want, receive_stdout),
+            #                (want, receive_stdout))
+            if failmode == "noclobber":
+                self.failUnlessIn("Error: "
+                                  "refusing to overwrite existing 'testdir'{NL}"
+                                  .format(NL=NL), receive_stdout)
+            else:
+                self.failUnlessIn("Error: "
+                                  "insufficient free space (0B) for directory ({size:d}B){NL}"
+                                  .format(NL=NL, size=size), receive_stdout)
+
+        if failmode == "noclobber":
+            fn = os.path.join(receive_dir, receive_name)
             self.failUnless(os.path.exists(fn))
             with open(fn, "r") as f:
                 self.failUnlessEqual(f.read(), PRESERVE)
 
+        # check server stats
+        self._rendezvous.get_stats()
+
+    def test_fail_file_noclobber(self):
+        return self._do_test_fail("file", "noclobber")
+    def test_fail_directory_noclobber(self):
+        return self._do_test_fail("directory", "noclobber")
+    def test_fail_file_toobig(self):
+        return self._do_test_fail("file", "toobig")
+    def test_fail_directory_toobig(self):
+        return self._do_test_fail("directory", "toobig")
+
 class NotWelcome(ServerBase, unittest.TestCase):
     def setUp(self):
         self._setup_relay(error="please upgrade XYZ")

src/wormhole/test/test_util.py

@@ -38,3 +38,11 @@ class Utils(unittest.TestCase):
         d = util.bytes_to_dict(b)
         self.assertIsInstance(d, dict)
         self.assertEqual(d, {"a": "b", "c": 2})
+
+class Space(unittest.TestCase):
+    def test_free_space(self):
+        free = util.estimate_free_space(".")
+        self.assert_(isinstance(free, (int, type(None))), repr(free))
+        # some platforms (I think the VMs used by travis are in this
+        # category) return 0, and windows will return None, so don't assert
+        # anything more specific about the return value

src/wormhole/util.py

@@ -1,5 +1,5 @@
 # No unicode_literals
-import json, unicodedata
+import os, json, unicodedata
 from binascii import hexlify, unhexlify
 
 def to_bytes(u):
@@ -24,3 +24,15 @@ def bytes_to_dict(b):
     d = json.loads(b.decode("utf-8"))
     assert isinstance(d, dict)
     return d
+
+def estimate_free_space(target):
+    # f_bfree is the blocks available to a root user. It might be more
+    # accurate to use f_bavail (blocks available to non-root user), but we
+    # don't know which user is running us, and a lot of installations don't
+    # bother with reserving extra space for root, so let's just stick to the
+    # basic (larger) estimate.
+    try:
+        s = os.statvfs(os.path.dirname(os.path.abspath(target)))
+        return s.f_frsize * s.f_bfree
+    except AttributeError:
+        return None