Created
September 4, 2022 08:25
-
-
Save moyix/71896182bc5937fb8a1a882d765bc8ac to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
import os | |
import sys | |
import subprocess as sp | |
import tempfile | |
import hashlib | |
script_dir = os.path.dirname(os.path.realpath(__file__)) | |
sys.path.append(script_dir) | |
from fast_check_for_ffast_math import check_file | |
# MAX_SIZE = 30 * 1024 * 1024 * 1024 # 30GB | |
def dirsize(d): | |
size = sp.run(['du', '-csb', d], | |
universal_newlines=True, | |
stdout=sp.PIPE, stderr=sp.DEVNULL | |
).stdout.splitlines()[-1].split()[0] | |
return int(size) | |
def extract_with_unzip(wheel, dest): | |
os.makedirs(dest, exist_ok=True) | |
zip_out = sp.run( | |
['unzip', wheel, '*.so*', '-d', dest], | |
stdout=sp.PIPE, stderr=sp.PIPE, check=False | |
) | |
def remove(d): | |
sp.run(['rm', '-rf', d], check=False) | |
def hash_file(path): | |
BUF_SIZE = 65536 | |
sha256 = hashlib.sha256() | |
with open(path, 'rb') as f: | |
while True: | |
data = f.read(BUF_SIZE) | |
if not data: | |
break | |
sha256.update(data) | |
return sha256.hexdigest() | |
def main(args): | |
JOB_NUM = int(args[1]) | |
LOG_DIR = '/fastdata/pypi_wheels/logs' | |
HASH_LOG = os.path.join(LOG_DIR, f'hashes.{JOB_NUM}.txt') | |
MATH_LOG = os.path.join(LOG_DIR, f'ffast_math.{JOB_NUM}.txt') | |
EXTRACT_DIR = '/fastdata/pypi_wheels/extract' | |
os.makedirs(EXTRACT_DIR, exist_ok=True) | |
os.makedirs(LOG_DIR, exist_ok=True) | |
with open(HASH_LOG, 'w') as hash_log, open(MATH_LOG, 'w') as math_log: | |
for file in args[2:]: | |
with tempfile.TemporaryDirectory( | |
prefix=os.path.join(EXTRACT_DIR, 'extract_')) as tmpdir: | |
per_wheel_dir = os.path.join(tmpdir, file) | |
extract_with_unzip(file, per_wheel_dir) | |
for root, dirs, files in os.walk(per_wheel_dir): | |
for f in files: | |
path = os.path.join(root, f) | |
hash_log.write(f'{file}\t{hash_file(path)}\t{path}\n') | |
check_file(path, math_log) | |
remove(per_wheel_dir) | |
if __name__ == '__main__': | |
main(sys.argv) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import sys | |
import mmap | |
from elftools.elf.elffile import ELFFile, ELFError | |
import struct | |
set_fast_math_code = bytes.fromhex('0fae5c24fc814c24fc408000000fae5424fcc3') | |
def load_bytes_from_elf(bindata, elf, vaddr, size): | |
try: | |
paddr = next(iter(elf.address_offsets(vaddr))) | |
except StopIteration: | |
return None | |
return bindata[paddr:paddr+size] | |
def check_for_ffast_math(elf, bindata, addr): | |
# +4 because we may have an endbr64 instruction (f30f1efa) in the prologue | |
code = load_bytes_from_elf(bindata, elf, addr, len(set_fast_math_code)+4) | |
if code is None: | |
return False | |
return set_fast_math_code in code | |
def check_file(filename, out=sys.stdout): | |
print(f"{filename} ", end="", flush=True, file=out) | |
try: | |
try: | |
elf = ELFFile(open(filename,'rb')) | |
except ELFError: | |
print("is not an ELF file", file=out) | |
return | |
if (arch := elf.get_machine_arch()) != 'x64': | |
print(f"is {arch} not x64", file=out) | |
return | |
bindata = mmap.mmap(elf.stream.fileno(), 0, access=mmap.ACCESS_READ) | |
try: | |
s = elf.get_section_by_name('.init_array') | |
except ELFError: | |
print("error parsing ELF file", file=out) | |
return | |
if s is not None and (ia_data := s.data()): | |
constructors = struct.unpack(f'<{len(ia_data)//8}Q', ia_data) | |
else: | |
constructors = [] | |
i = len(constructors) | |
print(f"({i} constructor{'s'[:i^1]}) ", end="", flush=True, file=out) | |
for addr in constructors: | |
if check_for_ffast_math(elf, bindata, addr): | |
print(f"contains ffast-math constructor at {hex(addr)}", file=out) | |
break | |
else: | |
print("is clean", file=out) | |
except Exception as e: | |
print(f"error: {e}", file=out) | |
if __name__ == "__main__": | |
for filename in sys.argv[1:]: | |
check_file(filename) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment