truf-server/app/remote_worker_bootstrap.py

"""Stdlib-only integrity boundary for packaged remote worker clients."""

import hashlib
import json
import os
import runpy
import stat
import sys

sys.dont_write_bytecode = True
if not sys.dont_write_bytecode:
    raise RuntimeError('worker bootstrap could not disable bytecode writes')


MAX_MANIFEST_BYTES = 1024 * 1024
MAX_MANIFEST_FILES = 512
WORKER_PACKAGE_SCHEMA = 3
WORKER_PROTOCOL_VERSION = 2


def _canonical(path):
    return os.path.normcase(os.path.realpath(os.path.abspath(os.fspath(path))))


def _is_reparse_point(path):
    details = os.lstat(path)
    if stat.S_ISLNK(details.st_mode):
        return True
    attributes = getattr(details, 'st_file_attributes', 0)
    reparse_attribute = getattr(stat, 'FILE_ATTRIBUTE_REPARSE_POINT', 0)
    return bool(attributes & reparse_attribute) or getattr(
        os.path, 'isjunction', lambda _path: False,
    )(path)


def _relative(value, label):
    value = str(value or '')
    if (
        not value or len(value) > 512 or '\\' in value or '\x00' in value
        or value.startswith('/') or value.endswith('/')
    ):
        raise RuntimeError(f'invalid worker package {label} path')
    if any(part in ('', '.', '..') for part in value.split('/')):
        raise RuntimeError(f'invalid worker package {label} path')
    return value


def _sha256(path):
    digest = hashlib.sha256()
    with open(path, 'rb', buffering=0) as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b''):
            digest.update(block)
    return digest.hexdigest()


def _load_manifest(path):
    details = os.stat(path, follow_symlinks=False)
    if _is_reparse_point(path) or not stat.S_ISREG(details.st_mode):
        raise RuntimeError('worker package manifest is not a regular file')
    with open(path, 'rb') as handle:
        payload = handle.read(MAX_MANIFEST_BYTES + 1)
    if len(payload) > MAX_MANIFEST_BYTES:
        raise RuntimeError('worker package manifest exceeds its byte bound')
    value = json.loads(payload.decode('utf-8', errors='strict'))
    if (
        not isinstance(value, dict)
        or type(value.get('schema')) is not int
        or value['schema'] != WORKER_PACKAGE_SCHEMA
        or type(value.get('protocol_version')) is not int
        or value['protocol_version'] != WORKER_PROTOCOL_VERSION
    ):
        raise RuntimeError('worker package manifest is invalid')
    return value


def _application_files(app_dir):
    files = set()

    def raise_walk_error(exc):
        raise RuntimeError(f'unable to inspect worker application root: {exc}') from exc

    for current, directories, names in os.walk(
        app_dir, followlinks=False, onerror=raise_walk_error,
    ):
        for name in directories:
            candidate = os.path.join(current, name)
            if _is_reparse_point(candidate) or name.lower() == '__pycache__':
                raise RuntimeError('worker application directory is unsupported')
        for name in names:
            candidate = os.path.join(current, name)
            details = os.stat(candidate, follow_symlinks=False)
            if _is_reparse_point(candidate) or not stat.S_ISREG(details.st_mode):
                raise RuntimeError('worker application file is not regular')
            files.add(os.path.relpath(candidate, app_dir).replace(os.sep, '/'))
    return files


def _verify_application(package_root, manifest):
    app_root = _relative(manifest.get('app_root'), 'application root')
    app_dir = _canonical(os.path.join(package_root, *app_root.split('/')))
    if app_dir != _canonical(os.path.dirname(__file__)) or _is_reparse_point(app_dir):
        raise RuntimeError('worker package application root is not canonical')
    values = manifest.get('files')
    if not isinstance(values, dict) or not 1 <= len(values) <= MAX_MANIFEST_FILES:
        raise RuntimeError('worker package file set is invalid')
    expected = set()
    for name, entry in values.items():
        name = _relative(name, 'file name')
        if not isinstance(entry, dict) or set(entry) != {'path', 'sha256'}:
            raise RuntimeError('worker package file entry is invalid')
        relative = _relative(entry.get('path'), f'file {name}')
        if relative != f'{app_root}/{name}':
            raise RuntimeError('worker package file path is not canonical')
        digest = str(entry.get('sha256') or '')
        if len(digest) != 64 or any(ch not in '0123456789abcdef' for ch in digest):
            raise RuntimeError('worker package file digest is invalid')
        path = _canonical(os.path.join(package_root, *relative.split('/')))
        try:
            contained = os.path.commonpath((app_dir, path)) == app_dir
        except ValueError:
            contained = False
        if not contained or _is_reparse_point(path) or _sha256(path) != digest:
            raise RuntimeError('worker package application integrity check failed')
        expected.add(name)
    if _application_files(app_dir) != expected:
        raise RuntimeError('worker package application file set drifted')
    return app_dir


def main():
    if not (
        sys.flags.isolated and sys.flags.no_site and sys.flags.dont_write_bytecode
    ):
        raise RuntimeError(
            'remote worker bootstrap requires isolated no-site bytecode-free startup (-I -S -B)'
        )
    sys.dont_write_bytecode = True
    if len(sys.argv) < 2 or sys.argv[1] != '--':
        raise RuntimeError('usage: remote_worker_bootstrap.py -- <worker args>')

    package_root = _canonical(os.path.dirname(os.path.dirname(__file__)))
    manifest = _load_manifest(os.path.join(package_root, 'worker-package.json'))
    app_dir = _verify_application(package_root, manifest)
    dependency_dir = os.path.join(app_dir, 'dependencies')
    if (
        not os.path.isdir(dependency_dir) or _is_reparse_point(dependency_dir)
        or _canonical(dependency_dir) == app_dir
    ):
        raise RuntimeError('package-local worker dependencies are unavailable')

    entrypoint = os.path.join(app_dir, 'worker_cli.py')
    sys.path.insert(0, dependency_dir)
    sys.path.insert(0, app_dir)
    sys.argv = [entrypoint, *sys.argv[2:]]
    runpy.run_path(entrypoint, run_name='__main__')


if __name__ == '__main__':
    try:
        main()
    except Exception as exc:
        raise SystemExit('remote worker bootstrap rejected launch') from exc