creal.py

#!/usr/bin/env python3
import os, sys, shutil, re, time, tempfile, signal, random, string, argparse
from datetime import datetime
from glob import glob
from enum import Enum, auto
from diopter.compiler import (
    CompilationSetting,
    CompilerExe,
    OptLevel,
    SourceProgram,
    Language,
    ObjectCompilationOutput
)
from diopter.sanitizer import Sanitizer
from diopter.utils import TempDirEnv
import subprocess as sp
from synthesizer.synthesizer import Synthesizer, SynthesizerError
from utils.compcert import CComp as this_CComp
from pathlib import Path
from datetime import datetime
from termcolor import colored

def print_red(msg):
    print(colored(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + ' >', 'yellow'), colored(msg, 'red'), flush=True)
def print_green(msg):
    print(colored(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + ' >', 'yellow'), colored(msg, 'green'), flush=True)
def print_blue(msg):
    print(colored(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + ' >', 'yellow'), colored(msg, 'blue'), flush=True)

def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for _ in range(size))


DEBUG = 0
"""CONFIG"""
FUNCTION_DB_FILE = os.path.join(os.path.dirname(__file__), './databaseconstructor/functions_pointer_global_io.json')
MIN_PROGRAM_SIZE = 8000 # programs shorter than this many bytes are too boring to test
NUM_MUTANTS = 10 # number of mutants generated by the synthesizer per seed.
COMPILER_TIMEOUT = 200
PROG_TIMEOUT = 10
CCOMP_TIMEOUT = 60 # compcert timeout
CSMITH_USER_OPTIONS = "--no-volatiles --no-volatile-pointers --no-unions"
CSMITH_TIMEOUT = 20
CREDUCE_JOBS = 1
"""TOOL"""
CSMITH_HOME = os.environ["CSMITH_HOME"]

if not os.path.exists(os.path.join(CSMITH_HOME, 'include/csmith.h')):
    print_red('CSMITH_HOME is not set correctly, cannot find csmith.h in "$CSMITH_HOME/include/".')
    sys.exit(1)

CC = CompilationSetting(
            compiler=CompilerExe.get_system_gcc(),
            opt_level=OptLevel.O3,
            flags=("-march=native",f"-I{CSMITH_HOME}/include"),
            )
SAN_SAN = Sanitizer(checked_warnings=False, use_ccomp_if_available=False) # sanitizers only
SAN_CCOMP = this_CComp.get_system_ccomp() # CompCert only

"""Global vars"""

class CompCode(Enum):
    """Compile status
    """
    OK      =   auto()  # ok
    Timeout =   auto()  # timeout during compilation
    Sanfail =   auto()  # sanitization failed
    Crash   =   auto()  # compiler crash
    Error   =   auto()  # compiler error
    WrongEval=  auto()  # inconsistent results across compilers but consistent within the same compiler
    Wrong   =   auto()  # inconsistent results across compilers/opts

def generate_random_string(len:int=5) -> str:
    """Generate a random string of length len"""
    return ''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(len))

