Commit d65da2db authored by Rusty Russell's avatar Rusty Russell

Much nicer example: genetic algo to approximate jpeg

(And cute image!)
parent 3fb9ac93
CFLAGS=-g -Wall -Wstrict-prototypes -Wold-style-definition -Wmissing-prototypes -Wmissing-declarations -I../../..
CFLAGS=-g -O3 -Wall -Wstrict-prototypes -Wold-style-definition -Wmissing-prototypes -Wmissing-declarations -I../../..
all: find_md5 md5_worker dns_lookup
find_md5: md5_server.c ../../../libccan.a
$(CC) $(CFLAGS) -o $@ $^
md5_worker: md5_worker.c ../../../libccan.a
$(CC) $(CFLAGS) -o $@ $^
all: dns_lookup
dns_lookup: dns_lookup.c ../../../libccan.a
$(CC) $(CFLAGS) -o $@ $^
$(CC) $(CFLAGS) -o $@ $^
arabella: arabella.c ../../../libccan.a
$(CC) $(CFLAGS) -o $@ $^ -ljpeg -lm
/* Antithread example to approximate picture with triangles using
* genetic algorithm. Example for a 2 cpu system:
* ./arabella arabella.jpg out out.save 2
*/
#include <stdio.h>
#include <jpeglib.h>
#include <ccan/talloc/talloc.h>
#include <err.h>
#include <assert.h>
#include <limits.h>
#include <string.h>
#include <stdint.h>
#include <math.h>
#include <sys/select.h>
#include <ccan/str/str.h>
#include <ccan/antithread/antithread.h>
#include <sys/types.h>
#include <unistd.h>
// Define this to run 100 times without dumping state
//#define BENCHMARK
/* How many drawings in entire population. */
#define POPULATION_SIZE 1000
/* How many generations without 1% improvement before we terminate. */
#define PLATEAU_GENS 200
/* An image buffer to render into. */
struct image {
unsigned height, width;
unsigned int stride;
/* RGB data */
unsigned char *buffer;
};
/* A drawing is a (fixed) number of triangles. */
struct drawing {
struct triangle *tri;
unsigned int num_tris;
unsigned long score;
};
/* Here's a triangle. */
struct triangle {
struct {
unsigned int x, y;
} coord[3];
/* RGBA */
unsigned char color[4];
unsigned char mult;
uint16_t add[3];
};
/* Get pointer into image at a specific location. */
static unsigned char *image_at(struct image *image,
unsigned int x, unsigned int y)
{
return image->buffer + y * image->stride + x * 3;
}
/* Blend a dot into this location. */
static void add_dot(unsigned char *buf,
unsigned char mult, const uint16_t add[])
{
unsigned int c;
/* /256 isn't quite right, but it's much faster than /255 */
for (c = 0; c < 3; c++)
buf[c] = (buf[c] * mult + add[c]) / 256;
}
/* Code based on example taken from:
* http://www.devmaster.net/forums/showthread.php?t=1094 */
static void add_flat_triangle(struct image *image,
int x0, int y0, int x1, int y1, int x2, int y2,
unsigned char mult, const uint16_t add[])
{
unsigned char *buf;
unsigned int i, j;
// compute slopes for the two triangle legs
float dx0 = (float)(x2 - x0) / (y2 - y0);
float dx1 = (float)(x2 - x1) / (y2 - y1);
int yRange = 0;
float lx = (float) x0, rx = (float) x1;
if (y0 < y2) {
yRange = y2 - y0;
buf = image_at(image, 0, y0);
} else {
yRange = y0 - y2;
buf = image_at(image, 0, y2);
lx = rx = (float)x2;
}
for (i=0; i < yRange; ++i) {
for (j=(int)(lx); j<(int)((rx) + 1.0f); ++j)
add_dot(buf + 3 * j, mult, add);
lx += dx0;
rx += dx1;
buf += image->stride;
}
}
static void swap(int *a, int *b)
{
int tmp = *a;
*a = *b;
*b = tmp;
}
static void paint_triangle(struct image *image, const struct triangle *tri)
{
int i0 = 0, i1 = 1, i2 = 2;
int x0, y0, x1, y1, x2, y2;
/* Could do this on triangle creation. */
// sort verts by height
if (tri->coord[i0].y > tri->coord[i1].y) swap(&i0, &i1);
if (tri->coord[i0].y > tri->coord[i2].y) swap(&i0, &i2);
if (tri->coord[i1].y > tri->coord[i2].y) swap(&i1, &i2);
x0 = tri->coord[i0].x, y0 = tri->coord[i0].y;
x1 = tri->coord[i1].x, y1 = tri->coord[i1].y;
x2 = tri->coord[i2].x, y2 = tri->coord[i2].y;
// test for easy cases, else split trinagle in two and render both halfs
if (y1 == y2) {
if (x1 > x2) swap(&x1, &x2);
add_flat_triangle(image, x1, y1, x2, y2, x0, y0,
tri->mult, tri->add);
} else if (y0 == y1) {
if (x0 > x1) swap(&x0, &x1);
add_flat_triangle(image, x0, y0, x1, y1, x2, y2,
tri->mult, tri->add);
} else {
// compute x pos of the vert that builds the splitting line with x1
int tmp_x = x0 + (int)(0.5f + (float)(y1-y0) * (float)(x2-x0) / (float)(y2-y0));
if (x1 > tmp_x) swap(&x1, &tmp_x);
add_flat_triangle(image, x1, y1, tmp_x, y1, x0, y0,
tri->mult, tri->add);
add_flat_triangle(image, x1, y1, tmp_x, y1, x2, y2,
tri->mult, tri->add);
}
}
/* Create a new image, allocated off context. */
static struct image *new_image(const void *ctx,
unsigned width, unsigned height, unsigned stride)
{
struct image *image = talloc(ctx, struct image);
image->width = width;
image->height = height;
image->stride = stride;
image->buffer = talloc_zero_array(image, unsigned char,
stride * height);
return image;
}
/* Taken from JPEG example code. Quality is 1 to 100. */
static void write_jpeg_file(const struct image *image,
const char *filename, int quality)
{
struct jpeg_compress_struct cinfo;
struct jpeg_error_mgr jerr;
FILE *outfile;
JSAMPROW row_pointer[1];
int row_stride;
cinfo.err = jpeg_std_error(&jerr);
jpeg_create_compress(&cinfo);
if ((outfile = fopen(filename, "wb")) == NULL)
err(1, "can't open %s", filename);
jpeg_stdio_dest(&cinfo, outfile);
cinfo.image_width = image->width;
cinfo.image_height = image->height;
cinfo.input_components = 3;
cinfo.in_color_space = JCS_RGB;
jpeg_set_defaults(&cinfo);
jpeg_set_quality(&cinfo, quality, TRUE);
jpeg_start_compress(&cinfo, TRUE);
row_stride = image->width * 3;
while (cinfo.next_scanline < cinfo.image_height) {
row_pointer[0] = image->buffer + cinfo.next_scanline*row_stride;
(void) jpeg_write_scanlines(&cinfo, row_pointer, 1);
}
jpeg_finish_compress(&cinfo);
fclose(outfile);
jpeg_destroy_compress(&cinfo);
}
static struct image *read_jpeg_file(const void *ctx, const char *filename)
{
struct jpeg_decompress_struct cinfo;
struct image *image;
struct jpeg_error_mgr jerr;
FILE *infile;
int row_stride;
if ((infile = fopen(filename, "rb")) == NULL)
err(1, "can't open %s", filename);
cinfo.err = jpeg_std_error(&jerr);
jpeg_create_decompress(&cinfo);
jpeg_stdio_src(&cinfo, infile);
jpeg_read_header(&cinfo, TRUE);
jpeg_start_decompress(&cinfo);
row_stride = cinfo.output_width * cinfo.output_components;
image = new_image(ctx,
cinfo.output_width, cinfo.output_height, row_stride);
while (cinfo.output_scanline < cinfo.output_height) {
JSAMPROW row = &image->buffer[cinfo.output_scanline*row_stride];
jpeg_read_scanlines(&cinfo, &row, 1);
}
(void) jpeg_finish_decompress(&cinfo);
jpeg_destroy_decompress(&cinfo);
fclose(infile);
return image;
}
/* Higher means closer to perfect match. We assume images same size. */
static unsigned long compare_images(const struct image *a,
const struct image *b)
{
unsigned long result = 0;
unsigned i;
/* Huge images won't work here. We'd need to get cleverer. */
assert(a->height * a->stride < ULONG_MAX / 8);
for (i = 0; i < a->height * a->stride; i++) {
if (a->buffer[i] > b->buffer[i])
result += a->buffer[i] - b->buffer[i];
else
result += b->buffer[i] - a->buffer[i];
}
return result;
}
/* Precalculate the alpha adds and multiplies for this color/alpha combo. */
static void calc_multipliers(struct triangle *tri)
{
/* Multiply by 255 - alpha. */
tri->mult = (255 - tri->color[3]);
/* Add alpha * color */
tri->add[0] = (tri->color[0] * tri->color[3]);
tri->add[1] = (tri->color[1] * tri->color[3]);
tri->add[2] = (tri->color[2] * tri->color[3]);
}
/* Render the image of this drawing, and return it. */
static struct image *image_of_drawing(const void *ctx,
const struct drawing *drawing,
const struct image *master)
{
struct image *image;
unsigned int i;
image = new_image(ctx, master->width, master->height, master->stride);
for (i = 0; i < drawing->num_tris; i++)
paint_triangle(image, &drawing->tri[i]);
return image;
}
/* Render the image and compare with the master. */
static void score_drawing(struct drawing *drawing,
const struct image *master)
{
struct image *image;
/* We don't allocate it off the drawing, since we don't need
* it inside the shared area. */
image = image_of_drawing(NULL, drawing, master);
drawing->score = compare_images(image, master);
talloc_free(image);
}
/* Create a new drawing allocated off context (which is the antithread
* pool context, so this is all allocated inside the pool so the
* antithreads can access it).
*/
static struct drawing *new_drawing(const void *ctx, unsigned int num_tris)
{
struct drawing *drawing = talloc_zero(ctx, struct drawing);
drawing->num_tris = num_tris;
drawing->tri = talloc_array(drawing, struct triangle, num_tris);
return drawing;
}
/* Make a random change to the drawing: frob one triangle. */
static void mutate_drawing(struct drawing *drawing,
const struct image *master)
{
unsigned int r = random();
struct triangle *tri = &drawing->tri[r % drawing->num_tris];
r /= drawing->num_tris;
r %= 10;
if (r < 6) {
if (r % 2)
tri->coord[r/2].x = random() % master->width;
else
tri->coord[r/2].y = random() % master->height;
} else {
tri->color[r - 6] = random() % 256;
}
calc_multipliers(tri);
}
/* Breed two drawings together, and throw in a mutation. */
static struct drawing *breed_drawing(const void *ctx,
const struct drawing *a,
const struct drawing *b,
const struct image *master)
{
unsigned int i;
struct drawing *drawing;
unsigned int r = random(), randmax = RAND_MAX;
assert(a->num_tris == b->num_tris);
drawing = new_drawing(ctx, a->num_tris);
for (i = 0; i < a->num_tris; i++) {
switch (r & 1) {
case 0:
/* Take from A. */
drawing->tri[i] = a->tri[i];
break;
case 1:
/* Take from B. */
drawing->tri[i] = b->tri[i];
break;
}
r >>= 1;
randmax >>= 1;
if (randmax == 0) {
r = random();
randmax = RAND_MAX;
}
}
mutate_drawing(drawing, master);
score_drawing(drawing, master);
return drawing;
}
/* This is our anti-thread. It does the time-consuming operation of
* breeding two drawings together and scoring the result. */
static void *breeder(struct at_pool *atp, const struct image *master)
{
const struct drawing *a, *b;
/* For simplicity, controller just hands us two pointers in
* separate writes. It could put them in the pool for us to
* read. */
while ((a = at_read_parent(atp)) != NULL) {
struct drawing *child;
b = at_read_parent(atp);
child = breed_drawing(at_pool_ctx(atp), a, b, master);
at_tell_parent(atp, child);
}
/* Unused: we never exit. */
return NULL;
}
/* We keep a very rough count of how much work each athread has. This
* works fine since fairly all rendering takes about the same time.
*
* Better alternative would be to put all the pending work somewhere
* in the shared area and notify any idle thread. The threads would
* keep looking in that shared area until they can't see any more
* work, then they'd at_tell_parent() back. */
struct athread_work {
struct athread *at;
unsigned int pending;
};
/* It's assumed that there isn't more than num results pending. */
static unsigned gather_results(struct athread_work *athreads,
unsigned int num_threads,
struct drawing **drawing,
unsigned int num,
bool block)
{
unsigned int i, maxfd = 0, done = 0;
struct timeval zero = { .tv_sec = 0, .tv_usec = 0 };
/* If it mattered, we could cache this fd mask and maxfd. */
for (i = 0; i < num_threads; i++) {
if (at_fd(athreads[i].at) > maxfd)
maxfd = at_fd(athreads[i].at);
}
do {
fd_set set;
FD_ZERO(&set);
for (i = 0; i < num_threads; i++)
FD_SET(at_fd(athreads[i].at), &set);
if (select(maxfd+1, &set, NULL, NULL, block ? NULL : &zero) < 0)
err(1, "Selecting on antithread results");
for (i = 0; i < num_threads; i++) {
if (!FD_ISSET(at_fd(athreads[i].at), &set))
continue;
*drawing = at_read(athreads[i].at);
if (!*drawing)
err(1, "Error with thread %u", i);
drawing++;
num--;
athreads[i].pending--;
done++;
}
} while (block && num);
return done;
}
/* Hand work to an antithread to breed these two together. */
static void tell_some_breeder(struct athread_work *athreads,
unsigned int num_threads,
const struct drawing *a, const struct drawing *b)
{
unsigned int i, best = 0;
/* Find least loaded thread. */
for (i = 1; i < num_threads; i++) {
if (athreads[i].pending < athreads[best].pending)
best = i;
}
at_tell(athreads[best].at, a);
at_tell(athreads[best].at, b);
athreads[best].pending++;
}
/* We seed initial triangles colours from the master image. */
static const unsigned char *initial_random_color(const struct image *master)
{
return master->buffer + (random() % (master->height * master->width))*3;
}
/* Create an initial random drawing. */
static struct drawing *random_drawing(const void *ctx,
const struct image *master,
unsigned int num_tris)
{
struct drawing *drawing = new_drawing(ctx, num_tris);
unsigned int i;
for (i = 0; i < drawing->num_tris; i++) {
unsigned int c;
struct triangle *tri = &drawing->tri[i];
for (c = 0; c < 3; c++) {
tri->coord[c].x = random() % master->width;
tri->coord[c].y = random() % master->height;
}
memcpy(tri->color, initial_random_color(master), 3);
tri->color[3] = (random() % 255) + 1;
calc_multipliers(tri);
}
score_drawing(drawing, master);
return drawing;
}
/* Read in a drawing from the saved state file. */
static struct drawing *read_drawing(const void *ctx, FILE *in,
const struct image *master)
{
struct drawing *drawing;
unsigned int i;
if (fscanf(in, "%u triangles:\n", &i) != 1)
errx(1, "Reading saved state");
drawing = new_drawing(ctx, i);
for (i = 0; i < drawing->num_tris; i++) {
unsigned int color[4];
if (fscanf(in, "%u,%u,%u,%u,%u,%u,%u,%u,%u,%u\n",
&drawing->tri[i].coord[0].x,
&drawing->tri[i].coord[0].y,
&drawing->tri[i].coord[1].x,
&drawing->tri[i].coord[1].y,
&drawing->tri[i].coord[2].x,
&drawing->tri[i].coord[2].y,
&color[0], &color[1], &color[2], &color[3]) != 10)
errx(1, "Reading saved state");
drawing->tri[i].color[0] = color[0];
drawing->tri[i].color[1] = color[1];
drawing->tri[i].color[2] = color[2];
drawing->tri[i].color[3] = color[3];
calc_multipliers(&drawing->tri[i]);
}
score_drawing(drawing, master);
return drawing;
}
/* Comparison function for sorting drawings best to worst. */
static int compare_drawing_scores(const void *_a, const void *_b)
{
struct drawing **a = (void *)_a, **b = (void *)_b;
return (*a)->score - (*b)->score;
}
/* Save one drawing to state file */
static void dump_drawings(struct drawing **drawing, const char *outname)
{
FILE *out;
unsigned int i, j;
char *tmpout = talloc_asprintf(NULL, "%s.tmp", outname);
out = fopen(tmpout, "w");
if (!out)
err(1, "Opening %s", tmpout);
fprintf(out, "POPULATION_SIZE=%u\n", POPULATION_SIZE);
for (i = 0; i < POPULATION_SIZE; i++) {
fprintf(out, "%u triangles:\n", drawing[i]->num_tris);
for (j = 0; j < drawing[i]->num_tris; j++) {
fprintf(out, "%u,%u,%u,%u,%u,%u,%u,%u,%u,%u\n",
drawing[i]->tri[i].coord[0].x,
drawing[i]->tri[i].coord[0].y,
drawing[i]->tri[i].coord[1].x,
drawing[i]->tri[i].coord[1].y,
drawing[i]->tri[i].coord[2].x,
drawing[i]->tri[i].coord[2].y,
drawing[i]->tri[j].color[0],
drawing[i]->tri[j].color[1],
drawing[i]->tri[j].color[2],
drawing[i]->tri[j].color[3]);
}
}
if (fclose(out) != 0)
err(1, "Failure closing %s", tmpout);
if (rename(tmpout, outname) != 0)
err(1, "Renaming %s over %s", tmpout, outname);
talloc_free(tmpout);
}
/* Save state file. */
static void dump_state(struct drawing *drawing[POPULATION_SIZE],
const struct image *master,
const char *outpic,
const char *outstate,
unsigned int gen)
{
char *out = talloc_asprintf(NULL, "%s.%08u.jpg", outpic, gen);
struct image *image;
printf("Dumping gen %u to %s & %s\n", gen, out, outstate);
dump_drawings(drawing, outstate);
image = image_of_drawing(out, drawing[0], master);
write_jpeg_file(image, out, 80);
talloc_free(out);
}
/* Biassed coinflip moves us towards top performers. I didn't spend
* too much time on it, but 1/32 seems to give decent results (see
* breeding-algos.gnumeric). */
static struct drawing *select_good_drawing(struct drawing *drawing[],
unsigned int population)
{
if (population == 1)
return drawing[0];
if (random() % 32)
return select_good_drawing(drawing, population/2);
return select_good_drawing(drawing + population/2, population/2);
}
static void usage(void)
{
errx(1, "usage: <infile> <outfile> <statefile> <numtriangles> <numcpus> [<instatefile>]");
}
int main(int argc, char *argv[])
{
struct image *master;
unsigned int gen, since_prev_best, num_threads, i;
struct drawing *drawing[POPULATION_SIZE];
unsigned long prev_best, poolsize;
struct at_pool *atp;
struct athread_work *athreads;
if (argc != 6 && argc != 7)
usage();
/* Room for triangles and master image, with some spare.
* We should really read master image header first, so we can be
* more precise than "about 3MB". ccan/alloc also needs some
* more work to be more efficient. */
poolsize = (POPULATION_SIZE + POPULATION_SIZE/4) * (sizeof(struct drawing) + atoi(argv[4]) * sizeof(struct triangle)) * 2 + 1000 * 1000 * 3;
atp = at_pool(poolsize);
if (!atp)
err(1, "Creating pool of %lu bytes", poolsize);
/* Auto-free the pool and anything hanging off it (eg. threads). */
talloc_steal(talloc_autofree_context(), atp);
/* Read in file */
master = read_jpeg_file(at_pool_ctx(atp), argv[1]);
if (argc == 6) {
printf("Creating initial population");
fflush(stdout);
for (i = 0; i < POPULATION_SIZE; i++) {
drawing[i] = random_drawing(at_pool_ctx(atp),
master, atoi(argv[4]));
printf(".");
fflush(stdout);
}
printf("\n");
} else {
FILE *state;
char header[100];
state = fopen(argv[5], "r");
if (!state)
err(1, "Opening %s", argv[5]);
fflush(stdout);
fgets(header, 100, state);
printf("Loading initial population from %s: %s", argv[5],
header);
for (i = 0; i < POPULATION_SIZE; i++) {
drawing[i] = read_drawing(at_pool_ctx(atp),
state, master);
printf(".");
fflush(stdout);
}
}
num_threads = atoi(argv[5]);
if (!num_threads)
usage();
/* Hang the threads off the pool (not *in* the pool). */
athreads = talloc_array(atp, struct athread_work, num_threads);
for (i = 0; i < num_threads; i++) {
athreads[i].pending = 0;
athreads[i].at = at_run(atp, breeder, master);
if (!athreads[i].at)
err(1, "Creating antithread %u", i);
}
since_prev_best = 0;
/* Worse than theoretically worst case. */
prev_best = master->height * master->stride * 256;
for (gen = 0; since_prev_best < PLATEAU_GENS; gen++) {
unsigned int j, done = 0;
struct drawing *new[POPULATION_SIZE/4];
qsort(drawing, POPULATION_SIZE, sizeof(drawing[0]),
compare_drawing_scores);
printf("Best %lu, worst %lu\n",
drawing[0]->score, drawing[POPULATION_SIZE-1]->score);
/* Probability of being chosen to breed depends on
* rank. We breed over lowest 1/4 population. */
for (j = 0; j < POPULATION_SIZE / 4; j++) {
struct drawing *best1, *best2;
best1 = select_good_drawing(drawing, POPULATION_SIZE);
best2 = select_good_drawing(drawing, POPULATION_SIZE);
tell_some_breeder(athreads, num_threads, best1, best2);
/* We reap during loop, so return pipes don't fill.
* See "Better alternative" above. */
done += gather_results(athreads, num_threads,
new + done, j - done, false);
}
/* Collate final results. */
gather_results(athreads, num_threads, new+done, j-done, true);
/* Overwrite bottom 1/4 */
for (j = POPULATION_SIZE * 3 / 4; j < POPULATION_SIZE; j++) {
talloc_free(drawing[j]);
drawing[j] = new[j - POPULATION_SIZE * 3 / 4];
}
/* We dump on every 1% improvement in score. */
if (drawing[0]->score < prev_best * 0.99) {
#ifndef BENCHMARK
dump_state(drawing, master, argv[2], argv[3], gen);
#endif
prev_best = drawing[0]->score;
since_prev_best = 0;
} else
since_prev_best++;
#ifdef BENCHMARK
if (gen == 100)
exit(0);
#endif
}
/* Dump final state */
printf("No improvement over %lu for %u gens\n",
prev_best, since_prev_best);
dump_state(drawing, master, argv[2], argv[3], gen);
return 0;
}
/* Tries to find data with a given MD5 (up to N bits). */
#include "ccan/antithread/antithread.h"
#include "ccan/string/string.h"
#include "ccan/talloc/talloc.h"
#include "md5_finder.h"
#include <err.h>
#include <sys/select.h>
#include <stdio.h>
#include <stdlib.h>
static void usage(void)
{
errx(1, "Usage: md5calc <hexstring> <numcpus>");
}
static void parse_hexstring(const char *string, struct md5_search *md5s)
{
unsigned int i;
if (strstarts(string, "0x") || strstarts(string, "0X"))
string += 2;
for (i = 0; i < MD5_HASH_WORDS; i++) {
unsigned int n[4], j;
int ret;
ret = sscanf(string, "%02x%02x%02x%02x",
&n[0], &n[1], &n[2], &n[3]);
string += 8;
if (ret == EOF)
break;
for (j = 0; j < ret; j++) {
md5s->mask[MD5_HASH_WORDS-i-1] |= (0xFF << (8*j));
md5s->md5[MD5_HASH_WORDS-i-1] |= (n[j] << (8*j));
}
if (ret != 4)
break;
}
}
static void init_pattern(u8 *pattern, unsigned int num_bytes, u64 total)
{
unsigned int i;
for (i = 0; i < num_bytes; i++) {
pattern[i] = 'A' + (total % 26);
total /= 26;
}
}
#define PATTERN_BYTES 32
int main(int argc, char *argv[])
{
struct at_pool *atp;
struct md5_search md5s;
unsigned int i, maxfd, numathreads = argc == 3 ? atoi(argv[2]) : 0;
u64 total = 0;
fd_set fds;
char *cmdline[] = { "./md5_worker", NULL };
struct athread *at[numathreads];
if (numathreads == 0)
usage();
memset(&md5s, 0, sizeof(md5s));
parse_hexstring(argv[1], &md5s);
md5s.num_tries = 1024*1024;
md5s.num_bytes = PATTERN_BYTES;
/* *2 to allow for allocation inefficiency. */
atp = at_pool((sizeof(md5s) + PATTERN_BYTES) * (numathreads + 1) * 2);
if (!atp)
err(1, "Can't create pool");
/* Free pool on exit. */
// talloc_steal(talloc_autofree_context(), atp);
FD_ZERO(&fds);
maxfd = 0;
for (i = 0; i < numathreads; i++) {
at[i] = at_spawn(atp, NULL, cmdline);
if (!at[i])
err(1, "Can't spawn child");
FD_SET(at_fd(at[i]), &fds);
if (at_fd(at[i]) > maxfd)
maxfd = at_fd(at[i]);
}
for (;;) {
struct md5_search *m, *res;
fd_set in = fds;
/* Shouldn't fail! */
m = talloc(at_pool_ctx(atp), struct md5_search);
*m = md5s;
md5s.num_tries++;
m->pattern = talloc_array(m, u8, m->num_bytes);
init_pattern(m->pattern, m->num_bytes, total);
select(maxfd+1, &in, NULL, NULL, NULL);
for (i = 0; i < numathreads; i++)
if (FD_ISSET(at_fd(at[i]), &in))
break;
if (i == numathreads)
errx(1, "Select returned, but noone ready?");
res = at_read(at[i]);
if (res == NULL) {
warn("Thread died?");
FD_CLR(at_fd(at[i]), &fds);
continue;
}
if (res != INITIAL_POINTER) {
if (res->success) {
printf("Success! '%.*s'\n",
res->num_bytes, (char *)res->pattern);
exit(0);
}
m->num_tries++;
talloc_free(res);
}
at_tell(at[i], m);
total += m->num_tries;
}
}
/* Worker thread: tries to find data with given MD5. */
#include "ccan/antithread/antithread.h"
#include "md5_finder.h"
#include <netinet/in.h>
#include <err.h>
#include <string.h>
/*
* Cryptographic API.
*
* MD5 Message Digest Algorithm (RFC1321).
*
* Derived from cryptoapi implementation, originally based on the
* public domain implementation written by Colin Plumb in 1993.
*
* Copyright (c) Cryptoapi developers.
* Copyright (c) 2002 James Morris <jmorris@intercode.com.au>
*
* This program is free software; you can redistribute it and/or modify it
* under the terms of the GNU General Public License as published by the Free
* Software Foundation; either version 2 of the License, or (at your option)
* any later version.
*/
#define MD5_DIGEST_SIZE 16
#define MD5_HMAC_BLOCK_SIZE 64
#define MD5_BLOCK_WORDS 16
#define F1(x, y, z) (z ^ (x & (y ^ z)))
#define F2(x, y, z) F1(z, x, y)
#define F3(x, y, z) (x ^ y ^ z)
#define F4(x, y, z) (y ^ (x | ~z))
#define MD5STEP(f, w, x, y, z, in, s) \
(w += f(x, y, z) + in, w = (w<<s | w>>(32-s)) + x)
struct md5_ctx {
u32 hash[MD5_HASH_WORDS];
u32 block[MD5_BLOCK_WORDS];
u64 byte_count;
};
static void md5_transform(u32 *hash, u32 const *in)
{
u32 a, b, c, d;
a = hash[0];
b = hash[1];
c = hash[2];
d = hash[3];
MD5STEP(F1, a, b, c, d, in[0] + 0xd76aa478, 7);
MD5STEP(F1, d, a, b, c, in[1] + 0xe8c7b756, 12);
MD5STEP(F1, c, d, a, b, in[2] + 0x242070db, 17);
MD5STEP(F1, b, c, d, a, in[3] + 0xc1bdceee, 22);
MD5STEP(F1, a, b, c, d, in[4] + 0xf57c0faf, 7);
MD5STEP(F1, d, a, b, c, in[5] + 0x4787c62a, 12);
MD5STEP(F1, c, d, a, b, in[6] + 0xa8304613, 17);
MD5STEP(F1, b, c, d, a, in[7] + 0xfd469501, 22);
MD5STEP(F1, a, b, c, d, in[8] + 0x698098d8, 7);
MD5STEP(F1, d, a, b, c, in[9] + 0x8b44f7af, 12);
MD5STEP(F1, c, d, a, b, in[10] + 0xffff5bb1, 17);
MD5STEP(F1, b, c, d, a, in[11] + 0x895cd7be, 22);
MD5STEP(F1, a, b, c, d, in[12] + 0x6b901122, 7);
MD5STEP(F1, d, a, b, c, in[13] + 0xfd987193, 12);
MD5STEP(F1, c, d, a, b, in[14] + 0xa679438e, 17);
MD5STEP(F1, b, c, d, a, in[15] + 0x49b40821, 22);
MD5STEP(F2, a, b, c, d, in[1] + 0xf61e2562, 5);
MD5STEP(F2, d, a, b, c, in[6] + 0xc040b340, 9);
MD5STEP(F2, c, d, a, b, in[11] + 0x265e5a51, 14);
MD5STEP(F2, b, c, d, a, in[0] + 0xe9b6c7aa, 20);
MD5STEP(F2, a, b, c, d, in[5] + 0xd62f105d, 5);
MD5STEP(F2, d, a, b, c, in[10] + 0x02441453, 9);
MD5STEP(F2, c, d, a, b, in[15] + 0xd8a1e681, 14);
MD5STEP(F2, b, c, d, a, in[4] + 0xe7d3fbc8, 20);
MD5STEP(F2, a, b, c, d, in[9] + 0x21e1cde6, 5);
MD5STEP(F2, d, a, b, c, in[14] + 0xc33707d6, 9);
MD5STEP(F2, c, d, a, b, in[3] + 0xf4d50d87, 14);
MD5STEP(F2, b, c, d, a, in[8] + 0x455a14ed, 20);
MD5STEP(F2, a, b, c, d, in[13] + 0xa9e3e905, 5);
MD5STEP(F2, d, a, b, c, in[2] + 0xfcefa3f8, 9);
MD5STEP(F2, c, d, a, b, in[7] + 0x676f02d9, 14);
MD5STEP(F2, b, c, d, a, in[12] + 0x8d2a4c8a, 20);
MD5STEP(F3, a, b, c, d, in[5] + 0xfffa3942, 4);
MD5STEP(F3, d, a, b, c, in[8] + 0x8771f681, 11);
MD5STEP(F3, c, d, a, b, in[11] + 0x6d9d6122, 16);
MD5STEP(F3, b, c, d, a, in[14] + 0xfde5380c, 23);
MD5STEP(F3, a, b, c, d, in[1] + 0xa4beea44, 4);
MD5STEP(F3, d, a, b, c, in[4] + 0x4bdecfa9, 11);
MD5STEP(F3, c, d, a, b, in[7] + 0xf6bb4b60, 16);
MD5STEP(F3, b, c, d, a, in[10] + 0xbebfbc70, 23);
MD5STEP(F3, a, b, c, d, in[13] + 0x289b7ec6, 4);
MD5STEP(F3, d, a, b, c, in[0] + 0xeaa127fa, 11);
MD5STEP(F3, c, d, a, b, in[3] + 0xd4ef3085, 16);
MD5STEP(F3, b, c, d, a, in[6] + 0x04881d05, 23);
MD5STEP(F3, a, b, c, d, in[9] + 0xd9d4d039, 4);
MD5STEP(F3, d, a, b, c, in[12] + 0xe6db99e5, 11);
MD5STEP(F3, c, d, a, b, in[15] + 0x1fa27cf8, 16);
MD5STEP(F3, b, c, d, a, in[2] + 0xc4ac5665, 23);
MD5STEP(F4, a, b, c, d, in[0] + 0xf4292244, 6);
MD5STEP(F4, d, a, b, c, in[7] + 0x432aff97, 10);
MD5STEP(F4, c, d, a, b, in[14] + 0xab9423a7, 15);
MD5STEP(F4, b, c, d, a, in[5] + 0xfc93a039, 21);
MD5STEP(F4, a, b, c, d, in[12] + 0x655b59c3, 6);
MD5STEP(F4, d, a, b, c, in[3] + 0x8f0ccc92, 10);
MD5STEP(F4, c, d, a, b, in[10] + 0xffeff47d, 15);
MD5STEP(F4, b, c, d, a, in[1] + 0x85845dd1, 21);
MD5STEP(F4, a, b, c, d, in[8] + 0x6fa87e4f, 6);
MD5STEP(F4, d, a, b, c, in[15] + 0xfe2ce6e0, 10);
MD5STEP(F4, c, d, a, b, in[6] + 0xa3014314, 15);
MD5STEP(F4, b, c, d, a, in[13] + 0x4e0811a1, 21);
MD5STEP(F4, a, b, c, d, in[4] + 0xf7537e82, 6);
MD5STEP(F4, d, a, b, c, in[11] + 0xbd3af235, 10);
MD5STEP(F4, c, d, a, b, in[2] + 0x2ad7d2bb, 15);
MD5STEP(F4, b, c, d, a, in[9] + 0xeb86d391, 21);
hash[0] += a;
hash[1] += b;
hash[2] += c;
hash[3] += d;
}
/* XXX: this stuff can be optimized */
static inline void le32_to_cpu_array(u32 *buf, unsigned int words)
{
while (words--) {
*buf = ntohl(*buf);
buf++;
}
}
static inline void cpu_to_le32_array(u32 *buf, unsigned int words)
{
while (words--) {
*buf = htonl(*buf);
buf++;
}
}
static inline void md5_transform_helper(struct md5_ctx *ctx)
{
le32_to_cpu_array(ctx->block, sizeof(ctx->block) / sizeof(u32));
md5_transform(ctx->hash, ctx->block);
}
static void md5_init(struct md5_ctx *mctx)
{
mctx->hash[0] = 0x67452301;
mctx->hash[1] = 0xefcdab89;
mctx->hash[2] = 0x98badcfe;
mctx->hash[3] = 0x10325476;
mctx->byte_count = 0;
}
static void md5_update(struct md5_ctx *mctx, const u8 *data, unsigned int len)
{
const u32 avail = sizeof(mctx->block) - (mctx->byte_count & 0x3f);
mctx->byte_count += len;
if (avail > len) {
memcpy((char *)mctx->block + (sizeof(mctx->block) - avail),
data, len);
return;
}
memcpy((char *)mctx->block + (sizeof(mctx->block) - avail),
data, avail);
md5_transform_helper(mctx);
data += avail;
len -= avail;
while (len >= sizeof(mctx->block)) {
memcpy(mctx->block, data, sizeof(mctx->block));
md5_transform_helper(mctx);
data += sizeof(mctx->block);
len -= sizeof(mctx->block);
}
memcpy(mctx->block, data, len);
}
static void md5_final(struct md5_ctx *mctx)
{
const unsigned int offset = mctx->byte_count & 0x3f;
char *p = (char *)mctx->block + offset;
int padding = 56 - (offset + 1);
*p++ = 0x80;
if (padding < 0) {
memset(p, 0x00, padding + sizeof (u64));
md5_transform_helper(mctx);
p = (char *)mctx->block;
padding = 56;
}
memset(p, 0, padding);
mctx->block[14] = mctx->byte_count << 3;
mctx->block[15] = mctx->byte_count >> 29;
le32_to_cpu_array(mctx->block, (sizeof(mctx->block) -
sizeof(u64)) / sizeof(u32));
md5_transform(mctx->hash, mctx->block);
cpu_to_le32_array(mctx->hash, sizeof(mctx->hash) / sizeof(u32));
}
static bool bits_match(const u32 a[MD5_HASH_WORDS],
const u32 b[MD5_HASH_WORDS],
const u32 mask[MD5_HASH_WORDS])
{
unsigned int i;
for (i = 0; i < MD5_HASH_WORDS; i++) {
if ((a[i] & mask[i]) != (b[i] & mask[i]))
return false;
}
#if 0
printf("mask = %08x%08x%08x%08x\n"
"a = %08x%08x%08x%08x\n"
"b = %08x%08x%08x%08x\n",
mask[0], mask[1], mask[2], mask[3],
a[0], a[1], a[2], a[3],
b[0], b[1], b[2], b[3]);
#endif
return true;
}
static void inc_pattern(u8 *pattern, unsigned int len)
{
unsigned int i;
for (i = 0; i < len; i++) {
pattern[i]++;
if (pattern[i] <= 'Z')
break;
pattern[i] = 'A';
}
}
int main(int argc, char *argv[])
{
struct at_pool *atp = at_get_pool(&argc, argv, NULL);
struct md5_search *md5s;
if (!atp)
err(1, "Not a worker thread?");
/* Tell parent we're ready. */
at_tell_parent(atp, INITIAL_POINTER);
while ((md5s = at_read_parent(atp)) != NULL) {
unsigned int i;
md5s->success = false;
for (i = 0; i < md5s->num_tries; i++) {
struct md5_ctx ctx;
md5_init(&ctx);
md5_update(&ctx, md5s->pattern, md5s->num_bytes);
md5_final(&ctx);
if (bits_match(ctx.hash, md5s->md5, md5s->mask)) {
md5s->success = true;
break;
}
inc_pattern(md5s->pattern, md5s->num_bytes);
}
at_tell_parent(atp, md5s);
}
return 0;
}
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment