336 lines
12 KiB
Python
336 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import errno
|
|
import os
|
|
from collections import OrderedDict
|
|
from typing import BinaryIO
|
|
|
|
from .formats import PatchReader
|
|
from .proto import Control, SyncOp, SyncOpType, TlcContainer, TlcFile
|
|
from .wire import BLOCK_SIZE
|
|
|
|
|
|
class PatchApplyError(Exception):
|
|
pass
|
|
|
|
|
|
_MODE_MASK = 0o7777
|
|
|
|
|
|
class FilePool:
|
|
def __init__(self, paths: list[str], max_open: int = 128):
|
|
self._paths = list(paths)
|
|
self._handles: OrderedDict[int, BinaryIO] = OrderedDict()
|
|
self._max_open = max(1, int(max_open))
|
|
|
|
def _touch(self, index: int, handle: BinaryIO) -> None:
|
|
self._handles.pop(index, None)
|
|
self._handles[index] = handle
|
|
|
|
def _evict_if_needed(self) -> None:
|
|
while len(self._handles) >= self._max_open:
|
|
_, handle = self._handles.popitem(last=False)
|
|
handle.close()
|
|
|
|
def open(self, index: int) -> BinaryIO:
|
|
if index < 0 or index >= len(self._paths):
|
|
raise PatchApplyError(f"file index out of range: {index}")
|
|
handle = self._handles.get(index)
|
|
if handle is None:
|
|
try:
|
|
self._evict_if_needed()
|
|
handle = open(self._paths[index], "rb")
|
|
except OSError as exc:
|
|
if exc.errno in (errno.EMFILE, errno.ENFILE):
|
|
self.close()
|
|
handle = open(self._paths[index], "rb")
|
|
else:
|
|
raise
|
|
self._handles[index] = handle
|
|
else:
|
|
self._touch(index, handle)
|
|
return handle
|
|
|
|
def size(self, index: int) -> int:
|
|
if index < 0 or index >= len(self._paths):
|
|
raise PatchApplyError(f"file index out of range: {index}")
|
|
return os.path.getsize(self._paths[index])
|
|
|
|
def close(self) -> None:
|
|
for handle in self._handles.values():
|
|
handle.close()
|
|
self._handles.clear()
|
|
|
|
|
|
def _copy_range(
|
|
dst: BinaryIO, src: BinaryIO, length: int, buffer_size: int = 32 * 1024
|
|
) -> None:
|
|
remaining = length
|
|
while remaining > 0:
|
|
chunk = src.read(min(buffer_size, remaining))
|
|
if not chunk:
|
|
break
|
|
dst.write(chunk)
|
|
remaining -= len(chunk)
|
|
|
|
|
|
def _copy_all(dst: BinaryIO, src: BinaryIO, buffer_size: int = 32 * 1024) -> None:
|
|
while True:
|
|
chunk = src.read(buffer_size)
|
|
if not chunk:
|
|
return
|
|
dst.write(chunk)
|
|
|
|
|
|
def _compute_num_blocks(file_size: int) -> int:
|
|
return (file_size + BLOCK_SIZE - 1) // BLOCK_SIZE
|
|
|
|
|
|
def _normalize_op_type(op_type: SyncOpType | int | None) -> SyncOpType | int:
|
|
return op_type if op_type is not None else SyncOpType.BLOCK_RANGE
|
|
|
|
|
|
def _is_full_file_op(op: SyncOp, target_size: int, output_size: int) -> bool:
|
|
op_type = _normalize_op_type(op.type)
|
|
if op_type != SyncOpType.BLOCK_RANGE:
|
|
return False
|
|
block_index = 0 if op.block_index is None else op.block_index
|
|
if block_index != 0:
|
|
return False
|
|
if target_size != output_size:
|
|
return False
|
|
block_span = 0 if op.block_span is None else op.block_span
|
|
return block_span == _compute_num_blocks(output_size)
|
|
|
|
|
|
def _apply_file_mode(path: str, file: TlcFile | None) -> None:
|
|
if file is None or file.mode is None:
|
|
return
|
|
mode = int(file.mode) & _MODE_MASK
|
|
try:
|
|
os.chmod(path, mode)
|
|
except (PermissionError, OSError):
|
|
pass
|
|
|
|
|
|
def apply_rsync_ops(ops: list[SyncOp], target_pool: FilePool, output: BinaryIO) -> None:
|
|
for op in ops:
|
|
op_type = _normalize_op_type(op.type)
|
|
if op_type == SyncOpType.DATA:
|
|
output.write(op.data or b"")
|
|
continue
|
|
|
|
if op_type != SyncOpType.BLOCK_RANGE:
|
|
raise PatchApplyError(f"unsupported sync op type: {op_type}")
|
|
|
|
file_index = 0 if op.file_index is None else op.file_index
|
|
block_index = 0 if op.block_index is None else op.block_index
|
|
block_span = 0 if op.block_span is None else op.block_span
|
|
|
|
file_size = target_pool.size(file_index)
|
|
last_block_index = block_index + block_span - 1
|
|
last_block_size = BLOCK_SIZE
|
|
if BLOCK_SIZE * (last_block_index + 1) > file_size:
|
|
last_block_size = file_size % BLOCK_SIZE
|
|
op_size = (block_span - 1) * BLOCK_SIZE + last_block_size
|
|
|
|
src = target_pool.open(file_index)
|
|
src.seek(block_index * BLOCK_SIZE)
|
|
_copy_range(output, src, op_size)
|
|
|
|
|
|
def apply_bsdiff_controls(
|
|
controls: list[Control], old: BinaryIO, output: BinaryIO
|
|
) -> None:
|
|
old_offset = 0
|
|
for ctrl in controls:
|
|
if ctrl.eof:
|
|
break
|
|
|
|
add = ctrl.add or b""
|
|
copy = ctrl.copy or b""
|
|
seek = ctrl.seek or 0
|
|
|
|
if add:
|
|
old.seek(old_offset)
|
|
old_chunk = old.read(len(add))
|
|
if len(old_chunk) != len(add):
|
|
raise PatchApplyError("unexpected EOF while reading bsdiff add data")
|
|
out = bytes(((a + b) & 0xFF) for a, b in zip(old_chunk, add))
|
|
output.write(out)
|
|
old_offset += len(add)
|
|
|
|
if copy:
|
|
output.write(copy)
|
|
|
|
old_offset += seek
|
|
|
|
|
|
def _ensure_parent(path: str) -> None:
|
|
parent = os.path.dirname(path)
|
|
if parent:
|
|
os.makedirs(parent, exist_ok=True)
|
|
|
|
|
|
def apply_patch(
|
|
patch_reader: PatchReader, target_paths: list[str], output_paths: list[str]
|
|
) -> None:
|
|
pool = FilePool(target_paths)
|
|
try:
|
|
expected_files = len(output_paths)
|
|
if patch_reader.source_container is not None:
|
|
expected_files = len(patch_reader.source_container.files)
|
|
entry_iter = patch_reader.iter_file_entries()
|
|
for expected_index in range(expected_files):
|
|
try:
|
|
entry = next(entry_iter)
|
|
except StopIteration:
|
|
raise PatchApplyError(
|
|
f"corrupted patch: expected {expected_files} file entries, got {expected_index}"
|
|
)
|
|
header_index = (
|
|
0
|
|
if entry.sync_header.file_index is None
|
|
else int(entry.sync_header.file_index)
|
|
)
|
|
if header_index != expected_index:
|
|
raise PatchApplyError(
|
|
f"corrupted patch: expected file index {expected_index}, got {header_index}"
|
|
)
|
|
out_index = header_index
|
|
if out_index < 0 or out_index >= len(output_paths):
|
|
raise PatchApplyError(f"output index out of range: {out_index}")
|
|
|
|
out_path = output_paths[out_index]
|
|
_ensure_parent(out_path)
|
|
|
|
if entry.is_rsync():
|
|
if entry.sync_ops is None:
|
|
raise PatchApplyError("missing rsync ops")
|
|
if entry.sync_ops:
|
|
op = entry.sync_ops[0]
|
|
target_index = 0 if op.file_index is None else op.file_index
|
|
target_file = None
|
|
output_file = None
|
|
if patch_reader.target_container and patch_reader.source_container:
|
|
if 0 <= target_index < len(patch_reader.target_container.files):
|
|
target_file = patch_reader.target_container.files[
|
|
target_index
|
|
]
|
|
if 0 <= out_index < len(patch_reader.source_container.files):
|
|
output_file = patch_reader.source_container.files[out_index]
|
|
if target_file is not None and output_file is not None:
|
|
target_size = (
|
|
int(target_file.size) if target_file.size is not None else 0
|
|
)
|
|
output_size = (
|
|
int(output_file.size) if output_file.size is not None else 0
|
|
)
|
|
if _is_full_file_op(op, target_size, output_size):
|
|
src = pool.open(target_index)
|
|
src.seek(0)
|
|
with open(out_path, "wb") as out:
|
|
_copy_all(out, src)
|
|
_apply_file_mode(out_path, output_file)
|
|
continue
|
|
|
|
with open(out_path, "wb") as out:
|
|
apply_rsync_ops(entry.sync_ops, pool, out)
|
|
output_file = None
|
|
if patch_reader.source_container and 0 <= out_index < len(
|
|
patch_reader.source_container.files
|
|
):
|
|
output_file = patch_reader.source_container.files[out_index]
|
|
_apply_file_mode(out_path, output_file)
|
|
continue
|
|
|
|
if entry.is_bsdiff():
|
|
if entry.bsdiff_header is None or entry.bsdiff_controls is None:
|
|
raise PatchApplyError("missing bsdiff data")
|
|
target_index = (
|
|
0
|
|
if entry.bsdiff_header.target_index is None
|
|
else entry.bsdiff_header.target_index
|
|
)
|
|
with open(out_path, "wb") as out:
|
|
old = pool.open(target_index)
|
|
apply_bsdiff_controls(entry.bsdiff_controls, old, out)
|
|
expected_size = None
|
|
output_file = None
|
|
if patch_reader.source_container and 0 <= out_index < len(
|
|
patch_reader.source_container.files
|
|
):
|
|
output_file = patch_reader.source_container.files[out_index]
|
|
if output_file.size is not None:
|
|
expected_size = int(output_file.size)
|
|
final_size = out.tell()
|
|
if expected_size is not None and final_size != expected_size:
|
|
raise PatchApplyError(
|
|
f"corrupted patch: expected output size {expected_size}, got {final_size}"
|
|
)
|
|
_apply_file_mode(out_path, output_file)
|
|
continue
|
|
|
|
raise PatchApplyError("unknown file patch type")
|
|
finally:
|
|
pool.close()
|
|
|
|
|
|
def _container_file_paths(container: TlcContainer, root: str) -> list[str]:
|
|
paths = []
|
|
for f in container.files:
|
|
if f.path is None:
|
|
raise PatchApplyError("container file missing path")
|
|
paths.append(os.path.join(root, f.path))
|
|
return paths
|
|
|
|
|
|
def _ensure_dirs(container: TlcContainer, root: str) -> None:
|
|
for d in container.dirs:
|
|
if d.path is None:
|
|
continue
|
|
path = os.path.join(root, d.path)
|
|
os.makedirs(path, exist_ok=True)
|
|
if d.mode is not None:
|
|
mode = int(d.mode) & _MODE_MASK
|
|
try:
|
|
os.chmod(path, mode)
|
|
except (PermissionError, OSError):
|
|
pass
|
|
|
|
|
|
def _ensure_symlinks(container: TlcContainer, root: str) -> None:
|
|
for s in container.symlinks:
|
|
if s.path is None or s.dest is None:
|
|
continue
|
|
path = os.path.join(root, s.path)
|
|
if os.path.lexists(path):
|
|
continue
|
|
parent = os.path.dirname(path)
|
|
if parent:
|
|
os.makedirs(parent, exist_ok=True)
|
|
try:
|
|
os.symlink(s.dest, path)
|
|
except (AttributeError, OSError):
|
|
continue
|
|
|
|
|
|
def apply_patch_to_folders(patch_path: str, target_root: str, output_root: str) -> None:
|
|
with PatchReader.open(
|
|
patch_path,
|
|
container_decoder=(TlcContainer.from_bytes, TlcContainer.from_bytes),
|
|
) as reader:
|
|
if reader.target_container is None or reader.source_container is None:
|
|
raise PatchApplyError("missing containers in patch")
|
|
|
|
target_container: TlcContainer = reader.target_container
|
|
source_container: TlcContainer = reader.source_container
|
|
|
|
_ensure_dirs(source_container, output_root)
|
|
_ensure_symlinks(source_container, output_root)
|
|
|
|
target_paths = _container_file_paths(target_container, target_root)
|
|
output_paths = _container_file_paths(source_container, output_root)
|
|
|
|
apply_patch(reader, target_paths, output_paths)
|