framework/replay: Allow specifying the MinIO bucket containing the traces

Previously we were inferring it from the URL, but this becomes more
complicated if the download URL isn't that of the MinIO server (eg.
there's a proxy in the middle).

Signed-off-by: Tomeu Vizoso <tomeu.vizoso@collabora.com>
Reviewed-by: Rohan Garg <rohan.garg@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/piglit/-/merge_requests/509>

5 files changed, 20 insertions, 7 deletions

diff --git a/framework/replay/download_utils.py b/framework/replay/download_utils.py
index 47aca71f2..3bc682600 100644
--- a/framework/replay/download_utils.py
+++ b/framework/replay/download_utils.py
@@ -84,19 +84,14 @@ def get_minio_credentials(url):
             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,
+                                                                OPTIONS.download['minio_bucket'],
                                                                 resource)
     signature = sign_with_hmac(minio_secret, to_sign)
 
@@ -125,6 +120,9 @@ def ensure_file(file_path):
         file_path), end=' ', flush=True)
 
     if OPTIONS.download['minio_host']:
+        assert OPTIONS.download['minio_bucket']
+        assert OPTIONS.download['role_session_name']
+        assert OPTIONS.download['jwt']
         headers = get_authorization_headers(url, file_path)
     else:
         headers = None
diff --git a/framework/replay/options.py b/framework/replay/options.py
index a2d5953bc..084d46d36 100644
--- a/framework/replay/options.py
+++ b/framework/replay/options.py
@@ -64,6 +64,7 @@ class _Options(object):  # pylint: disable=too-many-instance-attributes
     download.force -- Forces downloading even if the destination file already
                       exists.
     download.minio_host -- Name of MinIO server from which to download traces
+    download.minio_bucket -- Name of bucket in MinIO server containing the traces
     download.role_session_name -- Role session name for authentication with MinIO
     download.jwt -- JWT token for authentication with MinIO
     """
@@ -75,8 +76,8 @@ class _Options(object):  # pylint: disable=too-many-instance-attributes
         self.results_path = None
         self.download = {'url': None,
                          'force': False,
-                         'minio': False,
                          'minio_host': '',
+                         'minio_bucket': '',
                          'role_session_name': '',
                          'jwt': ''
         }
diff --git a/framework/replay/programs/compare.py b/framework/replay/programs/compare.py
index 6648adc65..c5372f076 100644
--- a/framework/replay/programs/compare.py
+++ b/framework/replay/programs/compare.py
@@ -39,6 +39,7 @@ def _from_yaml(args):
     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['minio_bucket'] = args.download_minio_bucket
     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
@@ -53,6 +54,7 @@ def _trace(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['minio_bucket'] = args.download_minio_bucket
     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
@@ -84,6 +86,7 @@ def compare(input_):
                  parsers.DOWNLOAD_URL,
                  parsers.DOWNLOAD_FORCE,
                  parsers.DOWNLOAD_MINIO_HOST,
+                 parsers.DOWNLOAD_MINIO_BUCKET,
                  parsers.DOWNLOAD_ROLE_SESSION_NAME,
                  parsers.DOWNLOAD_JWT,
                  parsers.DB_PATH,
@@ -107,6 +110,7 @@ def compare(input_):
                  parsers.YAML,
                  parsers.DOWNLOAD_FORCE,
                  parsers.DOWNLOAD_MINIO_HOST,
+                 parsers.DOWNLOAD_MINIO_BUCKET,
                  parsers.DOWNLOAD_ROLE_SESSION_NAME,
                  parsers.DOWNLOAD_JWT,
                  parsers.DB_PATH,
diff --git a/framework/replay/programs/download.py b/framework/replay/programs/download.py
index 926353de3..a9e84b3bc 100644
--- a/framework/replay/programs/download.py
+++ b/framework/replay/programs/download.py
@@ -37,6 +37,7 @@ 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['minio_bucket'] = args.download_minio_bucket
     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
@@ -50,6 +51,7 @@ def download(input_):
     parser = argparse.ArgumentParser(parents=[parsers.DOWNLOAD_URL,
                                               parsers.DOWNLOAD_FORCE,
                                               parsers.DOWNLOAD_MINIO_HOST,
+                                              parsers.DOWNLOAD_MINIO_BUCKET,
                                               parsers.DOWNLOAD_ROLE_SESSION_NAME,
                                               parsers.DOWNLOAD_JWT,
                                               parsers.DB_PATH])
diff --git a/framework/replay/programs/parsers.py b/framework/replay/programs/parsers.py
index 2325d5e84..e225e3790 100644
--- a/framework/replay/programs/parsers.py
+++ b/framework/replay/programs/parsers.py
@@ -75,6 +75,14 @@ DOWNLOAD_MINIO_HOST.add_argument(
     default=None,
     help=('name of MinIO server from which to download traces'))
 
+DOWNLOAD_MINIO_BUCKET = argparse.ArgumentParser(add_help=False)
+DOWNLOAD_MINIO_BUCKET.add_argument(
+    '-b', '--minio_bucket',
+    dest='download_minio_bucket',
+    required=False,
+    default=None,
+    help=('Name of bucket in MinIO server containing the traces'))
+
 DOWNLOAD_ROLE_SESSION_NAME = argparse.ArgumentParser(add_help=False)
 DOWNLOAD_ROLE_SESSION_NAME.add_argument(
     '-r', '--role-session-name',