| 1 | # SPDX-License-Identifier: GPL-2.0-or-later |
| 2 | # |
| 3 | # Copyright 2024 Red Hat, Inc. |
| 4 | # |
| 5 | # This work is licensed under the terms of the GNU GPL, version 2 or |
| 6 | # later. See the COPYING file in the top-level directory. |
| 7 | ''' |
| 8 | Test utilities for fetching & caching assets |
| 9 | ''' |
| 10 | |
| 11 | import hashlib |
| 12 | import logging |
| 13 | import os |
| 14 | import stat |
| 15 | import sys |
| 16 | import time |
| 17 | import unittest |
| 18 | import urllib.request |
| 19 | from time import sleep |
| 20 | from pathlib import Path |
| 21 | from shutil import copyfileobj |
| 22 | from urllib.error import HTTPError, URLError |
| 23 | |
| 24 | |
| 25 | class AssetError(Exception): |
| 26 | '''This exception will be raised if an asset is not usable''' |
| 27 | def __init__(self, asset, msg, transient=False): |
| 28 | self.url = asset.url |
| 29 | self.msg = msg |
| 30 | self.transient = transient |
| 31 | |
| 32 | def __str__(self): |
| 33 | return f"{self.url}: {self.msg}" |
| 34 | |
| 35 | |
| 36 | class Asset: |
| 37 | ''' |
| 38 | This class is used to represent an asset that gets downloaded from |
| 39 | the internet and will be stored in the local asset cache. |
| 40 | Instances of this class must be declared as class level variables |
| 41 | starting with a name "ASSET_". This enables the pre-caching logic |
| 42 | to easily find all referenced assets and download them prior to |
| 43 | execution of the tests. |
| 44 | ''' |
| 45 | def __init__(self, url, hashsum): |
| 46 | self.url = url |
| 47 | self.hash = hashsum |
| 48 | cache_dir_env = os.getenv('QEMU_TEST_CACHE_DIR') |
| 49 | if cache_dir_env: |
| 50 | self.cache_dir = Path(cache_dir_env, "download") |
| 51 | else: |
| 52 | self.cache_dir = Path(Path("~").expanduser(), |
| 53 | ".cache", "qemu", "download") |
| 54 | self.cache_file = Path(self.cache_dir, hashsum) |
| 55 | self.log = logging.getLogger('qemu-test') |
| 56 | |
| 57 | def __repr__(self): |
| 58 | return f"Asset: url={self.url} hash={self.hash} cache={self.cache_file}" |
| 59 | |
| 60 | def __str__(self): |
| 61 | return str(self.cache_file) |
| 62 | |
| 63 | def _check(self, cache_file): |
| 64 | if self.hash is None: |
| 65 | return True |
| 66 | if len(self.hash) == 64: |
| 67 | hl = hashlib.sha256() |
| 68 | else: |
| 69 | raise AssetError(self, "unsupported hash type") |
| 70 | |
| 71 | # Calculate the hash of the file: |
| 72 | with open(cache_file, 'rb') as file: |
| 73 | while True: |
| 74 | chunk = file.read(1 << 20) |
| 75 | if not chunk: |
| 76 | break |
| 77 | hl.update(chunk) |
| 78 | |
| 79 | return self.hash == hl.hexdigest() |
| 80 | |
| 81 | def valid(self): |
| 82 | '''Check whether the file exists in the cache and has the right hash''' |
| 83 | if os.getenv("QEMU_TEST_REFRESH_CACHE", None) is not None: |
| 84 | self.log.info("Force refresh of asset %s", self.url) |
| 85 | return False |
| 86 | |
| 87 | return self.cache_file.exists() and self._check(self.cache_file) |
| 88 | |
| 89 | def fetchable(self): |
| 90 | '''Check whether we are allowed to download assets from the internet''' |
| 91 | return not os.environ.get("QEMU_TEST_NO_DOWNLOAD", False) |
| 92 | |
| 93 | def available(self): |
| 94 | '''Check whether the asset is either in the cache or fetchable''' |
| 95 | return self.valid() or self.fetchable() |
| 96 | |
| 97 | def _wait_for_other_download(self, tmp_cache_file): |
| 98 | # Another thread already seems to download the asset, so wait until |
| 99 | # it is done, while also checking the size to see whether it is stuck |
| 100 | try: |
| 101 | current_size = tmp_cache_file.stat().st_size |
| 102 | new_size = current_size |
| 103 | except: |
| 104 | if os.path.exists(self.cache_file): |
| 105 | return True |
| 106 | raise |
| 107 | waittime = lastchange = 600 |
| 108 | while waittime > 0: |
| 109 | sleep(1) |
| 110 | waittime -= 1 |
| 111 | try: |
| 112 | new_size = tmp_cache_file.stat().st_size |
| 113 | except: |
| 114 | if os.path.exists(self.cache_file): |
| 115 | return True |
| 116 | raise |
| 117 | if new_size != current_size: |
| 118 | lastchange = waittime |
| 119 | current_size = new_size |
| 120 | elif lastchange - waittime > 90: |
| 121 | return False |
| 122 | |
| 123 | self.log.debug("Time out while waiting for %s!", tmp_cache_file) |
| 124 | raise TimeoutError(f"Time out while waiting for {tmp_cache_file}") |
| 125 | |
| 126 | def _save_time_stamp(self): |
| 127 | ''' |
| 128 | Update the time stamp of the asset in the cache. Unfortunately, we |
| 129 | cannot use the modification or access time of the asset file itself, |
| 130 | since e.g. the functional jobs in the gitlab CI reload the files |
| 131 | from the gitlab cache and thus always have recent file time stamps, |
| 132 | so we have to save our asset time stamp to a separate file instead. |
| 133 | ''' |
| 134 | self.cache_file.with_suffix(".stamp").write_text(f"{int(time.time())}") |
| 135 | |
| 136 | def _try_to_fetch(self, tmp_cache_file): |
| 137 | for _retries in range(3): |
| 138 | try: |
| 139 | with tmp_cache_file.open("xb") as dst: |
| 140 | with urllib.request.urlopen(self.url) as resp: |
| 141 | copyfileobj(resp, dst) |
| 142 | length_hdr = resp.getheader("Content-Length") |
| 143 | |
| 144 | # Verify downloaded file size against length metadata, if |
| 145 | # available. |
| 146 | if length_hdr is not None: |
| 147 | length = int(length_hdr) |
| 148 | fsize = tmp_cache_file.stat().st_size |
| 149 | if fsize != length: |
| 150 | self.log.error("Unable to download %s: " |
| 151 | "connection closed before " |
| 152 | "transfer complete (%d/%d)", |
| 153 | self.url, fsize, length) |
| 154 | tmp_cache_file.unlink() |
| 155 | continue |
| 156 | break |
| 157 | except FileExistsError: |
| 158 | self.log.debug("%s already exists, " |
| 159 | "waiting for other thread to finish...", |
| 160 | tmp_cache_file) |
| 161 | if self._wait_for_other_download(tmp_cache_file): |
| 162 | return True |
| 163 | self.log.debug("%s seems to be stale, " |
| 164 | "deleting and retrying download...", |
| 165 | tmp_cache_file) |
| 166 | tmp_cache_file.unlink() |
| 167 | continue |
| 168 | except HTTPError as e: |
| 169 | tmp_cache_file.unlink() |
| 170 | self.log.error("Unable to download %s: HTTP error %d", |
| 171 | self.url, e.code) |
| 172 | # Treat 404 as fatal, since it is highly likely to |
| 173 | # indicate a broken test rather than a transient |
| 174 | # server or networking problem |
| 175 | if e.code == 404: |
| 176 | raise AssetError(self, "Unable to download: " |
| 177 | f"HTTP error {e.code}") from e |
| 178 | continue |
| 179 | except URLError as e: |
| 180 | # This is typically a network/service level error |
| 181 | # eg urlopen error [Errno 110] Connection timed out> |
| 182 | tmp_cache_file.unlink() |
| 183 | self.log.error("Unable to download %s: URL error %s", |
| 184 | self.url, e.reason) |
| 185 | raise AssetError(self, |
| 186 | f"Unable to download: URL error{e.reason}", |
| 187 | transient=True) from e |
| 188 | except ConnectionError as e: |
| 189 | # A socket connection failure, such as dropped conn |
| 190 | # or refused conn |
| 191 | tmp_cache_file.unlink() |
| 192 | self.log.error("Unable to download %s: Connection error %s", |
| 193 | self.url, e) |
| 194 | continue |
| 195 | except Exception as e: |
| 196 | tmp_cache_file.unlink() |
| 197 | raise AssetError(self, f"Unable to download: {e}", |
| 198 | transient=True) from e |
| 199 | return False |
| 200 | |
| 201 | def fetch(self): |
| 202 | '''Download the asset from the internet''' |
| 203 | if not self.cache_dir.exists(): |
| 204 | self.cache_dir.mkdir(parents=True, exist_ok=True) |
| 205 | |
| 206 | if self.valid(): |
| 207 | self.log.debug("Using cached asset %s for %s", |
| 208 | self.cache_file, self.url) |
| 209 | self._save_time_stamp() |
| 210 | return str(self.cache_file) |
| 211 | |
| 212 | if not self.fetchable(): |
| 213 | raise AssetError(self, |
| 214 | "Asset cache is invalid and downloads disabled") |
| 215 | |
| 216 | self.log.info("Downloading %s to %s...", self.url, self.cache_file) |
| 217 | tmp_cache_file = self.cache_file.with_suffix(".download") |
| 218 | |
| 219 | if self._try_to_fetch(tmp_cache_file): |
| 220 | return str(self.cache_file) |
| 221 | |
| 222 | if not os.path.exists(tmp_cache_file): |
| 223 | raise AssetError(self, "Download retries exceeded", transient=True) |
| 224 | |
| 225 | try: |
| 226 | # Set these just for informational purposes. Note that |
| 227 | # setxattr is Linux-only; as this is only informational |
| 228 | # we can simply skip it on other platforms. |
| 229 | if hasattr(os, "setxattr"): |
| 230 | os.setxattr(str(tmp_cache_file), "user.qemu-asset-url", |
| 231 | self.url.encode('utf8')) |
| 232 | os.setxattr(str(tmp_cache_file), "user.qemu-asset-hash", |
| 233 | self.hash.encode('utf8')) |
| 234 | except OSError as e: |
| 235 | self.log.debug("Unable to set xattr on %s: %s", tmp_cache_file, e) |
| 236 | |
| 237 | if not self._check(tmp_cache_file): |
| 238 | tmp_cache_file.unlink() |
| 239 | raise AssetError(self, f"Hash does not match {self.hash}") |
| 240 | tmp_cache_file.replace(self.cache_file) |
| 241 | self._save_time_stamp() |
| 242 | # Remove write perms to stop tests accidentally modifying them |
| 243 | os.chmod(self.cache_file, stat.S_IRUSR | stat.S_IRGRP) |
| 244 | |
| 245 | self.log.info("Cached %s at %s", self.url, self.cache_file) |
| 246 | return str(self.cache_file) |
| 247 | |
| 248 | @staticmethod |
| 249 | def precache_test(test): |
| 250 | ''' |
| 251 | Look for variables starting with "ASSET_" and try to fetch the asset |
| 252 | that is specified there. |
| 253 | ''' |
| 254 | log = logging.getLogger('qemu-test') |
| 255 | log.setLevel(logging.DEBUG) |
| 256 | handler = logging.StreamHandler(sys.stdout) |
| 257 | handler.setLevel(logging.DEBUG) |
| 258 | formatter = logging.Formatter( |
| 259 | '%(asctime)s - %(name)s - %(levelname)s - %(message)s') |
| 260 | handler.setFormatter(formatter) |
| 261 | log.addHandler(handler) |
| 262 | for name, asset in vars(test.__class__).items(): |
| 263 | if name.startswith("ASSET_") and isinstance(asset, Asset): |
| 264 | try: |
| 265 | asset.fetch() |
| 266 | except AssetError as e: |
| 267 | if not e.transient: |
| 268 | raise |
| 269 | log.error("%s: skipping asset precache", e) |
| 270 | |
| 271 | log.removeHandler(handler) |
| 272 | |
| 273 | @staticmethod |
| 274 | def precache_suite(suite): |
| 275 | ''' |
| 276 | Iterate through all tests/suites in a suite and precache their assets |
| 277 | ''' |
| 278 | for test in suite: |
| 279 | if isinstance(test, unittest.TestSuite): |
| 280 | Asset.precache_suite(test) |
| 281 | elif isinstance(test, unittest.TestCase): |
| 282 | Asset.precache_test(test) |
| 283 | |
| 284 | @staticmethod |
| 285 | def precache_suites(path, cache_tstamp): |
| 286 | ''' |
| 287 | Get the available test suite and precache their assets |
| 288 | ''' |
| 289 | loader = unittest.loader.defaultTestLoader |
| 290 | tests = loader.loadTestsFromNames([path], None) |
| 291 | |
| 292 | with open(cache_tstamp, "w", encoding='utf-8'): |
| 293 | Asset.precache_suite(tests) |