Spaces:
Running on Zero
Running on Zero
File size: 8,148 Bytes
0802907 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """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)
|