diff --git a/servefile/servefile.py b/servefile/servefile.py index 4363de9..39d2789 100755 --- a/servefile/servefile.py +++ b/servefile/servefile.py @@ -21,6 +21,7 @@ import select import socket from subprocess import Popen, PIPE import sys +import tempfile import time # fix imports for python2/python3 @@ -611,7 +612,24 @@ class FilePutter(BaseHTTPServer.BaseHTTPRequestHandler): # create FieldStorage object for multipart parsing env = os.environ env['REQUEST_METHOD'] = "POST" - fstorage = cgi.FieldStorage(fp=self.rfile, headers=self.headers, environ=env) + targetDir = self.targetDir + + class CustomFieldStorage(cgi.FieldStorage): + + def make_file(self, *args, **kwargs): + """Overwritten to use a named file and the upload directory + + Python 2.7 has an unused "binary" argument while Python 3 does + not have any arguments. Python 2.7 does not have a + self._binary_file attribute. + """ + if sys.version_info.major == 2 or self._binary_file: + return tempfile.NamedTemporaryFile("wb+", dir=targetDir) + else: + return tempfile.NamedTemporaryFile( + "w+", encoding=self.encoding, newline='\n', dir=targetDir) + + fstorage = CustomFieldStorage(fp=self.rfile, headers=self.headers, environ=env) if "file" not in fstorage: self.sendResponse(400, "No file found in request.") return @@ -621,14 +639,21 @@ class FilePutter(BaseHTTPServer.BaseHTTPRequestHandler): self.sendResponse(400, "Filename was empty or invalid") return - # write file down to disk, send a 200 afterwards - target = open(destFileName, "wb") - bytesLeft = length - while bytesLeft > 0: - bytesToRead = min(self.blockSize, bytesLeft) - target.write(fstorage["file"].file.read(bytesToRead)) - bytesLeft -= bytesToRead - target.close() + # put the file at the right place, send 200 afterwards + if getattr(fstorage["file"].file, "name", None): + # the sent file was large, so we can just hard link the temporary + # file and are done + os.link(fstorage["file"].file.name, destFileName) + else: + # write file to disk. it was small enough so no temporary file was + # created + target = open(destFileName, "wb") + bytesLeft = length + while bytesLeft > 0: + bytesToRead = min(self.blockSize, bytesLeft) + target.write(fstorage["file"].file.read(bytesToRead)) + bytesLeft -= bytesToRead + target.close() self.sendResponse(200, "OK! Thanks for uploading") print("Received file '%s' from %s." % (destFileName, self.client_address[0])) diff --git a/tests/test_servefile.py b/tests/test_servefile.py index 8400122..1d7e0f7 100644 --- a/tests/test_servefile.py +++ b/tests/test_servefile.py @@ -357,6 +357,20 @@ def test_upload_size_limit(run_servefile, tmp_path): assert r.status_code == 200 +def test_upload_large_file(run_servefile, tmp_path): + # small files end up in BytesIO while large files get temporary files. this + # test makes sure we hit the large file codepath at least once + uploaddir = tmp_path / 'upload' + run_servefile(['-u', str(uploaddir)]) + + data = "asdf" * 1024 + files = {'file': ('more_data.txt', data)} + r = _retry_while(ConnectionError, make_request)(method='post', files=files) + assert r.status_code == 200 + with open(str(uploaddir / 'more_data.txt')) as f: + assert f.read() == data + + def test_tar_mode(run_servefile, datadir): d = { 'foo': {