Skip to content

Commit

Permalink
fix: allow closed file handles in calc_sha256sum
Browse files Browse the repository at this point in the history
Since we tend to allow filenames and file handles interchangably in this
code base, we might run into the edge case of a closed file handle being
input to calc_sha256sum. We should then open the file before getting the
hash.
  • Loading branch information
hermankolden authored and bkueng committed Jul 4, 2024
1 parent 8bcee7b commit a84bf52
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 6 deletions.
11 changes: 6 additions & 5 deletions pyulog/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,15 +114,16 @@ def calc_sha256sum(log_file):
if log_file is None:
return None
if isinstance(log_file, str):
log_file = open(log_file, 'rb') # pylint: disable=consider-using-with
log_context = log_file
file_context = open(log_file, 'rb') # pylint: disable=consider-using-with
elif log_file.closed:
file_context = open(log_file.name, 'rb')
else:
log_context = contextlib.nullcontext()
file_context = contextlib.nullcontext(log_file)

file_hash = hashlib.sha256()
with log_context:
with file_context as open_file:
while True:
block = log_file.read(4096)
block = open_file.read(4096)
file_hash.update(block)
if block == b'':
break
Expand Down
11 changes: 10 additions & 1 deletion test/test_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,8 @@ def test_unapplied_migrations(self):
@data('sample_log_small')
def test_sha256sum(self, test_case):
'''
Verify that the sha256sum set on save can be used to find the same file again.
Verify that the sha256sum set on save can be used to find the same file
again, using any of the approved file input methods.
'''

test_file = os.path.join(TEST_PATH, f'{test_case}.ulg')
Expand All @@ -151,6 +152,14 @@ def test_sha256sum(self, test_case):
digest = DatabaseULog.calc_sha256sum(test_file)
self.assertEqual(digest, dbulog.sha256sum)

test_file_handle = open(test_file, 'rb') # pylint: disable=consider-using-with
open_digest = DatabaseULog.calc_sha256sum(test_file_handle)
self.assertEqual(digest, open_digest)

test_file_handle.close()
closed_digest = DatabaseULog.calc_sha256sum(test_file_handle)
self.assertEqual(digest, closed_digest)

pk_from_digest = DatabaseULog.primary_key_from_sha256sum(self.db_handle, digest)
self.assertEqual(pk_from_digest, dbulog.primary_key)

Expand Down

0 comments on commit a84bf52

Please sign in to comment.