Spaces:
Running on Zero
Running on Zero
| """Authenticate the maintainer's sealed evaluator before publishing private files.""" | |
| from __future__ import annotations | |
| import base64 | |
| import hashlib | |
| import hmac | |
| import io | |
| import json | |
| import os | |
| from pathlib import Path | |
| import re | |
| import shutil | |
| import stat | |
| import tempfile | |
| import zipfile | |
| from cryptography.hazmat.primitives.ciphers.aead import AESGCM | |
| MAGIC = b"AUTORSE-EVAL-GCM1\n" | |
| MAX_BLOB_BYTES = 64 * 1024**2 | |
| MAX_EXPANDED_BYTES = 64 * 1024**2 | |
| MAX_MANIFEST_BYTES = 64 * 1024 | |
| PRIVATE_FILES = frozenset("""hstu/beauty/expected_user_ids.npy hstu/beauty/catalog.npy hstu/beauty/test_targets.npy | |
| tiger/beauty/expected_user_ids.npy tiger/beauty/catalog.npy tiger/beauty/test_targets.npy | |
| intsr/Amazon/expected_submission.npz intsr/Amazon/test_data_labeled.csv | |
| wukong/movielenslatest/test_data_labeled.csv rankmixer/criteo5m/expected_sample_ids.npy | |
| rankmixer/criteo5m/test_labels.npy onetrans/kuaisar_v2/labels.npz""".split()) | |
| FILES = PRIVATE_FILES | {"private-manifest.json"} | |
| ERROR = "Encrypted evaluator verification or initialization failed." | |
| def make_aad(expected_manifest_sha, public_bundle_sha, space_id) -> bytes: | |
| return json.dumps({"protocol": "autorse-evaluator-aesgcm-v1", | |
| "manifest_sha256": expected_manifest_sha, | |
| "public_bundle_sha256": public_bundle_sha, | |
| "space_id": space_id}, sort_keys=True, separators=(",", ":")).encode("utf-8") | |
| def _require(condition): | |
| if not condition: | |
| raise ValueError(ERROR) | |
| def _hash(data): | |
| return hashlib.sha256(data).hexdigest() | |
| def _sha(value): | |
| return isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) is not None | |
| def _unique_object(pairs): | |
| result = {} | |
| for key, value in pairs: | |
| _require(key not in result) | |
| result[key] = value | |
| return result | |
| def _unpack(plaintext, expected_manifest_sha, public_bundle_sha): | |
| _require(len(plaintext) <= MAX_BLOB_BYTES) | |
| with zipfile.ZipFile(io.BytesIO(plaintext)) as archive: | |
| infos = archive.infolist() | |
| _require(len(infos) == len(FILES) and {i.filename for i in infos} == FILES) | |
| _require(sum(i.file_size for i in infos) <= MAX_EXPANDED_BYTES) | |
| by_name = {} | |
| for info in infos: | |
| name = info.filename | |
| _require(name not in by_name and not name.startswith("/") and "\\" not in name | |
| and all(part not in ("", ".", "..") for part in name.split("/"))) | |
| mode = info.external_attr >> 16 | |
| _require(not info.is_dir() and stat.S_IFMT(mode) in (0, stat.S_IFREG) | |
| and not info.flag_bits & 1 and 0 <= info.file_size <= MAX_EXPANDED_BYTES | |
| and info.compress_type in (zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED)) | |
| by_name[name] = info | |
| def read(name, size): | |
| info = by_name[name] | |
| _require(info.file_size == size) | |
| with archive.open(info) as stream: | |
| value = stream.read(size + 1) | |
| _require(len(value) == size) | |
| return value | |
| manifest_info = by_name["private-manifest.json"] | |
| _require(manifest_info.file_size <= MAX_MANIFEST_BYTES) | |
| raw_manifest = read("private-manifest.json", manifest_info.file_size) | |
| _require(hmac.compare_digest(_hash(raw_manifest), expected_manifest_sha)) | |
| manifest = json.loads(raw_manifest, object_pairs_hook=_unique_object) | |
| _require(isinstance(manifest, dict) and manifest.get("schema_version") == 1 | |
| and manifest.get("visibility") == "PRIVATE_EVALUATOR_ONLY" | |
| and manifest.get("public_bundle_sha256") == public_bundle_sha) | |
| rows = manifest.get("files") | |
| _require(isinstance(rows, list) and len(rows) == len(PRIVATE_FILES)) | |
| result = {"private-manifest.json": raw_manifest} | |
| for row in rows: | |
| _require(isinstance(row, dict) and set(row) == {"path", "size", "sha256"}) | |
| name, size, sha = row["path"], row["size"], row["sha256"] | |
| _require(isinstance(name, str) and name in PRIVATE_FILES and name not in result | |
| and type(size) is int and 0 <= size <= MAX_EXPANDED_BYTES and _sha(sha)) | |
| value = read(name, size) | |
| _require(hmac.compare_digest(_hash(value), sha)) | |
| result[name] = value | |
| _require(set(result) == FILES) | |
| return result | |
| def _safe_chain(directory): | |
| for path in (directory, *directory.parents): | |
| _require(not path.is_symlink() and (not path.exists() or path.is_dir())) | |
| def _existing_matches(root, payload): | |
| entries = list(root.rglob("*")) | |
| _require(not any(p.is_symlink() for p in entries)) | |
| if not entries: | |
| return False | |
| _require({p.relative_to(root).as_posix() for p in entries if p.is_file()} == FILES) | |
| for path in (root, *entries): | |
| metadata = path.stat() | |
| if path.is_dir(): | |
| _require(stat.S_IMODE(metadata.st_mode) == 0o700) | |
| else: | |
| name = path.relative_to(root).as_posix() | |
| _require(stat.S_ISREG(metadata.st_mode) and metadata.st_nlink == 1 | |
| and stat.S_IMODE(metadata.st_mode) == 0o600 | |
| and metadata.st_size == len(payload[name])) | |
| with path.open("rb") as stream: | |
| digest = hashlib.file_digest(stream, "sha256").hexdigest() | |
| _require(hmac.compare_digest(digest, _hash(payload[name]))) | |
| return True | |
| def _write_file(path, data): | |
| descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), 0o600) | |
| with os.fdopen(descriptor, "wb") as stream: | |
| os.fchmod(stream.fileno(), 0o600) | |
| stream.write(data) | |
| def provision_encrypted_data(*, blob_path: Path, private_root: Path, key_b64: str, | |
| expected_cipher_sha: str, expected_manifest_sha: str, | |
| public_bundle_sha: str, | |
| space_id: str = "Ray0311/AutoRSE-Evaluator") -> None: | |
| stage = None | |
| try: | |
| _require(_sha(expected_cipher_sha) and _sha(expected_manifest_sha) and _sha(public_bundle_sha) | |
| and isinstance(space_id, str) and 0 < len(space_id) <= 256 and isinstance(key_b64, str)) | |
| key = base64.b64decode(key_b64, validate=True) | |
| _require(len(key) == 32) | |
| blob_path, root = Path(blob_path), Path(private_root).absolute() | |
| _require(not blob_path.is_symlink() and blob_path.is_file() | |
| and len(MAGIC) + 12 + 16 <= blob_path.stat().st_size <= MAX_BLOB_BYTES) | |
| with blob_path.open("rb") as stream: | |
| blob = stream.read(MAX_BLOB_BYTES + 1) | |
| _require(len(blob) <= MAX_BLOB_BYTES and blob.startswith(MAGIC) | |
| and hmac.compare_digest(_hash(blob), expected_cipher_sha)) | |
| offset = len(MAGIC) | |
| plaintext = AESGCM(key).decrypt(blob[offset:offset + 12], blob[offset + 12:], | |
| make_aad(expected_manifest_sha, public_bundle_sha, space_id)) | |
| payload = _unpack(plaintext, expected_manifest_sha, public_bundle_sha) | |
| _safe_chain(root) | |
| if root.exists() and _existing_matches(root, payload): | |
| return | |
| missing = [] | |
| parent = root.parent | |
| while not parent.exists(): | |
| missing.append(parent) | |
| parent = parent.parent | |
| for parent in reversed(missing): | |
| parent.mkdir(mode=0o700) | |
| parent.chmod(0o700) | |
| _safe_chain(root) | |
| stage = Path(tempfile.mkdtemp(prefix="." + root.name + "-", dir=root.parent)) | |
| stage.chmod(0o700) | |
| for name, data in payload.items(): | |
| path = stage / name | |
| path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) | |
| _write_file(path, data) | |
| for directory in stage.rglob("*"): | |
| if directory.is_dir(): | |
| directory.chmod(0o700) | |
| _safe_chain(root) | |
| os.replace(stage, root) | |
| stage = None | |
| except Exception: | |
| raise ValueError(ERROR) from None | |
| finally: | |
| if stage is not None: | |
| shutil.rmtree(stage, ignore_errors=True) | |