fuzzer.j2

/// This file is autogenerated

#include "{{ class_header }}"

#include <cstdint>
#include <iterator> // for std::size
#include <random>
#include <iostream>

/// ------------------------------ DUMPED DATA ------------------------------

namespace target {

struct Argument {
    bool is_pointer;
    size_t size;
};

const size_t ARG_LIMIT = {{ arg_limit }};

struct AData {
    Argument list[ARG_LIMIT];
    size_t size;
    size_t bytesize;
};

/// -------------- CONSTRUCTORS --------------

struct CData {
    AData args;
    {{ class_name }} (*fn)(const uint8_t *);
};

{% for c in constructors %}
{{ class_name }} constr_{{ loop.index }}(const uint8_t *data) {
{% if length(c.args) == 0 %}    return {{ class_name }}();
{% else %}    size_t size = 0;

    {% for a in c.args -%}
    {{ a.name }} *arg_{{ loop.index }} = ({{ a.name }} *)(data + size);
    size += sizeof({{ a.name }});
    {% endfor %}

    return {{ class_name }}({% for a in c.args %}{% if loop.index != 0 %}, {% endif %}*arg_{{loop.index}}{% endfor %});
{% endif %}{{ "};" }}
{% endfor %}

const CData c_list[] = {
    {% for c in constructors -%}
    {
        .args = {
            .list = {
                {% for a in c.args %}{{ "{" }}{{ a.is_pointer }}, sizeof({{ a.name }})},
                {% endfor %}
            },
            .size = {{ length(c.args) }},
            .bytesize = 0{% for a in c.args %} + sizeof({{ a.name }}){% endfor %},
        },
        .fn = constr_{{ loop.index }},
    },{% endfor %}
};

constexpr size_t c_size = std::size(c_list);

/// -------------- METHODS --------------

struct MData {
    AData args;
    void (*fn)({{ class_name }} *, const uint8_t *);
};

{% for m in methods %}
void method_{{ loop.index }}({{ class_name }} *obj, const uint8_t *data) {
{% if length(m.args) == 0 %}    obj->{{ m.name }}();
{% else %}    size_t size = 0;

    {% for a in m.args -%}
    {{ a.name }} *arg_{{ loop.index }} = ({{ a.name }} *)(data + size);
    size += sizeof({{ a.name }});
    {% endfor %}

    obj->{{ m.name }}({% for a in m.args %}{% if loop.index != 0 %}, {% endif %}*arg_{{ loop.index }}{% endfor %});
{% endif %}{{ "};"}}
{% endfor %}

const MData m_list[] = {
    {% for m in methods -%}
    {
        .args = {
            .list = {
                {% for a in m.args %}{{ "{" }}{{ a.is_pointer }}, sizeof({{ a.name }})},
                {% endfor %}
            },
            .size = {{ length(m.args) }},
            .bytesize = 0{% for a in m.args %} + sizeof({{ a.name }}){% endfor %},
        },
        .fn = method_{{ loop.index }},
    },{% endfor %}
};

constexpr size_t m_size = std::size(m_list);

} // namespace

/// ------------------------------ CALL CHAIN AND MUTATOR ------------------------------

extern "C" size_t LLVMFuzzerMutate(uint8_t *data, size_t size, size_t max_size);

#include <iomanip>
void print_hex(const uint8_t *data, size_t size) {
    for (size_t i = 0; i < size; ++i)
        std::cout << std::hex << std::setfill('0') << std::setw(2) << (int)data[i];
}

/**
 * This is class for working with Call Chains (mutations and other userful functions)
 * Note that parsing takes time and for simple classes it's better to used standart binary mutations
*/
class CallChain {
public:
    /// MAX amount of methods allowed (class uses static memory)
    static const size_t MLIMIT = {{ cc_limit }}; 

    /// Evaluate CC while reading
    static int eval(const uint8_t *data, size_t size) {
        if (size == 0)
            return 1;
        
        size_t ptr = 0;
        auto c = target::c_list[data[ptr] % target::c_size];
        ptr++;

        if (ptr + c.args.bytesize > size)
            return 1;

        auto obj = c.fn(data + ptr);
        ptr += c.args.bytesize;

        if (ptr == size) // CC with only constructor
            return 1;

        auto m = target::m_list[data[ptr] % target::m_size];
        ptr++;

        for (size_t i = 0; i < MLIMIT; ++i) {
            if (ptr + m.args.bytesize > size)
                break;
            
            m.fn(&obj, data + ptr);
            ptr += m.args.bytesize;

            if (ptr == size) // No room for method id
                return 1;
            
            m = target::m_list[data[ptr] % target::m_size];
            ptr++;
        }

        return 0;
    }

