not-lain commited on
Commit
9d51341
1 Parent(s): 7a55abe

rebuild pr when untrusted for security reasons

Browse files
Files changed (1) hide show
  1. src/gradio_space_ci/webhook.py +33 -18
src/gradio_space_ci/webhook.py CHANGED
@@ -28,9 +28,6 @@ from huggingface_hub import (
28
  snapshot_download,
29
  space_info,
30
  upload_folder,
31
- delete_space_secret,
32
- delete_space_variable,
33
- delete_space_storage,
34
  )
35
  from huggingface_hub.repocard import RepoCard
36
  from huggingface_hub.utils import (
@@ -487,19 +484,37 @@ def set_config(space_id: str, pr_num: str) -> None:
487
  request_space_storage(ci_space_id, storage)
488
 
489
 
490
- def unset_config(space_id: str, pr_num: int) -> None:
491
- "a function to unset the configuration of an ephemeral space"
 
 
492
  ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
493
- variables: Dict[str, str] = EPHEMERAL_SPACES_CONFIG["variables"]
494
- secrets: Dict[str, str] = EPHEMERAL_SPACES_CONFIG["secrets"]
495
- # Unset space variables and secrets
496
- for key in variables.keys():
497
- delete_space_variable(ci_space_id, key)
498
- for key in secrets.keys():
499
- delete_space_secret(ci_space_id, key)
500
- # Reset hardware
501
- request_space_hardware(ci_space_id, SpaceHardware.CPU_BASIC)
502
- delete_space_storage(ci_space_id)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
503
 
504
 
505
  def handle_modification(space_id: str, discussion: Any) -> None:
@@ -509,8 +524,8 @@ def handle_modification(space_id: str, discussion: Any) -> None:
509
  details = get_discussion_details(repo_id=space_id, repo_type="space", discussion_num=discussion.num)
510
  event_author = details.events[-1]._event["author"]["name"] # username of that event
511
  if event_author not in EPHEMERAL_SPACES_CONFIG["trusted_authors"]:
512
- # Untrusted author, we unset the config as part or security reasons
513
- unset_config(space_id=space_id, pr_num=discussion.num)
514
 
515
 
516
  def handle_command(space_id: str, payload: WebhookPayload) -> None:
@@ -523,7 +538,7 @@ def handle_command(space_id: str, payload: WebhookPayload) -> None:
523
  set_config(space_id=space_id, pr_num=pr_num)
524
  notify_pr(space_id=space_id, pr_num=pr_num, action="trusted_pr")
525
  elif payload.comment.content == "/untrust_pr":
526
- unset_config(space_id=space_id, pr_num=pr_num)
527
  notify_pr(space_id=space_id, pr_num=pr_num, action="untrusted_pr")
528
 
529
 
 
28
  snapshot_download,
29
  space_info,
30
  upload_folder,
 
 
 
31
  )
32
  from huggingface_hub.repocard import RepoCard
33
  from huggingface_hub.utils import (
 
484
  request_space_storage(ci_space_id, storage)
485
 
486
 
487
+ def rebuild_space(space_id: str, pr_num: int) -> None:
488
+ "a function to rebuild the ephemeral space without config"
489
+ # This is useful to cut down on resource usage and to remove tokens from
490
+ # the ephemeral space
491
  ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
492
+ try:
493
+ delete_repo(repo_id=ci_space_id, repo_type="space")
494
+ except RepositoryNotFoundError:
495
+ pass
496
+ create_ephemeral_space(space_id=space_id, pr_num=pr_num)
497
+ # Download space codebase from PR revision
498
+ snapshot_path = Path(snapshot_download(repo_id=space_id, revision=f"refs/pr/{pr_num}", repo_type="space"))
499
+
500
+ # Overwrite README file in cache (/!\)
501
+ readme_path = snapshot_path / "README.md"
502
+ card = RepoCard.load(readme_path)
503
+ setattr(card.data, "synced_sha", snapshot_path.name) # latest sha
504
+ card.data.title = f"{card.data.title} (ephemeral #{pr_num})"
505
+ card.save(readme_path)
506
+
507
+ # Sync space codebase with PR revision
508
+ upload_folder(
509
+ repo_id=ci_space_id,
510
+ repo_type="space",
511
+ commit_message=f"Sync CI Space with PR {pr_num}.",
512
+ folder_path=snapshot_path,
513
+ delete_patterns="*",
514
+ )
515
+
516
+ # Delete readme file from cache (just in case)
517
+ readme_path.unlink(missing_ok=True)
518
 
519
 
520
  def handle_modification(space_id: str, discussion: Any) -> None:
 
524
  details = get_discussion_details(repo_id=space_id, repo_type="space", discussion_num=discussion.num)
525
  event_author = details.events[-1]._event["author"]["name"] # username of that event
526
  if event_author not in EPHEMERAL_SPACES_CONFIG["trusted_authors"]:
527
+ # Untrusted author, we rebuild the space
528
+ rebuild_space(space_id=space_id, pr_num=discussion.num)
529
 
530
 
531
  def handle_command(space_id: str, payload: WebhookPayload) -> None:
 
538
  set_config(space_id=space_id, pr_num=pr_num)
539
  notify_pr(space_id=space_id, pr_num=pr_num, action="trusted_pr")
540
  elif payload.comment.content == "/untrust_pr":
541
+ rebuild_space(space_id=space_id, pr_num=pr_num)
542
  notify_pr(space_id=space_id, pr_num=pr_num, action="untrusted_pr")
543
 
544