selfpatch-slr/selfpatch/src/targets.c

#include "targets.h"
#include "spslr_program.h"

#include <stdlib.h>
#include <time.h>

static void seed_rand_time() {
	srand(time(NULL));
}

static uint32_t rand_uint32() {
	return (uint32_t)rand();
}

struct Field {
	uint32_t initial_offset;
	uint32_t initial_idx;

	uint32_t offset;
	uint32_t size;
	uint32_t alignment;
	uint32_t flags;
};

struct FinalField {
	uint32_t initial_offset;
	uint32_t offset;
};

struct Target {
	uint32_t uid;
	uint32_t size;
	uint32_t field_count;

	uint32_t present_field_count;
	struct Field* fields;
	struct FinalField* final_fields; // After rand (indexable by initial layout)

	struct Target* next;
};

static struct Target* targets = NULL;

static struct Target* find_target(uint32_t uid) {
	struct Target* cur = targets;
	while (cur) {
		if (cur->uid == uid)
			return cur;

		cur = cur->next;
	}

	return NULL;
}

static int add_target(uint32_t uid, uint32_t size, uint32_t fieldcnt) {
	if (find_target(uid))
		return -1;

	struct Target* new_target = (struct Target*)malloc(sizeof(struct Target));
	if (!new_target)
		return -1;

	new_target->uid = uid;
	new_target->size = size;
	new_target->field_count = fieldcnt;

	new_target->present_field_count = 0;
	new_target->final_fields = NULL;
	new_target->fields = (struct Field*)malloc(sizeof(struct Field) * fieldcnt);

	if (!new_target->fields) {
		free(new_target);
		return -1;
	}

	new_target->next = targets;
	targets = new_target;
	return 0;
}

static void clear_targets() {
	struct Target* cur = targets;
	while (cur) {
		if (cur->fields)
			free(cur->fields);
		if (cur->final_fields)
			free(cur->final_fields);

		struct Target* tmp = cur;
		cur = cur->next;

		free(tmp);
	}

	targets = NULL;
}

int spslr_target(uint32_t uid, uint32_t size, uint32_t fieldcnt) {
	if (find_target(uid))
		return -1;

	return add_target(uid, size, fieldcnt);
}

int spslr_field(uint32_t target, uint32_t offset, uint32_t size, uint32_t alignment, uint32_t flags) {
	struct Target* t = find_target(target);
	if (!t)
		return -1;

	if (offset + size > t->size)
		return -1;

	if (t->present_field_count >= t->field_count)
		return -1;

	if (offset % alignment != 0)
		return -1;

	if (t->present_field_count) {
		struct Field* pred = &t->fields[t->present_field_count - 1];
		if (offset < pred->offset + pred->size)
			return -1;
	}

	uint32_t idx = t->present_field_count++;

	struct Field* f = &t->fields[idx];
	f->initial_offset = offset;
	f->initial_idx = idx;
	f->offset = offset;
	f->size = size;
	f->alignment = alignment;
	f->flags = flags;

	return 0;
}

struct ShuffleRegion {
	uint32_t begin, end;
	uint32_t fill_begin, fill_end;
};

static void target_get_origin_region(const struct Target* target, struct ShuffleRegion* region, uint32_t idx) {
	region->fill_begin = target->fields[idx].offset;
	region->fill_end = region->fill_begin + target->fields[idx].size;

	if (idx <= 0)
		region->begin = 0;
	else
		region->begin = target->fields[idx - 1].offset + target->fields[idx - 1].size;

	if (idx >= target->field_count - 1)
		region->end = target->size;
	else
		region->end = target->fields[idx + 1].offset;
}

