/* SGX530 1.2.1 Android binary patcher
 * Daniel Beer <dlbeer@gmail.com>
 * 26 Jun 2015
 *
 * This file is in the public domain.
 */

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

/************************************************************************
 * SHA256 implementation.
 */

#define SHA256_BLOCK_SIZE		64
#define SHA256_HASH_SIZE		32

static const uint32_t sha256_round_k[64] = {
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
    0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
    0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,
    0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
    0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
    0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,
    0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,
    0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
    0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
};

static inline uint32_t rot32(uint32_t x, int bits)
{
    return (x >> bits) | (x << (32 - bits));
}

static inline uint32_t rnet32(const uint8_t *x)
{
	uint32_t r = 0;

	r = (r << 8) | x[0];
	r = (r << 8) | x[1];
	r = (r << 8) | x[2];
	r = (r << 8) | x[3];

	return r;
}

static inline void wnet32(uint8_t *x, uint32_t r)
{
	x[3] = r;
	r >>= 8;
	x[2] = r;
	r >>= 8;
	x[1] = r;
	r >>= 8;
	x[0] = r;
}

const uint8_t sha256_init_state[SHA256_HASH_SIZE] = {
    0x6a, 0x09, 0xe6, 0x67, 0xbb, 0x67, 0xae, 0x85,
    0x3c, 0x6e, 0xf3, 0x72, 0xa5, 0x4f, 0xf5, 0x3a,
    0x51, 0x0e, 0x52, 0x7f, 0x9b, 0x05, 0x68, 0x8c,
    0x1f, 0x83, 0xd9, 0xab, 0x5b, 0xe0, 0xcd, 0x19,
};

static void sha256_transform(uint8_t *state, uint32_t *w)
{
	uint32_t x[8];
	unsigned int i;

	// Load state
	for (i = 0; i < 8; i++)
		x[i] = rnet32(state + i * 4);

	for (i = 0; i < 64; i++) {
		// Compute value of w[i + 16]. w[wrap(i)] is currently w[i]
		const uint32_t wi = w[i & 15];
		const uint32_t wi15 = w[(i + 1) & 15];
		const uint32_t wi2 = w[(i + 14) & 15];
		const uint32_t wi7 = w[(i + 9) & 15];
		const uint32_t s0 =
		    rot32(wi15, 7) ^ rot32(wi15, 18) ^ (wi15 >> 3);
		const uint32_t s1 =
		    rot32(wi2, 17) ^ rot32(wi2, 19) ^ (wi2 >> 10);

		// Round calculations
		const uint32_t S0 =
		    rot32(x[0], 2) ^ rot32(x[0], 13) ^ rot32(x[0], 22);
		const uint32_t S1 =
		    rot32(x[4], 6) ^ rot32(x[4], 11) ^ rot32(x[4], 25);
		const uint32_t ch = (x[4] & x[5]) ^ ((~x[4]) & x[6]);
		const uint32_t temp1 = x[7] + S1 + ch + sha256_round_k[i] + wi;
		const uint32_t maj =
		    (x[0] & x[1]) ^ (x[0] & x[2]) ^ (x[1] & x[2]);
		const uint32_t temp2 = S0 + maj;

		// Update round state
		x[7] = x[6];
		x[6] = x[5];
		x[5] = x[4];
		x[4] = x[3] + temp1;
		x[3] = x[2];
		x[2] = x[1];
		x[1] = x[0];
		x[0] = temp1 + temp2;

		// w[wrap(i)] becomes w[i + 16]
		w[i & 15] = wi + s0 + wi7 + s1;
	}

	for (i = 0; i < 8; i++)
		wnet32(state + i * 4, rnet32(state + i * 4) + x[i]);
}

static void sha256_init(uint8_t *state)
{
	memcpy(state, sha256_init_state, SHA256_HASH_SIZE);
}

static void sha256_update(uint8_t *state, const uint8_t *block)
{
	uint32_t w[16];
	unsigned int i;

	for (i = 0; i < 16; i++)
		w[i] = rnet32(block + i * 4);

	sha256_transform(state, w);
}

static void sha256_final(uint8_t *state,
			 const uint8_t *part_block, size_t total_len)
{
	const unsigned int last_len = total_len % SHA256_BLOCK_SIZE;
	const unsigned int last_words = last_len / 4;
	const unsigned int last_bytes = last_len % 4;
	uint32_t w[16];
	uint32_t r = 0;
	unsigned int i;

	memset(w, 0, sizeof(w));

	// Copy in whole words
	for (i = 0; i < last_words; i++)
		w[i] = rnet32(part_block + i * 4);

	// Copy the 0 <= x < 8 remaining words, plus a terminating 0x80,
	// into the last word.
	for (i = 0; i < last_bytes; i++)
		r = (r << 8) | part_block[last_words * 4 + i];
	r = (r << 8) | 0x80;
	for (i = last_bytes + 1; i < 4; i++)
		r <<= 8;

	w[last_words] = r;

	// If we don't have enough room for the size count, flip the block
	if (last_words >= 14) {
		sha256_transform(state, w);
		memset(w, 0, sizeof(w));
	}

	// Flip the last block with the size count
	w[15] = total_len * 8;
	sha256_transform(state, w);
}

