ec_reconstruct_fuzz.py

#!/usr/bin/env python
from argparse import ArgumentParser
from hashlib import md5
import math
import random
import sys

from pyeclib.ec_iface import ECDriver, VALID_EC_TYPES

SEGSIZE = 1048576
SEGNUM = 5


parser = ArgumentParser()
parser.add_argument('--ec-type', default=VALID_EC_TYPES[0],
                    choices=VALID_EC_TYPES)
parser.add_argument('--segment-size', type=int, default=SEGSIZE)
parser.add_argument('--num-segments', type=int, default=SEGNUM,
                    help='number of segments in made up body')
parser.add_argument('--body', help='path to body')
parser.add_argument('-k', type=int, default=6,
                    help='number of data elements')
parser.add_argument('-m', type=int, default=4,
                    help='number of parity elements')
parser.add_argument('i', type=int, help='fragment index to reconstruct')
parser.add_argument('j', type=int, nargs='+',
                    help='fragment indexes to use for reconstruct')


def make_fragment_payloads(driver, test_body, segsize=SEGSIZE):
    # split up the body into buffers
    chunks = [test_body[x:x + segsize]
              for x in range(0, len(test_body), segsize)]
    # encode the buffers into fragment payloads
    fragment_payloads = []
    for chunk in chunks:
        fragments = driver.encode(chunk)
        if not fragments:
            break
        fragment_payloads.append(fragments)
    return fragment_payloads


def main():
    args = parser.parse_args()
    if len(args.j) < args.k:
        return ('ERROR: need at least %s indexes '
                'given as j to reconstruct' % args.k)
    if args.body:
        test_body = open(args.body).read()
    else:
        num_segments = 'crazy body we just make up'[:args.num_segments]
        test_body = num_segments * args.segment_size
        # trim a random bit from the last segment
        test_body = test_body[:-random.randint(0, args.segment_size)]
    driver = ECDriver(k=args.k, m=args.m, ec_type=args.ec_type)

    # generate fragment payloads
    fragment_payloads = make_fragment_payloads(driver, test_body,
                                               segsize=args.segment_size)
    info = {
        'body length': len(test_body),
        'segment size': args.segment_size,
        'num segments': int(math.ceil(
            len(test_body) / float(args.segment_size))),
        'num fragments': len(fragment_payloads),
    }
    # python y u no ordereddict literal!?  [a:1, b:2] would be shawheat!
    key_order = [
        'body length',
        'segment size',
        'num segments',
        'num fragments'
    ]
    for k in key_order:
        print '%20s: %s' % (k, info.pop(k))
    # catch all for left over info not in key_order
    for k, v in info.items():
        print '%20s: %s' % (k, v)

    # this is the archive we expect to rebuild
    expected_archive = ''.join(frags[args.i] for frags in fragment_payloads)

    reconstructed_fragments = []
    for frags in fragment_payloads:
        fragment_payload = [frags[j] for j in args.j]
        fragment = driver.reconstruct(fragment_payload, [args.i])[0]
        reconstructed_fragments.append(fragment)

    reconstructed_archive = ''.join(reconstructed_fragments)

    if len(expected_archive) != len(reconstructed_archive):
        return 'ERROR: len(expected) %s != len(reconstructed) %s' % (
            len(expected_archive), len(reconstructed_archive))

    expected_archive_hash = md5(expected_archive).hexdigest()
    reconstructed_archive_hash = md5(reconstructed_archive).hexdigest()

    if expected_archive_hash != reconstructed_archive_hash:
        return 'ERROR: expected %s != reconstructed %s' % (
            expected_archive_hash, reconstructed_archive_hash)

    print 'SUCCESS: %s == %s' % (
        expected_archive_hash, reconstructed_archive_hash)


if __name__ == '__main__':
    sys.exit(main())

seg_fault_repro

$ ec_reconstruct_fuzz.py -k 3 -m 2 2 1 3 4 --ec-type isa_l_rs_vand
         body length: 4457535
        segment size: 1048576
        num segments: 5
       num fragments: 5
SUCCESS: 98daf604ced63494169f70be9b892a37 == 98daf604ced63494169f70be9b892a37
$ ec_reconstruct_fuzz.py -k 3 -m 2 1 1 3 4 --ec-type isa_l_rs_vand
         body length: 4588725
        segment size: 1048576
        num segments: 5
       num fragments: 5
Segmentation fault (core dumped)
$ ec_reconstruct_fuzz.py -k 3 -m 2 1 1 3 4 --ec-type jerasure_rs_vand
         body length: 4551565
        segment size: 1048576
        num segments: 5
       num fragments: 5
SUCCESS: 203f92a4796f89463b9ef1699b4e64f4 == 203f92a4796f89463b9ef1699b4e64f4