From 8e43033ed71e2e9658c7587e61be5659edf74ae9 Mon Sep 17 00:00:00 2001 From: Mathieu Pillard Date: Mon, 12 Apr 2021 21:27:42 +0200 Subject: [PATCH] Fix repack_fileupload() to clean up temporary directories once it's done (#16900) * Fix repack_fileupload() to clean up temporary directories once it's done Also remove unnecessary timers * Add back timers --- src/olympia/files/tasks.py | 36 +++++++++++++++------------ src/olympia/files/tests/test_tasks.py | 5 ++++ 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/src/olympia/files/tasks.py b/src/olympia/files/tasks.py index e30741797c..2ca9077846 100644 --- a/src/olympia/files/tasks.py +++ b/src/olympia/files/tasks.py @@ -28,22 +28,26 @@ def repack_fileupload(results, upload_pk): if upload.path.endswith('.xpi'): timer = StopWatch('files.tasks.repack_fileupload.') timer.start() - try: - tempdir = tempfile.mkdtemp() # *not* on TMP_PATH, we want local fs - extract_zip(upload.path, tempdir=tempdir) - except Exception as exc: - # Something bad happened, maybe we couldn't parse the zip file. - # @validation_task should ensure the exception is caught and - # transformed in a generic error message for the developer, so we - # just log it and re-raise. - log.exception( - 'Could not extract upload %s for repack.', upload_pk, exc_info=exc - ) - raise - timer.log_interval('1.extracted') - log.info('Zip from upload %s extracted, repackaging', upload_pk) - file_ = tempfile.NamedTemporaryFile(suffix='.zip', delete=False) - shutil.make_archive(os.path.splitext(file_.name)[0], 'zip', tempdir) + # tempdir must *not* be on TMP_PATH, we want local fs instead. It will be + # deleted automatically once we exit the context manager. + with tempfile.TemporaryDirectory(prefix='repack_fileupload_extract') as tempdir: + try: + extract_zip(upload.path, tempdir=tempdir) + except Exception as exc: + # Something bad happened, maybe we couldn't parse the zip file. + # @validation_task should ensure the exception is caught and + # transformed in a generic error message for the developer, so we + # just log it and re-raise. + log.exception( + 'Could not extract upload %s for repack.', upload_pk, exc_info=exc + ) + raise + timer.log_interval('1.extracted') + log.info('Zip from upload %s extracted, repackaging', upload_pk) + # We'll move the file to its final location below with move_stored_file(), + # so don't let tempfile delete it. + file_ = tempfile.NamedTemporaryFile(suffix='.zip', delete=False) + shutil.make_archive(os.path.splitext(file_.name)[0], 'zip', tempdir) with open(file_.name, 'rb') as f: upload.hash = 'sha256:%s' % get_sha256(f) timer.log_interval('2.repackaged') diff --git a/src/olympia/files/tests/test_tasks.py b/src/olympia/files/tests/test_tasks.py index 8c6d005e48..160fb41231 100644 --- a/src/olympia/files/tests/test_tasks.py +++ b/src/olympia/files/tests/test_tasks.py @@ -1,5 +1,6 @@ # -*- coding: utf-8 -*- from unittest import mock +import os import tempfile import zipfile @@ -41,6 +42,7 @@ class TestRepackFileUpload(UploadTest, TestCase): upload = self.get_upload('unicode-filenames.xpi') original_hash = upload.hash fake_results = {'errors': 0} + tmp_dir_before_repacking = sorted(os.listdir(tempfile.gettempdir())) repack_fileupload(fake_results, upload.pk) upload.reload() assert upload.hash.startswith('sha256:') @@ -70,6 +72,9 @@ class TestRepackFileUpload(UploadTest, TestCase): assert info.file_size == 717 assert info.compress_size < info.file_size assert info.compress_type == zipfile.ZIP_DEFLATED + # Check we cleaned up after ourselves: we shouldn't have left anything + # in /tmp. + assert tmp_dir_before_repacking == sorted(os.listdir(tempfile.gettempdir())) class TestHideDisabledFile(TestCase):