build_tools/benchmarks/post_benchmark_comment.py - 3p/openxla/iree - Git at Google

 #!/usr/bin/env python3
 # Copyright 2022 The IREE Authors
 #
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 """
 Posts benchmark results to gist and comments on pull requests.

 Requires the environment variables:

 - GITHUB_TOKEN: token from GitHub action that has write access on issues. See
   https://docs.github.com/en/actions/security-guides/automatic-token-authentication#permissions-for-the-github_token
 - COMMENT_BOT_USER: user name that posts the comment. Note this can be different
     from the user creates the gist.
 - GIST_BOT_TOKEN: token that has write access to gist. Gist will be posted as
   the owner of the token. See
   https://docs.github.com/en/rest/overview/permissions-required-for-fine-grained-personal-access-tokens#gists
 """

 import sys
 import pathlib

 # Add build_tools python dir to the search path.
 sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

 import argparse
 import http.client
 import json
 import os
 import requests
 from typing import Any, Optional

 from reporting import benchmark_comment

 GITHUB_IREE_API_PREFIX = "https://api.github.com/repos/openxla/iree"
 GITHUB_GIST_API = "https://api.github.com/gists"
 GITHUB_API_VERSION = "2022-11-28"


 class APIRequester(object):
   """REST API client that injects proper GitHub authentication headers."""

   def __init__(self, github_token: str):
     self._api_headers = {
         "Accept": "application/vnd.github+json",
         "Authorization": f"token {github_token}",
         "X-GitHub-Api-Version": GITHUB_API_VERSION,
     }
     self._session = requests.session()

   def get(self, endpoint: str, payload: Any = {}) -> requests.Response:
     return self._session.get(endpoint,
                              data=json.dumps(payload),
                              headers=self._api_headers)

   def post(self, endpoint: str, payload: Any = {}) -> requests.Response:
     return self._session.post(endpoint,
                               data=json.dumps(payload),
                               headers=self._api_headers)

   def patch(self, endpoint: str, payload: Any = {}) -> requests.Response:
     return self._session.patch(endpoint,
                                data=json.dumps(payload),
                                headers=self._api_headers)


 class GithubClient(object):
   """Helper to call Github REST APIs."""

   def __init__(self, requester: APIRequester):
     self._requester = requester

   def post_to_gist(self,
                    filename: str,
                    content: str,
                    verbose: bool = False) -> str:
     """Posts the given content to a new GitHub Gist and returns the URL to it."""

     response = self._requester.post(endpoint=GITHUB_GIST_API,
                                     payload={
                                         "public": True,
                                         "files": {
                                             filename: {
                                                 "content": content
                                             }
                                         }
                                     })
     if response.status_code != http.client.CREATED:
       raise RuntimeError(
           f"Failed to create on gist; error code: {response.status_code} - {response.text}"
       )

     response = response.json()
     if verbose:
       print(f"Gist posting response: {response}")

     if response["truncated"]:
       raise RuntimeError(f"Content is too large and was truncated")

     return response["html_url"]

   def get_previous_comment_on_pr(self,
                                  pr_number: int,
                                  comment_bot_user: str,
                                  comment_type_id: str,
                                  query_comment_per_page: int = 100,
                                  max_pages_to_search: int = 10,
                                  verbose: bool = False) -> Optional[int]:
     """Gets the previous comment's id from GitHub."""

     for page in range(1, max_pages_to_search + 1):
       response = self._requester.get(
           endpoint=f"{GITHUB_IREE_API_PREFIX}/issues/{pr_number}/comments",
           payload={
               "per_page": query_comment_per_page,
               "page": page,
               "sort": "updated",
               "direction": "desc"
           })
       if response.status_code != http.client.OK:
         raise RuntimeError(
             f"Failed to get PR comments from GitHub; error code: {response.status_code} - {response.text}"
         )

       comments = response.json()
       if verbose:
         print(f"Previous comment query response on page {page}: {comments}")

       # Find the most recently updated comment that matches.
       for comment in comments:
         if (comment["user"]["login"] == comment_bot_user and
             comment_type_id in comment["body"]):
           return comment["id"]

       if len(comments) < query_comment_per_page:
         break

     return None

   def update_comment_on_pr(self, comment_id: int, content: str):
     """Updates the content of the given comment id."""

     response = self._requester.patch(
         endpoint=f"{GITHUB_IREE_API_PREFIX}/issues/comments/{comment_id}",
         payload={"body": content})
     if response.status_code != http.client.OK:
       raise RuntimeError(
           f"Failed to comment on GitHub; error code: {response.status_code} - {response.text}"
       )

   def create_comment_on_pr(self, pr_number: int, content: str):
     """Posts the given content as comments to the current pull request."""

     response = self._requester.post(
         endpoint=f"{GITHUB_IREE_API_PREFIX}/issues/{pr_number}/comments",
         payload={"body": content})
     if response.status_code != http.client.CREATED:
       raise RuntimeError(
           f"Failed to comment on GitHub; error code: {response.status_code} - {response.text}"
       )

   def get_pull_request_head_commit(self, pr_number: int) -> str:
     """Get pull request head commit SHA."""

     response = self._requester.get(
         endpoint=f"{GITHUB_IREE_API_PREFIX}/pulls/{pr_number}")
     if response.status_code != http.client.OK:
       raise RuntimeError(
           f"Failed to fetch the pull request: {pr_number}; "
           f"error code: {response.status_code} - {response.text}")

     return response.json()["head"]["sha"]


 def _parse_arguments():
   parser = argparse.ArgumentParser()
   parser.add_argument("comment_json", type=pathlib.Path)
   parser.add_argument("--verbose", action="store_true")
   verification_parser = parser.add_mutually_exclusive_group(required=True)
   verification_parser.add_argument("--github_event_json", type=pathlib.Path)
   # Temporary option for buildkite pipeline.
   verification_parser.add_argument("--no_verify_pr", action="store_true")
   return parser.parse_args()


 def main(args: argparse.Namespace):
   github_token = os.environ.get("GITHUB_TOKEN")
   if github_token is None:
     raise ValueError("GITHUB_TOKEN must be set.")

   comment_bot_user = os.environ.get("COMMENT_BOT_USER")
   if comment_bot_user is None:
     raise ValueError("COMMENT_BOT_USER must be set.")

   gist_bot_token = os.environ.get("GIST_BOT_TOKEN")
   if gist_bot_token is None:
     raise ValueError("GIST_BOT_TOKEN must be set.")

   comment_data = benchmark_comment.CommentData(
       **json.loads(args.comment_json.read_text()))
   # Sanitize the pr number to make sure it is an integer.
   pr_number = int(comment_data.unverified_pr_number)

   pr_client = GithubClient(requester=APIRequester(github_token=github_token))
   if args.github_event_json is None:
     github_event = None
   else:
     github_event = json.loads(args.github_event_json.read_text())
     workflow_run_sha = github_event["workflow_run"]["head_sha"]
     pr_head_sha = pr_client.get_pull_request_head_commit(pr_number=pr_number)
     # We can't get the trusted PR number of a workflow run from GitHub API. So we
     # take the untrusted PR number from presubmit workflow and verify if the PR's
     # current head SHA matches the commit SHA in the workflow run. It assumes
     # that to generate the malicious comment data, attacker must modify the code
     # and has a new commit SHA. So if the PR head commit matches the workflow
     # run with attacker's commit, either the PR is created by the attacker or
     # other's PR has the malicious commit. In both cases posting malicious
     # comment is acceptable.
     #
     # Note that the collision of a target SHA1 is possible but GitHub has some
     # protections (https://github.blog/2017-03-20-sha-1-collision-detection-on-github-com/).
     # The assumption also only holds if files in GCS can't be overwritten (so the
     # comment data can't be modified without changing the code).
     # The check will also fail if the PR author pushes the new commit after the
     # workflow is triggered. But pushing the new commit means to cancel the
     # current CI run including the benchmarking. So it will unlikely fail for
     # that reason.
     if workflow_run_sha != pr_head_sha:
       raise ValueError(
           f"Workflow run SHA: {workflow_run_sha} does not match "
           f"the head SHA: {pr_head_sha} of the pull request: {pr_number}.")

   gist_client = GithubClient(requester=APIRequester(
       github_token=gist_bot_token))
   gist_url = gist_client.post_to_gist(
       filename=f'iree-full-benchmark-results-{pr_number}.md',
       content=comment_data.full_md,
       verbose=args.verbose)

   previous_comment_id = pr_client.get_previous_comment_on_pr(
       pr_number=pr_number,
       comment_bot_user=comment_bot_user,
       comment_type_id=comment_data.type_id,
       verbose=args.verbose)

   abbr_md = comment_data.abbr_md.replace(
       benchmark_comment.GIST_LINK_PLACEHORDER, gist_url)
   if github_event is not None:
     abbr_md += f'\n\n[Source Workflow Run]({github_event["workflow_run"]["html_url"]})'
   if previous_comment_id is not None:
     pr_client.update_comment_on_pr(comment_id=previous_comment_id,
                                    content=abbr_md)
   else:
     pr_client.create_comment_on_pr(pr_number=pr_number, content=abbr_md)


 if __name__ == "__main__":
   main(_parse_arguments())
	#!/usr/bin/env python3
	# Copyright 2022 The IREE Authors
	#
	# Licensed under the Apache License v2.0 with LLVM Exceptions.
	# See https://llvm.org/LICENSE.txt for license information.
	# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	"""
	Posts benchmark results to gist and comments on pull requests.

	Requires the environment variables:

	- GITHUB_TOKEN: token from GitHub action that has write access on issues. See
	https://docs.github.com/en/actions/security-guides/automatic-token-authentication#permissions-for-the-github_token
	- COMMENT_BOT_USER: user name that posts the comment. Note this can be different
	from the user creates the gist.
	- GIST_BOT_TOKEN: token that has write access to gist. Gist will be posted as
	the owner of the token. See
	https://docs.github.com/en/rest/overview/permissions-required-for-fine-grained-personal-access-tokens#gists
	"""

	import sys
	import pathlib

	# Add build_tools python dir to the search path.
	sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

	import argparse
	import http.client
	import json
	import os
	import requests
	from typing import Any, Optional

	from reporting import benchmark_comment

	GITHUB_IREE_API_PREFIX = "https://api.github.com/repos/openxla/iree"
	GITHUB_GIST_API = "https://api.github.com/gists"
	GITHUB_API_VERSION = "2022-11-28"


	class APIRequester(object):
	"""REST API client that injects proper GitHub authentication headers."""

	def __init__(self, github_token: str):
	self._api_headers = {
	"Accept": "application/vnd.github+json",
	"Authorization": f"token {github_token}",
	"X-GitHub-Api-Version": GITHUB_API_VERSION,
	}
	self._session = requests.session()

	def get(self, endpoint: str, payload: Any = {}) -> requests.Response:
	return self._session.get(endpoint,
	data=json.dumps(payload),
	headers=self._api_headers)

	def post(self, endpoint: str, payload: Any = {}) -> requests.Response:
	return self._session.post(endpoint,
	data=json.dumps(payload),
	headers=self._api_headers)

	def patch(self, endpoint: str, payload: Any = {}) -> requests.Response:
	return self._session.patch(endpoint,
	data=json.dumps(payload),
	headers=self._api_headers)


	class GithubClient(object):
	"""Helper to call Github REST APIs."""

	def __init__(self, requester: APIRequester):
	self._requester = requester

	def post_to_gist(self,
	filename: str,
	content: str,
	verbose: bool = False) -> str:
	"""Posts the given content to a new GitHub Gist and returns the URL to it."""

	response = self._requester.post(endpoint=GITHUB_GIST_API,
	payload={
	"public": True,
	"files": {
	filename: {
	"content": content
	}
	}
	})
	if response.status_code != http.client.CREATED:
	raise RuntimeError(
	f"Failed to create on gist; error code: {response.status_code} - {response.text}"
	)

	response = response.json()
	if verbose:
	print(f"Gist posting response: {response}")

	if response["truncated"]:
	raise RuntimeError(f"Content is too large and was truncated")

	return response["html_url"]

	def get_previous_comment_on_pr(self,
	pr_number: int,
	comment_bot_user: str,
	comment_type_id: str,
	query_comment_per_page: int = 100,
	max_pages_to_search: int = 10,
	verbose: bool = False) -> Optional[int]:
	"""Gets the previous comment's id from GitHub."""

	for page in range(1, max_pages_to_search + 1):
	response = self._requester.get(
	endpoint=f"{GITHUB_IREE_API_PREFIX}/issues/{pr_number}/comments",
	payload={
	"per_page": query_comment_per_page,
	"page": page,
	"sort": "updated",
	"direction": "desc"
	})
	if response.status_code != http.client.OK:
	raise RuntimeError(
	f"Failed to get PR comments from GitHub; error code: {response.status_code} - {response.text}"
	)

	comments = response.json()
	if verbose:
	print(f"Previous comment query response on page {page}: {comments}")

	# Find the most recently updated comment that matches.
	for comment in comments:
	if (comment["user"]["login"] == comment_bot_user and
	comment_type_id in comment["body"]):
	return comment["id"]

	if len(comments) < query_comment_per_page:
	break

	return None

	def update_comment_on_pr(self, comment_id: int, content: str):
	"""Updates the content of the given comment id."""

	response = self._requester.patch(
	endpoint=f"{GITHUB_IREE_API_PREFIX}/issues/comments/{comment_id}",
	payload={"body": content})
	if response.status_code != http.client.OK:
	raise RuntimeError(
	f"Failed to comment on GitHub; error code: {response.status_code} - {response.text}"
	)

	def create_comment_on_pr(self, pr_number: int, content: str):
	"""Posts the given content as comments to the current pull request."""

	response = self._requester.post(
	endpoint=f"{GITHUB_IREE_API_PREFIX}/issues/{pr_number}/comments",
	payload={"body": content})
	if response.status_code != http.client.CREATED:
	raise RuntimeError(
	f"Failed to comment on GitHub; error code: {response.status_code} - {response.text}"
	)

	def get_pull_request_head_commit(self, pr_number: int) -> str:
	"""Get pull request head commit SHA."""

	response = self._requester.get(
	endpoint=f"{GITHUB_IREE_API_PREFIX}/pulls/{pr_number}")
	if response.status_code != http.client.OK:
	raise RuntimeError(
	f"Failed to fetch the pull request: {pr_number}; "
	f"error code: {response.status_code} - {response.text}")

	return response.json()["head"]["sha"]


	def _parse_arguments():
	parser = argparse.ArgumentParser()
	parser.add_argument("comment_json", type=pathlib.Path)
	parser.add_argument("--verbose", action="store_true")
	verification_parser = parser.add_mutually_exclusive_group(required=True)
	verification_parser.add_argument("--github_event_json", type=pathlib.Path)
	# Temporary option for buildkite pipeline.
	verification_parser.add_argument("--no_verify_pr", action="store_true")
	return parser.parse_args()


	def main(args: argparse.Namespace):
	github_token = os.environ.get("GITHUB_TOKEN")
	if github_token is None:
	raise ValueError("GITHUB_TOKEN must be set.")

	comment_bot_user = os.environ.get("COMMENT_BOT_USER")
	if comment_bot_user is None:
	raise ValueError("COMMENT_BOT_USER must be set.")

	gist_bot_token = os.environ.get("GIST_BOT_TOKEN")
	if gist_bot_token is None:
	raise ValueError("GIST_BOT_TOKEN must be set.")

	comment_data = benchmark_comment.CommentData(
	**json.loads(args.comment_json.read_text()))
	# Sanitize the pr number to make sure it is an integer.
	pr_number = int(comment_data.unverified_pr_number)

	pr_client = GithubClient(requester=APIRequester(github_token=github_token))
	if args.github_event_json is None:
	github_event = None
	else:
	github_event = json.loads(args.github_event_json.read_text())
	workflow_run_sha = github_event["workflow_run"]["head_sha"]
	pr_head_sha = pr_client.get_pull_request_head_commit(pr_number=pr_number)
	# We can't get the trusted PR number of a workflow run from GitHub API. So we
	# take the untrusted PR number from presubmit workflow and verify if the PR's
	# current head SHA matches the commit SHA in the workflow run. It assumes
	# that to generate the malicious comment data, attacker must modify the code
	# and has a new commit SHA. So if the PR head commit matches the workflow
	# run with attacker's commit, either the PR is created by the attacker or
	# other's PR has the malicious commit. In both cases posting malicious
	# comment is acceptable.
	#
	# Note that the collision of a target SHA1 is possible but GitHub has some
	# protections (https://github.blog/2017-03-20-sha-1-collision-detection-on-github-com/).
	# The assumption also only holds if files in GCS can't be overwritten (so the
	# comment data can't be modified without changing the code).
	# The check will also fail if the PR author pushes the new commit after the
	# workflow is triggered. But pushing the new commit means to cancel the
	# current CI run including the benchmarking. So it will unlikely fail for
	# that reason.
	if workflow_run_sha != pr_head_sha:
	raise ValueError(
	f"Workflow run SHA: {workflow_run_sha} does not match "
	f"the head SHA: {pr_head_sha} of the pull request: {pr_number}.")

	gist_client = GithubClient(requester=APIRequester(
	github_token=gist_bot_token))
	gist_url = gist_client.post_to_gist(
	filename=f'iree-full-benchmark-results-{pr_number}.md',
	content=comment_data.full_md,
	verbose=args.verbose)

	previous_comment_id = pr_client.get_previous_comment_on_pr(
	pr_number=pr_number,
	comment_bot_user=comment_bot_user,
	comment_type_id=comment_data.type_id,
	verbose=args.verbose)

	abbr_md = comment_data.abbr_md.replace(
	benchmark_comment.GIST_LINK_PLACEHORDER, gist_url)
	if github_event is not None:
	abbr_md += f'\n\n[Source Workflow Run]({github_event["workflow_run"]["html_url"]})'
	if previous_comment_id is not None:
	pr_client.update_comment_on_pr(comment_id=previous_comment_id,
	content=abbr_md)
	else:
	pr_client.create_comment_on_pr(pr_number=pr_number, content=abbr_md)


	if __name__ == "__main__":
	main(_parse_arguments())