framework/replay: Set Authorization header when downloading from MinIO

This will allow piglit to test with traces hosted in repositories with
restricted access, such as those for which we don't have redistribution
rights.

v2: - Pass the MinIO host name explicitly (Andres)
    - Return credentials from function (Andres)

v3: - Increase credentials timeout to 1 hour (Andres)

Signed-off-by: Tomeu Vizoso <tomeu.vizoso@collabora.com>
Reviewed-by: Andres Gomez <agomez@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/piglit/-/merge_requests/500>

5 files changed, 134 insertions, 3 deletions

diff --git a/framework/replay/download_utils.py b/framework/replay/download_utils.py
index 1253de024..47aca71f2 100644
--- a/framework/replay/download_utils.py
+++ b/framework/replay/download_utils.py
@@ -23,10 +23,16 @@
 #
 # SPDX-License-Identifier: MIT
 
+import base64
+import hashlib
+import hmac
 import requests
+import xml.etree.ElementTree as ET
 
 from os import path
 from time import time
+from email.utils import formatdate
+from urllib.parse import urlparse
 
 from framework import core, exceptions
 from framework.replay.options import OPTIONS
@@ -34,6 +40,71 @@ from framework.replay.options import OPTIONS
 
 __all__ = ['ensure_file']
 
+minio_credentials = None
+
+def sign_with_hmac(key, message):
+    key = key.encode("UTF-8")
+    message = message.encode("UTF-8")
+
+    signature = hmac.new(key, message, hashlib.sha1).digest()
+
+    return base64.encodebytes(signature).strip().decode()
+
+def get_minio_credentials(url):
+    global minio_credentials
+
+    if minio_credentials is not None:
+        return (minio_credentials['AccessKeyId'],
+                minio_credentials['SecretAccessKey'],
+                minio_credentials['SessionToken'])
+
+    minio_credentials = {}
+
+    params = {'Action': 'AssumeRoleWithWebIdentity',
+              'Version': '2011-06-15',
+              'RoleArn': 'arn:aws:iam::123456789012:role/FederatedWebIdentityRole',
+              'RoleSessionName': OPTIONS.download['role_session_name'],
+              'DurationSeconds': 3600,
+              'WebIdentityToken': OPTIONS.download['jwt']}
+    r = requests.post('https://%s' % OPTIONS.download['minio_host'], params=params)
+    if r.status_code >= 400:
+        print(r.text)
+    r.raise_for_status()
+
+    root = ET.fromstring(r.text)
+    for attr in root.iter():
+        if attr.tag == '{https://sts.amazonaws.com/doc/2011-06-15/}AccessKeyId':
+            minio_credentials['AccessKeyId'] = attr.text
+        elif attr.tag == '{https://sts.amazonaws.com/doc/2011-06-15/}SecretAccessKey':
+            minio_credentials['SecretAccessKey'] = attr.text
+        elif attr.tag == '{https://sts.amazonaws.com/doc/2011-06-15/}SessionToken':
+            minio_credentials['SessionToken'] = attr.text
+
+    return (minio_credentials['AccessKeyId'],
+            minio_credentials['SecretAccessKey'],
+            minio_credentials['SessionToken'])
+
+def get_bucket(url):
+    o = urlparse(url)
+    return o.path[1:].split('/')[0]
+
+def get_authorization_headers(url, resource):
+    minio_key, minio_secret, minio_token = get_minio_credentials(url)
+
+    content_type = 'application/octet-stream'
+    date = formatdate(timeval=None, localtime=False, usegmt=True)
+    bucket = get_bucket(url)
+    to_sign = "GET\n\n\n%s\nx-amz-security-token:%s\n/%s/%s" % (date,
+                                                                minio_token,
+                                                                bucket,
+                                                                resource)
+    signature = sign_with_hmac(minio_secret, to_sign)
+
+    headers = {'Host': OPTIONS.download['minio_host'],
+               'Date': date,
+               'Authorization': 'AWS %s:%s' % (minio_key, signature),
+               'x-amz-security-token': minio_token}
+    return headers
 
 def ensure_file(file_path):
     destination_file_path = path.join(OPTIONS.db_path, file_path)
@@ -43,6 +114,8 @@ def ensure_file(file_path):
                 '{} missing'.format(destination_file_path))
         return
 
+    url = OPTIONS.download['url'].geturl()
+
     core.check_dir(path.dirname(destination_file_path))
 
     if not OPTIONS.download['force'] and path.exists(destination_file_path):
@@ -50,10 +123,19 @@ def ensure_file(file_path):
 
     print('[check_image] Downloading file {}'.format(
         file_path), end=' ', flush=True)
+
+    if OPTIONS.download['minio_host']:
+        headers = get_authorization_headers(url, file_path)
+    else:
+        headers = None
+
     download_time = time()
     with open(destination_file_path, 'wb') as file:
-        with requests.get(OPTIONS.download['url'].geturl() + file_path,
-                          allow_redirects=True, stream=True) as r:
+        with requests.get(url + file_path,
+                          allow_redirects=True, stream=True,
+                          headers=headers) as r:
+            if r.status_code >= 400:
+                print(r.text)
             r.raise_for_status()
             for chunk in r.iter_content(chunk_size=8194):
                 if chunk:
diff --git a/framework/replay/options.py b/framework/replay/options.py
index 407ae13b6..a2d5953bc 100644
--- a/framework/replay/options.py
+++ b/framework/replay/options.py
@@ -63,6 +63,9 @@ class _Options(object):  # pylint: disable=too-many-instance-attributes
     download.url -- The URL from which to download the files.
     download.force -- Forces downloading even if the destination file already
                       exists.