static uint32_t target_get_shuffle_options(const struct Target* target, uint32_t* options,
		const struct ShuffleRegion* origin, uint32_t alignment) {
	uint32_t count = 0;

	uint32_t current_field = 0;
	for (uint32_t offset = 0; offset < target->size; offset += alignment) {
		// Placing the origin region here or further would exceed struct boundaries
		uint32_t option_would_end = offset + (origin->fill_end - origin->fill_begin);
		if (option_would_end > target->size)
			break;

		// Fields that end before offset are irrelevant at this point
		while (current_field < target->field_count && target->fields[current_field].offset
			+ target->fields[current_field].size <= offset) current_field++;

		// If origin->fill_begin was placed at offset, could all fields that overlap with range swap with origin?
		int conflict = 0;
		uint32_t origin_region_ptr = origin->begin;
		for (uint32_t it = current_field; it < target->field_count; it++) {
			struct Field* f = &target->fields[it];
			if (f->offset >= option_would_end)
				break;

			if (f->flags & SPSLR_FLAG_FIELD_FIXED) {
				conflict = 1;
				break;
			}

			// Could field be placed in origin?
			if (origin_region_ptr % f->alignment != 0)
				origin_region_ptr += f->alignment - (origin_region_ptr % f->alignment);

			origin_region_ptr += f->size;

			if (origin_region_ptr > origin->end) {
				conflict = 1;
				break;
			}
		}

		if (!conflict)
			options[count++] = offset;
	}

	return count;
}

static void target_swap(struct Target* target, uint32_t origin_idx, const struct ShuffleRegion* origin_region,
		uint32_t new_offset) {
	int origin_pulled = 0;
	uint32_t origin_region_ptr = origin_region->begin;
	uint32_t option_fill_end = new_offset + (origin_region->fill_end - origin_region->fill_begin);
	for (uint32_t it = 0; it < target->field_count; it++) {
		// Fields that end before new_offset can stay there
		if (target->fields[it].offset + target->fields[it].size <= new_offset)
			continue;

		// Fields that start after where origin goes are irrelevant
		if (target->fields[it].offset >= option_fill_end)
			break;

		uint32_t falign = target->fields[it].alignment;
		if (origin_region_ptr % falign != 0)
			origin_region_ptr += falign - (origin_region_ptr % falign);

		// First field in swap region is exchanged with origin
		if (!origin_pulled) {
			origin_pulled = 1;

			struct Field tmp = target->fields[it];
			target->fields[it] = target->fields[origin_idx];
			target->fields[origin_idx] = tmp;

			// Update origin field
			target->fields[it].offset = new_offset;

			// Update swapped field
			target->fields[origin_idx].offset = origin_region_ptr;
			origin_region_ptr += target->fields[origin_idx].size;

			continue;
		}

		// For all other fields in swap region, pull all until origin_idx one forward and place at origin_idx
		struct Field tmp = target->fields[it];

		for (uint32_t pull_it = it + 1; pull_it <= origin_idx; pull_it++)
			target->fields[it - 1] = target->fields[it];

		target->fields[origin_idx] = tmp;
	}
}

static void target_shuffle_one(struct Target* target) {
	if (!target || !target->field_count)
		return;

	uint32_t origin = rand_uint32() % target->field_count;

	if (target->fields[origin].flags & SPSLR_FLAG_FIELD_FIXED)
		return;

	uint32_t origin_alignment = target->fields[origin].alignment;

	uint32_t max_options = target->size / origin_alignment;
	uint32_t* swap_options = (uint32_t*)malloc(sizeof(uint32_t) * max_options); // match fill_begin

	struct ShuffleRegion origin_region;
	target_get_origin_region(target, &origin_region, origin);

	uint32_t option_count = target_get_shuffle_options(target, swap_options, &origin_region, origin_alignment);

	if (!option_count) {
		free(swap_options);
		return;
	}

	uint32_t selected_option_idx = rand_uint32() % option_count;
	uint32_t selected_option = swap_options[selected_option_idx];

	free(swap_options);

	target_swap(target, origin, &origin_region, selected_option);
}

int spslr_randomize(uint32_t target) {
	struct Target* t = find_target(target);
	if (!t)
		return -1;

	if (t->field_count != t->present_field_count)
		return -1;

	seed_rand_time(); // Note -> this is obviously not sufficient

	uint32_t shuffle_count = t->field_count * 2;
	for (uint32_t i = 0; i < shuffle_count; i++)
		target_shuffle_one(t);

	// Compile array of final fields
	t->final_fields = (struct FinalField*)malloc(sizeof(struct FinalField) * t->field_count);
	if (!t->final_fields)
		return -1;

	for (uint32_t i = 0; i < t->field_count; i++) {
		const struct Field* f = &t->fields[i];
		struct FinalField* ff = &t->final_fields[f->initial_idx];

		ff->initial_offset = f->initial_offset;
		ff->offset = f->offset;
	}

	return 0;
}

void spslr_targets_clear() {
	clear_targets();
}