    CallChain(const uint8_t *data, size_t size) : d1(data), d2(nullptr), d_size(size) { update(); }
    
    CallChain(uint8_t *data, size_t size) : d1(data), d2(data), d_size(size) { update(); }

    inline bool is_const() const {
        return d2 == nullptr;
    }

    inline bool is_empty() const {
        return c_size == 0;
    }

    inline bool is_methodless() const {
        return m_size == 0;
    }

    /// Returns amount of methods (not counting constructor)
    inline size_t get_msize() const {
        return m_size;
    }

    /// Get used bytes size
    inline size_t get_dsize() const {
        return d_size;
    }

    /// Get constructor id
    size_t get_cid() const {
        if (!is_empty())
            return d1[0] % target::c_size;
        throw std::logic_error("Attempt to get constructor id of an empty CC");
    }

    /// Get method id
    size_t get_mid(size_t n) const {
        if (!is_methodless())
            return d1[m[n % m_size]] % target::m_size;
        throw std::logic_error("Attempt to get method id of an empty CC");
    }

    /// Is constructor argless
    bool is_c_argless() const {
        if (is_empty())
            throw std::logic_error("Attempt to check if constructor argless when there is no constructor");
        return c_size == 1;
    }

    /// Get arg bytearray for constructor
    std::tuple<const uint8_t *, size_t> get_c_arg(size_t arg_n) const {
        if (is_empty())
            return std::make_tuple(nullptr, 0);

        auto c_args = target::c_list[d1[0] % target::c_size].args;

        if (c_args.size == 0)
            return std::make_tuple(nullptr, 0);

        arg_n %= c_args.size;
        
        size_t i, bytes = 1;
        for (i = 0; i < arg_n; ++i)
            bytes += c_args.list[i].size;
        
        return std::make_tuple(d1 + bytes, c_args.list[i].size);
    }

    /// Get mutable arg bytearray for constructor
    std::tuple<uint8_t *, size_t> get_c_arg(size_t arg_n) {
        if (is_empty() || is_const())
            return std::make_tuple(nullptr, 0);

        auto c_args = target::c_list[d2[0] % target::c_size].args;

        if (c_args.size == 0)
            return std::make_tuple(nullptr, 0);

        arg_n %= c_args.size;
        
        size_t i, bytes = 1;
        for (i = 0; i < arg_n; ++i)
            bytes += c_args.list[i].size;
        
        return std::make_tuple(d2 + bytes, c_args.list[i].size);
    }

    bool is_m_argless(size_t n) const {
        if (is_methodless())
            throw std::logic_error("Attempt to check if method argless when there is no methods");
        auto m_args = target::m_list[d1[m[n % m_size]] % target::m_size].args;
        return m_args.size == 0; 
    }

    /// Get arg bytearray for method
    std::tuple<const uint8_t *, size_t> get_m_arg(size_t n, size_t arg_n) const {
        if (is_methodless())
            return std::make_tuple(nullptr, 0);

        n %= m_size;
        auto m_args = target::m_list[d1[m[n]] % target::m_size].args;

        if (m_args.size == 0)
            return std::make_tuple(nullptr, 0);

        arg_n %= m_args.size;
        
        size_t i, bytes = 1;
        for (i = 0; i < arg_n; ++i)
            bytes += m_args.list[i].size;
        
        return std::make_tuple(d1 + m[n] + bytes, m_args.list[i].size);
    }

    /// Get mutable arg bytearray for method
    std::tuple<uint8_t *, size_t> get_m_arg(size_t n, size_t arg_n) {
        if (is_methodless() || is_const())
            return std::make_tuple(nullptr, 0);

        n %= m_size;
        auto m_args = target::m_list[d2[m[n]] % target::m_size].args;

        if (m_args.size == 0)
            return std::make_tuple(nullptr, 0);

        arg_n %= m_args.size;
        
        size_t i, bytes = 1;
        for (i = 0; i < arg_n; ++i)
            bytes += m_args.list[i].size;
        
        return std::make_tuple(d2 + m[n] + bytes, m_args.list[i].size);
    }

    /// Get offset of method
    size_t get_m_offset(size_t n) const {
        if (is_empty())
            return 0;
        if (is_methodless())
            return c_size;
        
        return m[n % m_size];
    }

    /// Remove method from CC
    /// Returns new byte len
    size_t rm_m(size_t n) {
        if (is_empty() || is_methodless() || is_const())
            return d_size;

        n %= m_size;

        // If it's last then just cut it off
        if (n == m_size - 1) {
            d_size = m[n];
            m_size--;
            return d_size;
        }

        // Shift everything in d2
        size_t offset = m[n + 1] - m[n];
        for (size_t i = m[n + 1]; i < d_size; ++i)
            d2[i - offset] = d2[i];
        d_size -= offset;

        // Shift values of m and substract offset
        m_size--;
        for (size_t i = n; i < m_size; ++i)
            m[i] = m[i + 1] - offset;

        return d_size;
    }