+    download.minio_host -- Name of MinIO server from which to download traces
+    download.role_session_name -- Role session name for authentication with MinIO
+    download.jwt -- JWT token for authentication with MinIO
     """
 
     def __init__(self):
@@ -71,7 +74,11 @@ class _Options(object):  # pylint: disable=too-many-instance-attributes
         self.db_path = None
         self.results_path = None
         self.download = {'url': None,
-                         'force': False
+                         'force': False,
+                         'minio': False,
+                         'minio_host': '',
+                         'role_session_name': '',
+                         'jwt': ''
         }
 
     def clear(self):
diff --git a/framework/replay/programs/compare.py b/framework/replay/programs/compare.py
index 6647422a7..6648adc65 100644
--- a/framework/replay/programs/compare.py
+++ b/framework/replay/programs/compare.py
@@ -38,6 +38,9 @@ def _from_yaml(args):
     options.OPTIONS.device_name = args.device_name
     options.OPTIONS.keep_image = args.keep_image
     options.OPTIONS.download['force'] = args.force_download
+    options.OPTIONS.download['minio_host'] = args.download_minio_host
+    options.OPTIONS.download['role_session_name'] = args.download_role_session_name
+    options.OPTIONS.download['jwt'] = args.download_jwt
     options.OPTIONS.db_path = args.db_path
     options.OPTIONS.results_path = args.output
 
@@ -49,6 +52,9 @@ def _trace(args):
     options.OPTIONS.keep_image = args.keep_image
     options.OPTIONS.set_download_url(args.download_url)
     options.OPTIONS.download['force'] = args.force_download
+    options.OPTIONS.download['minio_host'] = args.download_minio_host
+    options.OPTIONS.download['role_session_name'] = args.download_role_session_name
+    options.OPTIONS.download['jwt'] = args.download_jwt
     options.OPTIONS.db_path = args.db_path
     options.OPTIONS.results_path = args.output
 
@@ -77,6 +83,9 @@ def compare(input_):
                  parsers.KEEP_IMAGE,
                  parsers.DOWNLOAD_URL,
                  parsers.DOWNLOAD_FORCE,
+                 parsers.DOWNLOAD_MINIO_HOST,
+                 parsers.DOWNLOAD_ROLE_SESSION_NAME,
+                 parsers.DOWNLOAD_JWT,
                  parsers.DB_PATH,
                  parsers.RESULTS_PATH],
         help=('Compares a specific trace given a checksum and a device.'))
@@ -97,6 +106,9 @@ def compare(input_):
                  parsers.KEEP_IMAGE,
                  parsers.YAML,
                  parsers.DOWNLOAD_FORCE,
+                 parsers.DOWNLOAD_MINIO_HOST,
+                 parsers.DOWNLOAD_ROLE_SESSION_NAME,
+                 parsers.DOWNLOAD_JWT,
                  parsers.DB_PATH,
                  parsers.RESULTS_PATH],
         help=('Compares from a traces description file listing traces '
diff --git a/framework/replay/programs/download.py b/framework/replay/programs/download.py
index 2a7311d81..926353de3 100644
--- a/framework/replay/programs/download.py
+++ b/framework/replay/programs/download.py
@@ -36,6 +36,9 @@ __all__ = ['download']
 def _ensure_file(args):
     options.OPTIONS.set_download_url(args.download_url)
     options.OPTIONS.download['force'] = args.force_download
+    options.OPTIONS.download['minio_host'] = args.download_minio_host
+    options.OPTIONS.download['role_session_name'] = args.download_role_session_name
+    options.OPTIONS.download['jwt'] = args.download_jwt
     options.OPTIONS.db_path = args.db_path
 
     return download_utils.ensure_file(args.file_path)
@@ -46,6 +49,9 @@ def download(input_):
     """ Parser for replayer download command """
     parser = argparse.ArgumentParser(parents=[parsers.DOWNLOAD_URL,
                                               parsers.DOWNLOAD_FORCE,
+                                              parsers.DOWNLOAD_MINIO_HOST,
+                                              parsers.DOWNLOAD_ROLE_SESSION_NAME,
+                                              parsers.DOWNLOAD_JWT,
                                               parsers.DB_PATH])
     parser.add_argument(
         'file_path',
diff --git a/framework/replay/programs/parsers.py b/framework/replay/programs/parsers.py
index 70ada2147..2325d5e84 100644
--- a/framework/replay/programs/parsers.py
+++ b/framework/replay/programs/parsers.py
@@ -67,6 +67,30 @@ DOWNLOAD_FORCE.add_argument(
     help=('forces downloading '
           'even if the destination file already exists'))
 
+DOWNLOAD_MINIO_HOST = argparse.ArgumentParser(add_help=False)
+DOWNLOAD_MINIO_HOST.add_argument(
+    '-m', '--minio_host',
+    dest='download_minio_host',
+    required=False,
+    default=None,
+    help=('name of MinIO server from which to download traces'))
+
+DOWNLOAD_ROLE_SESSION_NAME = argparse.ArgumentParser(add_help=False)
+DOWNLOAD_ROLE_SESSION_NAME.add_argument(
+    '-r', '--role-session-name',
+    dest='download_role_session_name',
+    required=False,
+    default=None,
+    help=('role session name for authentication with MinIO'))
+
+DOWNLOAD_JWT = argparse.ArgumentParser(add_help=False)
+DOWNLOAD_JWT.add_argument(
+    '-j', '--jwt',
+    dest='download_jwt',
+    required=False,
+    default=None,
+    help=('JWT token for authentication with MinIO'))
+
 DB_PATH = argparse.ArgumentParser(add_help=False)
 DB_PATH.add_argument(
     '-p', '--db-path',