[CIFuzz] Implement filestore based on github artifacts (#5943)

Implement filestore based on github actions' artifacts feature.
This uses the github api and the github actions API.
Also fix imports in github_actions_toolkit library that were broken by move to third_party directory.

11 files changed, 399 insertions, 10 deletions

diff --git a/infra/cifuzz/filestore/__init__.py b/infra/cifuzz/filestore/__init__.py
index ab0fe6e0..0acb4b9c 100644
--- a/infra/cifuzz/filestore/__init__.py
+++ b/infra/cifuzz/filestore/__init__.py
@@ -14,6 +14,10 @@
 """Module for a generic filestore."""
 
 
+class FilestoreError(Exception):
+  """Error using the filestore."""
+
+
 # pylint: disable=unused-argument,no-self-use
 class BaseFilestore:
   """Base class for a filestore."""
@@ -21,10 +25,14 @@ class BaseFilestore:
   def __init__(self, config):
     self.config = config
 
-  def upload_corpus(self, name, directory):
-    """Uploads the corpus located at |directory| to |name|."""
+  def upload_directory(self, name, directory):
+    """Uploads the |directory| to |name|."""
     raise NotImplementedError('Child class must implement method.')
 
   def download_corpus(self, name, dst_directory):
     """Downloads the corpus located at |name| to |dst_directory|."""
     raise NotImplementedError('Child class must implement method.')
+
+  def download_latest_build(self, name, dst_directory):
+    """Downloads the latest build with |name| to |dst_directory|."""
+    raise NotImplementedError('Child class must implement method.')
diff --git a/infra/cifuzz/filestore/github_actions/__init__.py b/infra/cifuzz/filestore/github_actions/__init__.py
new file mode 100644
index 00000000..fbcbb1d7
--- /dev/null
+++ b/infra/cifuzz/filestore/github_actions/__init__.py
@@ -0,0 +1,83 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of a filestore using Github actions artifacts."""
+import os
+import logging
+
+import http_utils
+import filestore
+from filestore.github_actions import github_api
+from third_party.github_actions_toolkit.artifact import artifact_client
+
+
+class GithubActionsFilestore(filestore.BaseFilestore):
+  """Implementation of BaseFilestore using Github actions artifacts. Relies on
+  github_actions_toolkit for using the GitHub actions API and the github_api
+  module for using GitHub's standard API. We need to use both because the GitHub
+  actions API is the only way to upload an artifact but it does not support
+  downloading artifacts from other runs. The standard GitHub API does support
+  this however."""
+
+  def __init__(self, config):
+    super().__init__(config)
+    self.github_api_http_headers = github_api.get_http_auth_headers(config)
+
+  def upload_directory(self, name, directory):  # pylint: disable=no-self-use
+    """Uploads |directory| as artifact with |name|."""
+    directory = os.path.abspath(directory)
+
+    # Get file paths.
+    file_paths = []
+    for root, _, curr_file_paths in os.walk(directory):
+      for file_path in curr_file_paths:
+        file_paths.append(os.path.join(root, file_path))
+
+    logging.debug('file_paths: %s', file_paths)
+
+    # TODO(metzman): Zip so that we can upload directories within directories
+    # and save time?
+
+    return artifact_client.upload_artifact(name, file_paths, directory)
+
+  def download_corpus(self, name, dst_directory):  # pylint: disable=unused-argument,no-self-use
+    """Downloads the corpus located at |name| to |dst_directory|."""
+    return self._download_artifact(name, dst_directory)
+
+  def _find_artifact(self, name):
+    """Finds an artifact using the GitHub API and returns it."""
+    logging.debug('listing artifact')
+    artifacts = self._list_artifacts()
+    artifact = github_api.find_artifact(name, artifacts)
+    logging.debug('Artifact: %s.', artifact)
+    return artifact
+
+  def _download_artifact(self, name, dst_directory):
+    """Downloads artifact with |name| to |dst_directory|."""
+    artifact = self._find_artifact(name)
+    if not artifact:
+      logging.warning('Could not download artifact: %s.', name)
+      return artifact
+    download_url = artifact['archive_download_url']
+    return http_utils.download_and_unpack_zip(
+        download_url, dst_directory, headers=self.github_api_http_headers)
+
+  def _list_artifacts(self):
+    """Returns a list of artifacts."""
+    return github_api.list_artifacts(self.config.project_repo_owner,
+                                     self.config.project_repo_name,
+                                     self.github_api_http_headers)
+
+  def download_latest_build(self, name, dst_directory):
+    """Downloads latest build with name |name| to |dst_directory|."""
+    return self._download_artifact(name, dst_directory)
diff --git a/infra/cifuzz/filestore/github_actions/github_actions_test.py b/infra/cifuzz/filestore/github_actions/github_actions_test.py
new file mode 100644
index 00000000..3204c151
--- /dev/null
+++ b/infra/cifuzz/filestore/github_actions/github_actions_test.py
@@ -0,0 +1,86 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for github_actions."""
+import os
+import sys
+import unittest
+from unittest import mock
+
+# pylint: disable=wrong-import-position
+INFRA_DIR = os.path.dirname(
+    os.path.dirname(os.path.dirname(os.path.dirname(
+        os.path.abspath(__file__)))))
+sys.path.append(INFRA_DIR)
+
+from filestore import github_actions
+import test_helpers
+
+# pylint: disable=protected-access,no-self-use
+
+
+class GithubActionsFilestoreTest(unittest.TestCase):
+  """Tests for GithubActionsFilestore."""
+
+  def setUp(self):
+    test_helpers.patch_environ(self)
+    self.github_token = 'example githubtoken'
+
+  def _get_expected_http_headers(self):
+    return {
+        'Authorization': 'token {token}'.format(token=self.github_token),
+        'Accept': 'application/vnd.github.v3+json',
+    }
+
+  @mock.patch('filestore.github_actions.github_api.list_artifacts')
+  def test_list_artifacts(self, mocked_list_artifacts):
+    """Tests that _list_artifacts works as intended."""
+    owner = 'exampleowner'
+    repo = 'examplerepo'
+    os.environ['GITHUB_REPOSITORY'] = '{owner}/{repo}'.format(owner=owner,
+                                                              repo=repo)
+    config = test_helpers.create_run_config(github_token=self.github_token)
+    filestore = github_actions.GithubActionsFilestore(config)
+    filestore._list_artifacts()
+    mocked_list_artifacts.assert_called_with(owner, repo,
+                                             self._get_expected_http_headers())
+
+  @mock.patch('logging.warning')
+  @mock.patch('filestore.github_actions.GithubActionsFilestore._list_artifacts',
+              return_value=None)
+  @mock.patch('filestore.github_actions.github_api.find_artifact',
+              return_value=None)
+  def test_download_latest_build_no_artifact(self, _, __, mocked_warning):
+    """Tests that download_latest_build returns None and doesn't exception when
+    find_artifact can't find an artifact."""
+    config = test_helpers.create_run_config(github_token=self.github_token)
+    filestore = github_actions.GithubActionsFilestore(config)
+    name = 'build-name'
+    build_dir = 'build-dir'
+    self.assertIsNone(filestore.download_latest_build(name, build_dir))
+    mocked_warning.assert_called_with('Could not download artifact: %s.', name)
+
+  @mock.patch('logging.warning')
+  @mock.patch('filestore.github_actions.GithubActionsFilestore._list_artifacts',
+              return_value=None)
+  @mock.patch('filestore.github_actions.github_api.find_artifact',
+              return_value=None)
+  def test_download_corpus_no_artifact(self, _, __, mocked_warning):
+    """Tests that download_corpus_build returns None and doesn't exception when
+    find_artifact can't find an artifact."""
+    config = test_helpers.create_run_config(github_token=self.github_token)
+    filestore = github_actions.GithubActionsFilestore(config)
+    name = 'corpus-name'
+    dst_dir = 'corpus-dir'
+    self.assertFalse(filestore.download_corpus(name, dst_dir))
+    mocked_warning.assert_called_with('Could not download artifact: %s.', name)
diff --git a/infra/cifuzz/filestore/github_actions/github_api.py b/infra/cifuzz/filestore/github_actions/github_api.py
new file mode 100644
index 00000000..32e1b392
--- /dev/null
+++ b/infra/cifuzz/filestore/github_actions/github_api.py
@@ -0,0 +1,106 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for dealing with the GitHub API. This is different from
+github_actions_toolkit which only deals with the actions API. We need to use
+both."""
+import logging
+import os
+import sys
+
+import requests
+
+# pylint: disable=wrong-import-position,import-error
+
+sys.path.append(
+    os.path.join(__file__, os.path.pardir, os.path.pardir, os.path.pardir,
+                 os.path.pardir))
+import retry
+
+_MAX_ITEMS_PER_PAGE = 100
+
+_GET_ATTEMPTS = 3
+_GET_BACKOFF = 1
+
+
+def get_http_auth_headers(config):
+  """Returns HTTP headers for authentication to the API."""
+  authorization = 'token {token}'.format(token=config.github_token)
+  return {
+      'Authorization': authorization,
+      'Accept': 'application/vnd.github.v3+json'
+  }
+
+
+def _get_artifacts_list_api_url(repo_owner, repo_name):
+  """Returns the artifacts_api_url for |repo_name| owned by |repo_owner|."""
+  return (f'https://api.github.com/repos/{repo_owner}/'
+          f'{repo_name}/actions/artifacts')
+
+
+@retry.wrap(_GET_ATTEMPTS, _GET_BACKOFF)
+def _do_get_request(*args, **kwargs):
+  """Wrapped version of requests.get that does retries."""
+  return requests.get(*args, **kwargs)
+
+
+def _get_items(url, headers):
+  """Generator that gets and yields items from a GitHub API endpoint (specified
+  by |URL|) sending |headers| with the get request."""
+  # Github API response pages are 1-indexed.
+  page_counter = 1
+
+  # Set to infinity so we run loop at least once.
+  total_num_items = float('inf')
+
+  item_num = 0
+  while item_num < total_num_items:
+    params = {'per_page': _MAX_ITEMS_PER_PAGE, 'page': str(page_counter)}
+    response = _do_get_request(url, params=params, headers=headers)
+    response_json = response.json()
+
+    if not response.status_code == 200:
+      # Check that request was successful.
+      logging.error('Request to %s failed. Code: %d. Response: %s',
+                    response.request.url, response.status_code, response_json)
+
+    if total_num_items == float('inf'):
+      # Set proper total_num_items
+      total_num_items = response_json['total_count']
+
+    # Get the key for the items we are after.
+    keys = [key for key in response_json.keys() if key != 'total_count']
+    assert len(keys) == 1, keys
+    items_key = keys[0]
+
+    for item in response_json[items_key]:
+      yield item
+      item_num += 1
+
+    page_counter += 1
+
+
+def find_artifact(artifact_name, artifacts):
+  """Find the artifact with the name |artifact_name| in |artifacts|."""
+  for artifact in artifacts:
+    # TODO(metzman): Handle multiple by making sure we download the latest.
+    if artifact['name'] == artifact_name and not artifact['expired']:
+      return artifact
+  return None
+
+
+def list_artifacts(owner, repo, headers):
+  """Returns a generator of all the artifacts for |owner/repo|."""
+  url = _get_artifacts_list_api_url(owner, repo)
+  logging.debug('Getting artifacts from: %s', url)
+  return _get_items(url, headers)
diff --git a/infra/cifuzz/filestore/github_actions/github_api_test.py b/infra/cifuzz/filestore/github_actions/github_api_test.py
new file mode 100644
index 00000000..adf7cf4f
--- /dev/null
+++ b/infra/cifuzz/filestore/github_actions/github_api_test.py
@@ -0,0 +1,33 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for github_api."""
+import unittest
+
+from filestore.github_actions import github_api
+import test_helpers
+
+
+class GetHttpAuthHeaders(unittest.TestCase):
+  """Tests for get_http_auth_headers."""
+
+  def test_get_http_auth_headers(self):
+    """Tests that get_http_auth_headers returns the correct result."""
+    github_token = 'example githubtoken'
+    run_config = test_helpers.create_run_config(github_token=github_token)
+    expected_headers = {
+        'Authorization': 'token {token}'.format(token=github_token),
+        'Accept': 'application/vnd.github.v3+json',
+    }
+    self.assertEqual(expected_headers,
+                     github_api.get_http_auth_headers(run_config))
diff --git a/infra/cifuzz/filestore_utils.py b/infra/cifuzz/filestore_utils.py
new file mode 100644
index 00000000..d100b851
--- /dev/null
+++ b/infra/cifuzz/filestore_utils.py
@@ -0,0 +1,25 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""External filestore interface. Cannot be depended on by filestore code."""
+import filestore
+import filestore.github_actions
+
+
+def get_filestore(config):
+  """Returns the correct filestore based on the platform in |config|.
+  Raises an exception if there is no correct filestore for the platform."""
+  # TODO(metzman): Force specifying of filestore.
+  if config.platform == config.Platform.EXTERNAL_GITHUB:
+    return filestore.github_actions.GithubActionsFilestore(config)
+  raise filestore.FilestoreError('Filestore doesn\'t support platform.')
diff --git a/infra/cifuzz/filestore_utils_test.py b/infra/cifuzz/filestore_utils_test.py
new file mode 100644
index 00000000..cb5e0d6e
--- /dev/null
+++ b/infra/cifuzz/filestore_utils_test.py
@@ -0,0 +1,48 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for filestore_utils."""
+import unittest
+from unittest import mock
+
+import parameterized
+
+import filestore
+from filestore import github_actions
+import filestore_utils
+import test_helpers
+
+
+class GetFilestoreTest(unittest.TestCase):
+  """Tests for get_filestore."""
+
+  @parameterized.parameterized.expand([
+      ({
+          'build_integration_path': '/',
+          'is_github': True,
+      }, github_actions.GithubActionsFilestore),
+  ])
+  def test_get_filestore(self, config_kwargs, filestore_cls):
+    """Tests that get_filestore returns the right filestore given a certain
+    platform."""
+    run_config = test_helpers.create_run_config(**config_kwargs)
+    filestore_impl = filestore_utils.get_filestore(run_config)
+    self.assertIsInstance(filestore_impl, filestore_cls)
+
+  def test_get_filestore_unsupported_platform(self):
+    """Tests that get_filestore exceptions given a platform it doesn't
+    support."""
+    with mock.patch('config_utils.BaseConfig.platform', return_value='other'):
+      run_config = test_helpers.create_run_config()
+      with self.assertRaises(filestore.FilestoreError):
+        filestore_utils.get_filestore(run_config)
diff --git a/infra/cifuzz/third_party/github_actions_toolkit/artifact/artifact_client.py b/infra/cifuzz/third_party/github_actions_toolkit/artifact/artifact_client.py
index 626b0cb3..ae5b3ab0 100644
--- a/infra/cifuzz/third_party/github_actions_toolkit/artifact/artifact_client.py
+++ b/infra/cifuzz/third_party/github_actions_toolkit/artifact/artifact_client.py
@@ -1,9 +1,9 @@
 """Public interface for artifact. Based on artifact-client.ts"""
 import logging
 
-from github_actions_toolkit.artifact import utils
-from github_actions_toolkit.artifact import upload_http_client
-from github_actions_toolkit.artifact import upload_specification
+from third_party.github_actions_toolkit.artifact import utils
+from third_party.github_actions_toolkit.artifact import upload_http_client
+from third_party.github_actions_toolkit.artifact import upload_specification
 
 
 def upload_artifact(name, files, root_directory, options=None):
diff --git a/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_http_client.py b/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_http_client.py
index 527f57c9..6e9ad932 100644
--- a/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_http_client.py
+++ b/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_http_client.py
@@ -10,9 +10,9 @@ import urllib.request
 
 import requests
 
-from github_actions_toolkit.artifact import config_variables
-from github_actions_toolkit.artifact import utils
-from github_actions_toolkit import http_client
+from third_party.github_actions_toolkit.artifact import config_variables
+from third_party.github_actions_toolkit.artifact import utils
+from third_party.github_actions_toolkit import http_client
 
 
 def upload_file(parameters):
diff --git a/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_specification.py b/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_specification.py
index 71690e12..0243e02d 100644
--- a/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_specification.py
+++ b/infra/cifuzz/third_party/github_actions_toolkit/artifact/upload_specification.py
@@ -2,7 +2,7 @@
 import logging
 import os
 
-from github_actions_toolkit.artifact import utils
+from third_party.github_actions_toolkit.artifact import utils
 
 
 class UploadSpecification:  # pylint: disable=too-few-public-methods
diff --git a/infra/cifuzz/third_party/github_actions_toolkit/artifact/utils.py b/infra/cifuzz/third_party/github_actions_toolkit/artifact/utils.py
index 6ec62969..67f404ee 100644
--- a/infra/cifuzz/third_party/github_actions_toolkit/artifact/utils.py
+++ b/infra/cifuzz/third_party/github_actions_toolkit/artifact/utils.py
@@ -1,7 +1,7 @@
 """Utility module. Based on utils.ts."""
 import logging
 
-from github_actions_toolkit.artifact import config_variables
+from third_party.github_actions_toolkit.artifact import config_variables
 
 MAX_API_ATTEMPTS = 5
 SLEEP_TIME = 1