    /// Insert method
    /// Returns new byte length
    size_t ins_m(size_t n, size_t id, size_t max_size) {
        id %= target::m_size;
        size_t args = target::m_list[id].args.bytesize;
        if (is_empty() || is_const() || d_size + 1 + args > max_size || m_size + 1 >= MLIMIT)
            return d_size;
        
        // mod + 1 because we want to add methods to the start and end
        n = m_size == 0 ? 0 : n % (m_size + 1);

        // If it's last then it's simple
        if (n == m_size) {
            // Create new empty call
            d2[d_size] = id;
            for (size_t i = 0; i < args; ++i)
                d2[d_size + 1 + i] = 0;
            
            // Update class info
            m[m_size] = d_size;
            d_size += 1 + args;
            m_size++;

            return d_size;
        }
        
        // Shift everything
        for (size_t i = d_size - 1; i >= m[n]; --i)
            d2[i + 1 + args] = d2[i];
        
        // Set call
        d2[m[n]] = id;
        for (size_t i = 0; i < args; ++i)
            d2[m[n] + 1 + i] = 0;
        
        d_size += 1 + args;
        m_size++;

        // Shift offset values of m and add (1 + args) to them
        for (size_t i = m_size; i >= n + 1; --i)
            m[i] = m[i - 1] + 1 + args;

        return d_size;
    }

    /// Set constructor
    size_t set_c(size_t id, size_t max_size) {
        if (is_empty() || is_const())
            return d_size;
        id %= target::c_size;
        
        // Calculate change
        int offset = target::c_list[id].args.bytesize - (c_size - 1);
        if (d_size + offset > max_size)
            return d_size;

        // The simplest way
        if (offset == 0) {
            d2[0] = id;
            for (size_t i = 1; i < c_size; ++i)
                d2[i] = 0;
            return d_size;
        }
        
        // Shrink/expand d
        if (offset < 0) {
            for (size_t i = c_size; i < d_size; ++i)
                d2[i + offset] = d2[i];
        } else {
            for (size_t i = d_size; i >= c_size; --i)
                d2[i + offset] = d2[i];
        }

        // Fix m
        d_size += offset;
        for (size_t i = 0; i < m_size; ++i)
            m[i] += offset;
        
        // Set call
        d2[0] = id;
        c_size += offset;
        for (size_t i = 1; i < c_size; ++i)
            d2[i] = 0;
        
        return d_size;
    }

    /// Set method
    size_t set_m(size_t n, size_t id, size_t max_size) {
        if (is_methodless() || is_const())
            return d_size;
        n %= m_size;
        id %= target::m_size;

        // Calculate change
        size_t old_args = target::m_list[d2[m[n]] % target::m_size].args.bytesize;
        int offset = target::m_list[id].args.bytesize - old_args;

        if (d_size + offset > max_size)
            return d_size;

        // If we don't need to shift we just set everything
        // Same goes if our element is the last one
        if (offset == 0 || n == m_size - 1) {
            d2[m[n]] = id;
            for (size_t i = 0; i < target::m_list[id].args.bytesize; ++i)
                d2[m[n] + 1 + i] = 0;
            d_size += offset;
            return d_size;
        }

        // Shrink/expand d
        if (offset < 0) {
            for (size_t i = m[n + 1]; i < d_size; ++i)
                d2[i + offset] = d2[i];
        } else {
            for (size_t i = d_size - 1; i >= m[n + 1]; --i)
                d2[i + offset] = d2[i];
        }

        // Fix m
        d_size += offset;
        for (size_t i = n + 1; i < m_size; ++i)
            m[i] += offset;
        
        // Set call
        d2[m[n]] = id;
        for (size_t i = m[n] + 1; i < m[n + 1]; ++i)
            d2[i] = 0;
        
        return d_size;
    }

private:
    /// Source buffer data
    const uint8_t *d1;
    uint8_t *d2;
    size_t d_size; // Note that d_size is data **used**, not allocated/stored

    /// Constructor section length
    size_t c_size;

    // Method offsets
    size_t m[MLIMIT];
    size_t m_size;

    /// Update method data depending on d1
    void update() {
        m_size = 0;
        c_size = 0;

        if (d_size == 0)
            return;

        size_t offset = 1 + target::c_list[d1[0] % target::c_size].args.bytesize;
        if (offset > d_size)
            return;
        
        c_size = offset;
        
        size_t ptr;
        for (ptr = offset; ptr < d_size && m_size < MLIMIT; ++m_size) {
            m[m_size] = ptr;
            ptr += 1 + target::m_list[d1[ptr] % target::m_size].args.bytesize;
        }

        // ptr points to last call end
        // if it's greater than d_size then we have to trow away last call
        if (ptr != d_size) {
            // Roll back to last ptr
            d_size = m[m_size - 1];
            m_size--;
        }
    }
};

