build_tools/benchmarks/post_benchmark_comment.py

#!/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())