static void sha256_blob(uint8_t *state, const uint8_t *data, size_t len)
{
	const size_t total = len;

	sha256_init(state);

	while (len >= SHA256_BLOCK_SIZE) {
		sha256_update(state, data);
		data += SHA256_BLOCK_SIZE;
		len -= SHA256_BLOCK_SIZE;
	}

	sha256_final(state, data, total);
}

/************************************************************************
 * Easy file manipulation ops
 */

static void perror_die(const char *prefix)
{
	perror(prefix);
	abort();
}

struct img {
	uint8_t		*data;
	size_t		len;
};

static void file_read(struct img *img, const char *filename)
{
	FILE *f = fopen(filename, "rb");

	if (!f)
		perror_die("fopen (read)");

	fseek(f, 0, SEEK_END);
	img->len = ftell(f);
	rewind(f);

	if (!img->len) {
		fclose(f);
		return;
	}

	img->data = malloc(img->len);
	if (!img->data)
		perror_die("malloc");

	if (fread(img->data, img->len, 1, f) != 1)
		perror_die("fread");

	fclose(f);
}

static void file_write(struct img *img, const char *filename)
{
	FILE *f = fopen(filename, "wb");

	if (!f)
		perror_die("fopen (write)");

	if (fwrite(img->data, img->len, 1, f) != 1)
		perror_die("fwrite");

	fclose(f);
}

/************************************************************************
 * Patcher
 */

/* 0xc160 is the start of glDiscardFramebufferEXT(). It's buggy and
 * crashes. We patch it to turn it into a no-op.
 */
static const size_t func_offset = 0xc160;
static const uint32_t ins_old = 0xe92d4ff7; // push {r0-r2, r4-r9, sl, fp, lr}
static const uint32_t ins_new = 0xe12fff1e; // bx lr

static inline uint32_t rle32(const uint8_t *x)
{
	uint32_t r = 0;

	r = (r << 8) | x[3];
	r = (r << 8) | x[2];
	r = (r << 8) | x[1];
	r = (r << 8) | x[0];

	return r;
}

static inline void wle32(uint8_t *x, uint32_t r)
{
	x[0] = r;
	r >>= 8;
	x[1] = r;
	r >>= 8;
	x[2] = r;
	r >>= 8;
	x[3] = r;
}

static uint32_t check_version(struct img *img)
{
	static const uint8_t hash_expect[SHA256_HASH_SIZE] = {
		0x6c, 0x15, 0xda, 0x9e, 0x84, 0xf8, 0x40, 0x0e,
		0xa0, 0xed, 0xfd, 0x19, 0x00, 0x4f, 0x08, 0x6c,
		0x0f, 0x2f, 0x72, 0xa9, 0xa4, 0x29, 0xe6, 0xf2,
		0x87, 0x89, 0x1d, 0x73, 0xa7, 0xbd, 0x0a, 0x96,
	};

	uint8_t hash[SHA256_HASH_SIZE];
	uint32_t save;
	unsigned int i;

	if (img->len != 2092185) {
		fprintf(stderr, "Incorrect file size\n");
		abort();
	}

	/* Hash the file as it would be if we'd never modified it */
	save = rle32(img->data + func_offset);
	wle32(img->data + func_offset, ins_old);
	sha256_blob(hash, img->data, img->len);
	wle32(img->data + func_offset, save);

	printf("Hash: ");
	for (i = 0; i < SHA256_HASH_SIZE; i++)
		printf("%02x", hash[i]);
	printf("\n");

	if (memcmp(hash, hash_expect, SHA256_HASH_SIZE)) {
		fprintf(stderr, "Hash does not match!\n");
		abort();
	}

	printf("Instruction @ 0x%x: 0x%08x\n",
	       (unsigned int)func_offset, save);
	if (save == ins_old)
		printf("  -- unpatched\n");
	else if (save == ins_new)
		printf("  -- patched\n");
	else
		printf("  -- unknown\n");

	return save;
}

int main(int argc, char **argv)
{
	struct img img = {0};
	uint32_t cur;
	int want_reverse = 0;

	if (argc < 2) {
		printf("Android SGX: "
		       "libGLESv2_POWERVR_SGX530_121.so patcher\n");
		printf("Daniel Beer <dlbeer@gmail.com>, "
		       "26 Jun 2015\n");
		printf("Usage: %s <file> [reverse]\n", argv[0]);
		return -1;
	}

	file_read(&img, argv[1]);
	printf("Read %d bytes\n", (int)img.len);

	cur = check_version(&img);
	want_reverse = (argc >= 3) && !strcmp(argv[2], "reverse");
	wle32(img.data + func_offset, want_reverse ? ins_old : ins_new);

	if (cur != rle32(img.data + func_offset)) {
		printf("Patching...\n");
		file_write(&img, argv[1]);
	}

	free(img.data);
	return 0;
}