extern "C" size_t LLVMFuzzerCustomMutator(uint8_t *data, size_t size, size_t max_size, unsigned int seed) {
    if (size == 0)
        return size;

    CallChain cc(data, size);
    if (cc.is_empty())
        return LLVMFuzzerMutate(data, size, max_size);

    std::mt19937 rng(seed);
    size_t s = size;
    // TODO: add mutation probabilities and set them within generator
    for (size_t i = 0; i < 8; ++i) // TODO: make LIMIT customizable (generator param)
        switch (rng() % 4) {
            // Remove method
            case 0: {
                s = cc.rm_m(rng());
                break;
            }
            // Insert method
            case 1: {
                s = cc.ins_m(rng(), rng(), max_size);
                break;
            }
            // Set argument(s)
            case 2: {
                uint8_t *arg;
                size_t arg_size;

                // 0 is used to indicate constructor
                size_t call_k = rng(), call_n = rng() % (cc.get_msize() + 1);
                if (call_n == 0) {
                    auto c_args = target::c_list[cc.get_cid()].args;

                    // 0 is used to indicate mutation of all arguments at once
                    call_k %= c_args.size + 1;
                    size_t k = call_k == 0 ? call_k : call_k - 1;

                    std::tie(arg, arg_size) = cc.get_c_arg(k);
                    
                    // If we want to mutate all args we need to expand size
                    if (call_k == 0)
                        arg_size = c_args.bytesize;
                } else {
                    call_n -= 1;
                    auto m_args = target::m_list[cc.get_mid(call_n)].args;

                    // 0 is used to indicate mutation of all arguments at once
                    call_k %= m_args.size + 1;
                    size_t k = call_k == 0 ? call_k : call_k - 1;

                    std::tie(arg, arg_size) = cc.get_m_arg(call_n, k);

                    // If we want to mutate all args we need to expand size
                    if (call_k == 0)
                        arg_size = m_args.bytesize;
                }

                // Check if there is such arg
                if (arg == nullptr)
                    continue;

                // Fill argument with random bytes
                for (size_t i = 0; i < arg_size; ++i)
                    arg[i] = rng();
                
                break;
            }
            // Set constructor/method
            case 3: {
                // Zero used to indicate constructor
                size_t call_n = rng() % (cc.get_msize() + 1);

                if (call_n == 0)
                    s = cc.set_c(rng(), max_size);
                else
                    s = cc.set_m(call_n - 1, rng(), max_size);

                break;
            }
            default: {/* Should be unreachable */}
        }
    // TODO: add probabilty to this mutation
    s = LLVMFuzzerMutate(data, size, max_size);

    return s;
}

extern "C" size_t LLVMFuzzerCustomCrossOver(
    const uint8_t *data1, size_t size1,
    const uint8_t *data2, size_t size2,
    uint8_t *out, size_t max_out_size,
    unsigned int seed
) {
    CallChain cc1(data1, size1), cc2(data2, size2);

    // Check if crossover is not applicable
    if (cc1.is_empty() || cc2.is_methodless()) {
        for (size_t i = 0; i < size1; ++i)
            out[i] = data1[i];
        return size1;
    }

    // Pick random split points
    std::mt19937 rng(seed);
    // Mod with +1 is because:
    // 0 - we take constructor only
    // cc1.get_msize() - we take whole CC
    size_t n1 = rng() % (cc1.get_msize() + 1);
    // Here we take methods from p2 to cc2.get_msize()
    // Note that we can't take 0 methods
    size_t n2 = rng() % cc2.get_msize();
    
    // Check if we go out of bounds
    size_t size = n1 + (size2 - n2);
    if (size > max_out_size) {
        for (size_t i = 0; i < size1; ++i)
            out[i] = data1[i];
        return size1;
    }

    // Copy first part (we should be careful to not get_m_offset for n1 == m_size)
    size_t limit = n1 == cc1.get_msize() ? size1 : cc1.get_m_offset(n1);
    size_t i;
    for (i = 0; i < limit; ++i)
        out[i] = data1[i];
    
    // Copy second part
    for (size_t j = cc2.get_m_offset(n2); j < size2 && i < size; ++j, ++i)
        out[i] = data2[j];

    return size;
}


/// ------------------------------ HARNESS FN ------------------------------

extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
    CallChain::eval(data, size);
    return 0;
}