demo.py

# Some example code for PyH2 hardware testing

import math
import os
import random
import sys

import afl

from hypothesis import given, strategies as st


"""
Setting up a demo test - standard Hypothesis test for a mock GCD unit.
"""


class GCD:
    def __init__(self, bitwidth, qsize):
        self.bitwidth = bitwidth
        self.qsize = qsize

    def __call__(self, x, y):
        return math.gcd(x, y)


GCD_UNITS = st.builds(
    GCD, bitwidth=st.sampled_from([8, 16, 32, 64, 128]), qsize=st.integers(1, 20),
)


def operands(bitwidth, max_length):
    values = st.integers(min_value=0, max_value=2 ** bitwidth)
    return st.lists(st.tuples(values, values), min_size=1, max_size=max_length)


@given(design=GCD_UNITS, data=st.data())
def test_designs(design, data):
    # Mimics the structure of a hardware test - we draw the design params
    # up front, then values which depend on the design.
    design_specific_op_strategy = operands(design.bitwidth, design.qsize)
    for x, y in data.draw(design_specific_op_strategy):
        divisor = design(x, y)
        assert x % divisor == 0
        assert y % divisor == 0


"""
Next, let's use python-afl to drive it:
"""


def shortlex_order(buf):
    return len(buf), buf


def get_buffers_for(target, how_many=100, nbytes=1000):
    corpus = set()
    while len(corpus) < how_many:
        rand_buf = random.getrandbits(nbytes * 8).to_bytes(nbytes, "big")
        # NOTE: this might raise an exception, in which case we've found a bug early...
        canonical_buf = target.hypothesis.fuzz_one_input(rand_buf)
        if canonical_buf is not None:
            corpus.add(canonical_buf)
    return corpus


if __name__ == "__main__":
    # If invoked as `python demo.py`, we'll actually run AFL.
    # First, let's ensure that we have a starter corpus:

    if not os.path.isdir("corpus"):
        bytestrings = get_buffers_for(test_designs)
        os.makedirs("corpus")
        for i, buf in enumerate(sorted(bytestrings, key=shortlex_order)):
            with open(f"corpus/{i:d}.data", "wb") as f:
                f.write(buf)

    # TODO: invoke afl-cmin here to remove redundant entries!

    # Now it's time to actually run afl-fuzz: we call the init() function
    # to indicate that this is the point to fork the process, then pass
    # the contents of sys.stdin to our fuzz_one_input
    afl.init()
    test_designs.hypothesis.fuzz_one_input(sys.stdin.buffer.read())

    # Once Hypothesis is done, we'll just kill the process without letting
    # Python run any cleanup code.  Slightly faster, and it should be fine...
    os._exit(0)


"""
With that excitement over, let's try Zac's 'predefined prefixes' trick:
"""


@given(GCD_UNITS)
def driver(_):
    pass


def run_prefix_trick(target, nbytes=1000, prefixes=100, examples_per_prefix=100):
    prefixes = get_buffers_for(driver, how_many=prefixes, nbytes=64)
    # Outer loop is _kinda_ like iterative deepening over designs
    for prefix in sorted(bytestrings, key=shortlex_order):
        # Inner loop is CRT over actions (but with parse-level heuristics)
        for _ in range(examples_per_prefix):
            rand_buf = random.getrandbits(nbytes * 8).to_bytes(nbytes, "big")
            target.hypothesis.fuzz_one_input(prefix + rand_buf)
    # If we found any bugs, they're probably for fairly small designs -
    # and just like the fuzz above, executing `test_designs()` will pick
    # them out of the Hypothesis example database and try to shrink them!