diff --git a/src/wormhole/cli/cmd_receive.py b/src/wormhole/cli/cmd_receive.py index ea24a9e..e458de3 100644 --- a/src/wormhole/cli/cmd_receive.py +++ b/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" % - (mode, destname)) + self._msg(u"Error: refusing to overwrite existing '%s'" % destname) raise TransferRejectedError() return abs_destname diff --git a/src/wormhole/test/test_scripts.py b/src/wormhole/test/test_scripts.py index 5627eca..e57f2fb 100644 --- a/src/wormhole/test/test_scripts.py +++ b/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,29 +460,46 @@ 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" - with open(os.path.join(send_dir, send_filename), "w") as f: - f.write(message) - send_cfg.what = receive_filename = send_filename - recv_cfg.what = receive_filename + if mode == "file": + 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) + size = os.stat(fn).st_size - PRESERVE = "don't clobber me\n" - clobberable = os.path.join(receive_dir, receive_filename) - with open(clobberable, "w") as f: - f.write(PRESERVE) + 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_name) + with open(clobberable, "w") as f: + f.write(PRESERVE) send_cfg.cwd = send_dir send_d = cmd_send.send(send_cfg) @@ -488,13 +507,17 @@ 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 - - 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") + # 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") send_stdout = send_cfg.stdout.getvalue() send_stderr = send_cfg.stderr.getvalue() @@ -513,28 +536,74 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase): (receive_stdout, receive_stderr)) # check sender - self.failUnlessIn("Sending {size:s} file named '{name}'{NL}" - .format(size=naturalsize(len(message)), - 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), - send_stdout) - self.failIfIn("File sent.. waiting for confirmation{NL}" - "Confirmation received. Transfer complete.{NL}" - .format(NL=NL), send_stdout) + if mode == "file": + self.failUnlessIn("Sending {size:s} file named '{name}'{NL}" + .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=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: " - "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) - self.failUnless(os.path.exists(fn)) - with open(fn, "r") as f: - self.failUnlessEqual(f.read(), PRESERVE) + if mode == "file": + self.failIfIn("Received file written to ", receive_stdout) + 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): diff --git a/src/wormhole/test/test_util.py b/src/wormhole/test/test_util.py index fb58adc..f15962e 100644 --- a/src/wormhole/test/test_util.py +++ b/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 diff --git a/src/wormhole/util.py b/src/wormhole/util.py index 967243e..b5e39fb 100644 --- a/src/wormhole/util.py +++ b/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