def run_cmd(cmd, timeout):
    if type(cmd) is not list:
        cmd = cmd.split(' ')
        cmd = list(filter(lambda x: x!='', cmd))
    # Start the subprocess
    process = sp.Popen(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
    # Wait for the subprocess to finish or timeout
    try:
        output, error = process.communicate(timeout=timeout)
        output = output.decode("utf-8")
    except sp.TimeoutExpired:
        # Timeout occurred, kill the process
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        finally:
            output = ''
        # A workaround to tmpxxx.exe as it sometimes escapes from os.killpg
        cmd_str = " ".join(cmd)
        time.sleep(2)
        if '.exe' in cmd_str:
            os.system(f"pkill -9 -f {cmd_str}")
        return 124, output

    # Return the exit code and stdout of the process
    return process.returncode, output

def write_bug_desc_to_file(to_file, data):
    with open(to_file, "a") as f:
        f.write(f"/* {data} */\n")

def read_checksum(data):
    res = re.findall(r'checksum = (.*)', data)
    if len(res) > 0:
        return res[0]
    return 'NO_CKSUM'

def check_sanitizers(src):
    """Check validity with sanitizers"""
    with open(src, 'r') as f:
        code = f.read()
    prog = SourceProgram(code=code, language=Language.C)
    preprog = CC.preprocess_program(prog, make_compiler_agnostic=True)
    if DEBUG:
        print(datetime.now().strftime("%d/%m/%Y %H:%M:%S"), "SAN.sanitize", flush=True)
    if not SAN_SAN.sanitize(preprog):
        return False
    return True

def check_ccomp(src, random_count=1):
    """
    Check validity with CompCert.
    src:str -> source file
    random_count:int -> the number of times using ccomp -random for checking
    """
    with open(src, 'r') as f:
        code = f.read()
    prog = SourceProgram(code=code, language=Language.C)
    preprog = CC.preprocess_program(prog, make_compiler_agnostic=True)
    if DEBUG:
        print(datetime.now().strftime("%d/%m/%Y %H:%M:%S"), "SAN.ccomp", flush=True)
    with TempDirEnv():
        try:
            ccomp_result = SAN_CCOMP.check_program(preprog, timeout=CCOMP_TIMEOUT, debug=DEBUG)
        except sp.TimeoutExpired:
            return False
        if ccomp_result is False:
            return False
    with TempDirEnv():
        for _ in range(random_count):
            try:
                ccomp_result_random = SAN_CCOMP.check_program(preprog, timeout=CCOMP_TIMEOUT, debug=DEBUG, additional_flags=["-random"])
            except sp.TimeoutExpired:
                return False
            if ccomp_result_random is False:
                return False
            # check for unspecified behavior
            if ccomp_result.stdout != ccomp_result_random.stdout:
                return False
    return True

def compile_and_run(compiler, src):
    cksum = ''
    tmp_f = tempfile.NamedTemporaryFile(suffix=".exe", delete=False)
    tmp_f.close()
    exe = tmp_f.name
    cmd = f"{compiler} {src} -I{CSMITH_HOME}/include -o {exe}"
    ret, out = run_cmd(cmd, COMPILER_TIMEOUT)
    if ret == 124: # another compile chance when timeout
        time.sleep(1)
        ret, out = run_cmd(cmd, COMPILER_TIMEOUT)
    if ret == 124: # we treat timeout as crash now.
        write_bug_desc_to_file(src, f"Compiler timeout! Can't compile with {compiler}")
        if os.path.exists(exe): os.remove(exe)
        return CompCode.Timeout, cksum
    if ret != 0:
        write_bug_desc_to_file(src, f"Compiler crash! Can't compile with {compiler}")
        if os.path.exists(exe): os.remove(exe)
        return CompCode.Crash, cksum
    ret, out = run_cmd(f"{exe}", PROG_TIMEOUT)
    cksum = read_checksum(out)
    write_bug_desc_to_file(src, f"EXITof {compiler}: {ret}")
    write_bug_desc_to_file(src, f"CKSMof {compiler}: {cksum}")
    if os.path.exists(exe): os.remove(exe)
    return CompCode.OK, cksum

def check_compile(src:str, compilers:list) -> CompCode:
    """Compile the program with a list of compilers and check their status
    """
    cksum_list = []
    for comp in compilers:
        if DEBUG:
            print(datetime.now().strftime("%d/%m/%Y %H:%M:%S"), "compiler_and_run: ", comp, flush=True)
        ret, cksum = compile_and_run(comp, src)
        if ret == CompCode.Crash:
            return CompCode.Crash
        if ret == CompCode.Timeout:
            return CompCode.Timeout
        if ret != CompCode.OK:
            return CompCode.Error
        cksum_list.append(cksum)
    if len(cksum_list) != len(compilers) or len(set(cksum_list)) != 1:
        maybe_WrongEval = True
        for i in range(len(compilers)):
            for j in range(i+1, len(compilers)):
                if compilers[i].split(' ')[0] == compilers[j].split(' ')[0] and cksum_list[i] != cksum_list[j]:
                    maybe_WrongEval = False
        if maybe_WrongEval:
            return CompCode.WrongEval
        return CompCode.Wrong
    return CompCode.OK

def run_one(compilers:list[str], dst_dir:Path, SYNER:Synthesizer) -> Path | None:
    """Run compiler testing
    """
    save_realsmith_dir = (dst_dir)
    succ_file_id = id_generator()
    src = str((dst_dir / f'{succ_file_id}_seed.c').absolute())
    print_blue('Generating seed...')
    while True:
        cmd = f"{CSMITH_HOME}/bin/csmith {CSMITH_USER_OPTIONS} --output {src}"
        ret, out = run_cmd(cmd, CSMITH_TIMEOUT)
        if ret != 0:
            print("csmith failed: generation.")
            continue
        # check size
        if os.path.getsize(src) < MIN_PROGRAM_SIZE:
            print("csmith failed: small program.")
            continue
        # check sanitization
        if check_sanitizers(src) and check_ccomp(src):
            break
        print("csmith failed: sanitization.")
    print_blue(f"Seed generated: {src}")
    ret = check_compile(src, compilers)
    print_blue('Synthesizing mutants...')
    # synthesize
    try:
        syn_files = SYNER.synthesizer(src_filename=src, num_mutant=NUM_MUTANTS, DEBUG=DEBUG)
    except:
        print('SynthesizerError!')
        os.remove(src)
        return 0

    print_green(f'Synthesizing done! Programs saved as {src.replace(".c", "_syn*.c")}')

    return 0


if __name__=='__main__':
    parser = argparse.ArgumentParser(description="Generate a number of realsmith mutants for evaluation.")
    parser.add_argument("--dst", required=True, type=Path, help="Destination directory for generated seeds.")
    parser.add_argument("--syn-prob", required=True, type=int, help="Synthesis probability")
    parser.add_argument("--num-mutants", required=True, type=int, help="The number of mutants per seed by realsmith")
    args = parser.parse_args()

    dst_dir = Path(args.dst)
    dst_dir.mkdir(parents=True, exist_ok=True)

    NUM_MUTANTS = args.num_mutants

    compilers = [
        "gcc -O0",
        "clang -O0"
    ]
    SYNER = Synthesizer(func_database=FUNCTION_DB_FILE, prob=args.syn_prob)
    with TempDirEnv() as tmp_dir:
        os.environ['TMPDIR'] = tmp_dir.absolute().as_posix()
        total = 0
        ret = run_one(compilers, dst_dir, SYNER)