Skip to content

Commit

Permalink
fix: version check before data_filter import
Browse files Browse the repository at this point in the history
- better error handling if untar fails
  • Loading branch information
18alantom committed Jan 24, 2024
1 parent 23bd717 commit 80f2e70
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 4 deletions.
3 changes: 2 additions & 1 deletion bench/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,8 @@ def get_cached(self) -> bool:
with tarfile.open(cache_path, mode) as tar:
try:
tar.extractall(app_path.parent, filter=get_app_cache_extract_filter())
except:
except Exception:
logger.exception(f"Cache extraction failed for {self.app_name}")
shutil.rmtree(app_path)
return False

Expand Down
12 changes: 9 additions & 3 deletions bench/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from glob import glob
from pathlib import Path
from shlex import split
from tarfile import data_filter, AbsoluteLinkError, TarInfo
from tarfile import AbsoluteLinkError, TarInfo
from typing import List, Optional, Tuple

# imports - third party imports
Expand Down Expand Up @@ -578,15 +578,21 @@ def get_app_cache_extract_filter(
): # -> Callable[[TarInfo, str], TarInfo | None]
state = dict(count=0, size=0)

if sys.version_info.major <=2 or sys.version_info.minor <=8:
def data_filter(m, p):
return m
else:
from tarfile import data_filter

def filter_function(member: TarInfo, dest_path: str) -> Optional[TarInfo]:
state["count"] += 1
state["size"] += member.size

if state["count"] > count_threshold:
raise Exception(f"Number of entries exceeds threshold ({state['count']})")
raise RuntimeError(f"Number of entries exceeds threshold ({state['count']})")

if state["size"] > size_threshold:
raise Exception(f"Extracted size exceeds threshold ({state['size']})")
raise RuntimeError(f"Extracted size exceeds threshold ({state['size']})")

try:
return data_filter(member, dest_path)
Expand Down

0 comments on commit 80f2e70

Please sign in to comment.