switch to rebuilding for a better security

src/gradio_space_ci/webhook.py

@@ -28,9 +28,6 @@ from huggingface_hub import (
     snapshot_download,
     space_info,
     upload_folder,
-    delete_space_secret,
-    delete_space_variable,
-    delete_space_storage,
 )
 from huggingface_hub.repocard import RepoCard
 from huggingface_hub.utils import (
@@ -487,19 +484,37 @@ def set_config(space_id: str, pr_num: str) -> None:
         request_space_storage(ci_space_id, storage)
 
 
-def unset_config(space_id: str, pr_num: int) -> None:
-    "a function to unset the configuration of an ephemeral space"
+def rebuild_space(space_id: str, pr_num: int) -> None:
+    "a function to rebuild the ephemeral space without config"
+    # This is useful to cut down on resource usage and to remove tokens from
+    # the ephemeral space
     ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
-    variables: Dict[str, str] = EPHEMERAL_SPACES_CONFIG["variables"]
-    secrets: Dict[str, str] = EPHEMERAL_SPACES_CONFIG["secrets"]
-    # Unset space variables and secrets
-    for key in variables.keys():
-        delete_space_variable(ci_space_id, key)
-    for key in secrets.keys():
-        delete_space_secret(ci_space_id, key)
-    # Reset hardware
-    request_space_hardware(ci_space_id, SpaceHardware.CPU_BASIC)
-    delete_space_storage(ci_space_id)
+    try:
+        delete_repo(repo_id=ci_space_id, repo_type="space")
+    except RepositoryNotFoundError:
+        pass
+    create_ephemeral_space(space_id=space_id, pr_num=pr_num)
+    # Download space codebase from PR revision
+    snapshot_path = Path(snapshot_download(repo_id=space_id, revision=f"refs/pr/{pr_num}", repo_type="space"))
+
+    # Overwrite README file in cache (/!\)
+    readme_path = snapshot_path / "README.md"
+    card = RepoCard.load(readme_path)
+    setattr(card.data, "synced_sha", snapshot_path.name)  # latest sha
+    card.data.title = f"{card.data.title} (ephemeral #{pr_num})"
+    card.save(readme_path)
+
+    # Sync space codebase with PR revision
+    upload_folder(
+        repo_id=ci_space_id,
+        repo_type="space",
+        commit_message=f"Sync CI Space with PR {pr_num}.",
+        folder_path=snapshot_path,
+        delete_patterns="*",
+    )
+
+    # Delete readme file from cache (just in case)
+    readme_path.unlink(missing_ok=True)
 
 
 def handle_modification(space_id: str, discussion: Any) -> None:
@@ -509,8 +524,8 @@ def handle_modification(space_id: str, discussion: Any) -> None:
     details = get_discussion_details(repo_id=space_id, repo_type="space", discussion_num=discussion.num)
     event_author = details.events[-1]._event["author"]["name"]  # username of that event
     if event_author not in EPHEMERAL_SPACES_CONFIG["trusted_authors"]:
-        # Untrusted author, we unset the config as part or security reasons
-        unset_config(space_id=space_id, pr_num=discussion.num)
+        # Untrusted author, we rebuild the space
+        rebuild_space(space_id=space_id, pr_num=discussion.num)
 
 
 def handle_command(space_id: str, payload: WebhookPayload) -> None:
@@ -523,7 +538,7 @@ def handle_command(space_id: str, payload: WebhookPayload) -> None:
             set_config(space_id=space_id, pr_num=pr_num)
             notify_pr(space_id=space_id, pr_num=pr_num, action="trusted_pr")
         elif payload.comment.content == "/untrust_pr":
-            unset_config(space_id=space_id, pr_num=pr_num)
+            rebuild_space(space_id=space_id, pr_num=pr_num)
             notify_pr(space_id=space_id, pr_num=pr_num, action="untrusted_pr")