181 lines
8.7 KiB
Python
181 lines
8.7 KiB
Python
|
"""Protected snapshot bytes are immutable to task UIDs and verified on reuse."""
|
||
|
|
import hashlib
|
||
|
|
import fcntl
|
||
|
|
import os
|
||
|
|
import select
|
||
|
|
from pathlib import Path
|
||
|
|
import subprocess
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
import uuid
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
|
||
|
|
@unittest.skipUnless(os.geteuid() == 0, 'protected artifact ownership requires root')
|
||
|
|
class FileArtifacts(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.temp = tempfile.TemporaryDirectory()
|
||
|
|
self.addCleanup(self.temp.cleanup)
|
||
|
|
self.root = Path(self.temp.name)
|
||
|
|
source = Path(__file__).resolve().parents[1] / 'crates/supervisor/src/runner_files.py'
|
||
|
|
self.module = {'ROOT': self.root}
|
||
|
|
exec(compile(source.read_text(), str(source), 'exec'), self.module)
|
||
|
|
|
||
|
|
def test_dedup_ownership_and_tamper_refusal(self):
|
||
|
|
data = b'private\0\xffsnapshot'
|
||
|
|
digest = hashlib.sha256(data).hexdigest()
|
||
|
|
store = self.module['file_snapshot_store']
|
||
|
|
directory = self.module['file_snapshot_directory']()
|
||
|
|
self.assertEqual(store(data, digest), digest)
|
||
|
|
inode = (directory / digest).stat().st_ino
|
||
|
|
self.assertEqual(store(data, digest), digest)
|
||
|
|
self.assertEqual((directory / digest).stat().st_ino, inode)
|
||
|
|
denied = subprocess.run(['python3', '-c', 'import sys;print(open(sys.argv[1],"rb").read())', str(directory / digest)],
|
||
|
|
user=1000, group=1000, extra_groups=[], capture_output=True)
|
||
|
|
self.assertNotEqual(denied.returncode, 0)
|
||
|
|
self.assertEqual(denied.stdout, b'')
|
||
|
|
(directory / digest).write_bytes(b'tampered')
|
||
|
|
with self.assertRaisesRegex(ValueError, 'INTEGRITY'):
|
||
|
|
store(data, digest)
|
||
|
|
with self.assertRaisesRegex(ValueError, 'INTEGRITY'):
|
||
|
|
self.module['file_snapshot_read'](directory, digest)
|
||
|
|
(directory / digest).unlink()
|
||
|
|
(directory / digest).symlink_to('/etc/passwd')
|
||
|
|
with self.assertRaisesRegex(ValueError, 'UNPROTECTED_ARTIFACT'):
|
||
|
|
store(data, digest)
|
||
|
|
|
||
|
|
def test_killed_publishers_release_pending_files_without_losing_published_bytes(self):
|
||
|
|
directory = self.module['file_snapshot_directory']()
|
||
|
|
source = Path(__file__).resolve().parents[1] / 'crates/supervisor/src/runner_files.py'
|
||
|
|
child = '''
|
||
|
|
import hashlib, os, sys
|
||
|
|
from pathlib import Path
|
||
|
|
module = {'ROOT': Path(sys.argv[1])}
|
||
|
|
exec(compile(Path(sys.argv[2]).read_text(), sys.argv[2], 'exec'), module)
|
||
|
|
hook = sys.argv[3]
|
||
|
|
original = getattr(os, hook)
|
||
|
|
def pause(*args):
|
||
|
|
original(*args)
|
||
|
|
print('ready', flush=True)
|
||
|
|
sys.stdin.read()
|
||
|
|
setattr(os, hook, pause)
|
||
|
|
data = hook.encode()
|
||
|
|
module['file_snapshot_store'](data, hashlib.sha256(data).hexdigest())
|
||
|
|
'''
|
||
|
|
for hook in ('fsync', 'link'):
|
||
|
|
with self.subTest(crash_after=hook):
|
||
|
|
process = subprocess.Popen(['python3', '-c', child, str(self.root), str(source), hook],
|
||
|
|
stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
|
|
try:
|
||
|
|
self.assertTrue(select.select([process.stdout], [], [], 10)[0], 'publisher did not reach fault point')
|
||
|
|
self.assertEqual(process.stdout.readline(), b'ready\n')
|
||
|
|
pending = list(directory.glob('.pending-*'))
|
||
|
|
self.assertEqual(len(pending), 1)
|
||
|
|
digest = hashlib.sha256(hook.encode()).hexdigest()
|
||
|
|
self.assertEqual((directory / digest).exists(), hook == 'link')
|
||
|
|
inode = pending[0].stat().st_ino
|
||
|
|
process.kill()
|
||
|
|
process.wait(timeout=10)
|
||
|
|
self.assertTrue(pending[0].exists())
|
||
|
|
self.module['file_snapshot_store'](hook.encode(), digest)
|
||
|
|
self.assertEqual(list(directory.glob('.pending-*')), [])
|
||
|
|
self.assertEqual(self.module['file_snapshot_read'](directory, digest), hook.encode())
|
||
|
|
if hook == 'link':
|
||
|
|
self.assertEqual((directory / digest).stat().st_ino, inode)
|
||
|
|
finally:
|
||
|
|
if process.poll() is None:
|
||
|
|
process.kill()
|
||
|
|
process.communicate(timeout=10)
|
||
|
|
|
||
|
|
def test_cleanup_recovers_quota_but_preserves_other_entries(self):
|
||
|
|
directory = self.module['file_snapshot_directory']()
|
||
|
|
data = b'quota recovery'
|
||
|
|
digest = hashlib.sha256(data).hexdigest()
|
||
|
|
pending = directory / ('.pending-' + uuid.uuid4().hex)
|
||
|
|
with os.fdopen(os.open(pending, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), 'wb') as stream:
|
||
|
|
stream.truncate(256 * 1024 * 1024)
|
||
|
|
self.module['file_snapshot_store'](data, digest)
|
||
|
|
self.assertFalse(pending.exists())
|
||
|
|
other = directory / '.pending-do-not-delete'
|
||
|
|
with os.fdopen(os.open(other, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), 'wb') as stream:
|
||
|
|
stream.truncate(256 * 1024 * 1024)
|
||
|
|
with self.assertRaisesRegex(ValueError, 'ARTIFACT_QUOTA'):
|
||
|
|
self.module['file_snapshot_store'](b'new', hashlib.sha256(b'new').hexdigest())
|
||
|
|
self.assertTrue(other.exists())
|
||
|
|
self.assertEqual(self.module['file_snapshot_read'](directory, digest), data)
|
||
|
|
|
||
|
|
def test_cleanup_waits_for_publisher_lock(self):
|
||
|
|
directory = self.module['file_snapshot_directory']()
|
||
|
|
source = Path(__file__).resolve().parents[1] / 'crates/supervisor/src/runner_files.py'
|
||
|
|
pending = directory / ('.pending-' + uuid.uuid4().hex)
|
||
|
|
guard = os.open(self.root / 'artifact-store.lock', os.O_WRONLY | os.O_CREAT, 0o600)
|
||
|
|
fcntl.flock(guard, fcntl.LOCK_EX)
|
||
|
|
process = None
|
||
|
|
try:
|
||
|
|
with os.fdopen(os.open(pending, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), 'wb') as stream:
|
||
|
|
stream.write(b'active publisher')
|
||
|
|
child = '''
|
||
|
|
import hashlib, sys
|
||
|
|
from pathlib import Path
|
||
|
|
module = {'ROOT': Path(sys.argv[1])}
|
||
|
|
exec(compile(Path(sys.argv[2]).read_text(), sys.argv[2], 'exec'), module)
|
||
|
|
print('ready', flush=True)
|
||
|
|
module['file_snapshot_store'](b'next', hashlib.sha256(b'next').hexdigest())
|
||
|
|
print('done', flush=True)
|
||
|
|
'''
|
||
|
|
process = subprocess.Popen(['python3', '-c', child, str(self.root), str(source)],
|
||
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
|
|
self.assertTrue(select.select([process.stdout], [], [], 10)[0])
|
||
|
|
self.assertEqual(process.stdout.readline(), b'ready\n')
|
||
|
|
self.assertFalse(select.select([process.stdout], [], [], 0.2)[0])
|
||
|
|
self.assertTrue(pending.exists())
|
||
|
|
fcntl.flock(guard, fcntl.LOCK_UN)
|
||
|
|
stdout, stderr = process.communicate(timeout=10)
|
||
|
|
self.assertEqual(process.returncode, 0, stderr)
|
||
|
|
self.assertEqual(stdout, b'done\n')
|
||
|
|
self.assertFalse(pending.exists())
|
||
|
|
finally:
|
||
|
|
os.close(guard)
|
||
|
|
if process is not None:
|
||
|
|
if process.poll() is None:
|
||
|
|
process.kill()
|
||
|
|
process.communicate(timeout=10)
|
||
|
|
|
||
|
|
def test_cleanup_sync_failure_refuses_publication_and_allows_retry(self):
|
||
|
|
directory = self.module['file_snapshot_directory']()
|
||
|
|
pending = directory / ('.pending-' + uuid.uuid4().hex)
|
||
|
|
os.close(os.open(pending, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600))
|
||
|
|
data = b'after cleanup'
|
||
|
|
digest = hashlib.sha256(data).hexdigest()
|
||
|
|
with patch.object(os, 'fsync', side_effect=OSError('cleanup sync failed')):
|
||
|
|
with self.assertRaisesRegex(OSError, 'cleanup sync failed'):
|
||
|
|
self.module['file_snapshot_store'](data, digest)
|
||
|
|
self.assertFalse((directory / digest).exists())
|
||
|
|
self.assertEqual(self.module['file_snapshot_store'](data, digest), digest)
|
||
|
|
self.assertEqual(self.module['file_snapshot_read'](directory, digest), data)
|
||
|
|
|
||
|
|
def test_failed_sync_never_turns_into_success_on_dedup_retry(self):
|
||
|
|
directory = self.module['file_snapshot_directory']()
|
||
|
|
data = b'durable snapshot'
|
||
|
|
digest = hashlib.sha256(data).hexdigest()
|
||
|
|
original = os.fsync
|
||
|
|
calls = 0
|
||
|
|
def fail_after_link(fd):
|
||
|
|
nonlocal calls
|
||
|
|
calls += 1
|
||
|
|
if calls == 2:
|
||
|
|
raise OSError('sync unavailable')
|
||
|
|
original(fd)
|
||
|
|
with patch.object(os, 'fsync', side_effect=fail_after_link):
|
||
|
|
with self.assertRaises(OSError):
|
||
|
|
self.module['file_snapshot_store'](data, digest)
|
||
|
|
self.assertTrue((directory / digest).exists())
|
||
|
|
with patch.object(os, 'fsync', side_effect=OSError('still unavailable')):
|
||
|
|
with self.assertRaises(OSError):
|
||
|
|
self.module['file_snapshot_store'](data, digest)
|
||
|
|
self.assertEqual(self.module['file_snapshot_store'](data, digest), digest)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|