/* radare - LGPL - Copyright 2009-2026 - pancake, dennis */

#include "lua.h"

// luac can not nest functions deeper than its LUAI_MAXCCALLS parser limit
#define LUA_MAX_PROTO_DEPTH 200

void lua_header_free(RLuaHeader *lhead) {
	if (lhead) {
		r_list_free (lhead->functionList);
		RVecRBinSymbol_fini (&lhead->symbols);
		free (lhead);
	}
}

static void lua_func_free(void *f) {
	LuaFunction *func = f;
	free (func->name_ptr);
	free (f);
}

static inline RLuaHeader *lua_header_new(void) {
	RLuaHeader *lh = R_NEW0 (RLuaHeader);
	lh->functionList = r_list_newf ((RListFree)lua_func_free);
	RVecRBinSymbol_init (&lh->symbols);
	return lh;
}

static inline bool has_bytes(const ut64 size, const ut64 offset, const ut64 len) {
	return offset <= size && len <= size - offset;
}

static ut64 parseProtos(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, LuaFunction *func, ParseStruct *parseStruct) {
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	if (lh->depth >= LUA_MAX_PROTO_DEPTH) {
		R_LOG_WARN ("Function prototypes nested deeper than %d levels", LUA_MAX_PROTO_DEPTH);
		return 0;
	}
	ut64 length = parseInt (data + offset);
	offset += lh->intSize;
	R_LOG_DEBUG ("Function has %" PFMT64x " Prototypes", length);

	int i;
	lh->depth++;
	for (i = 0; i < length; i++) {
		offset = lua53parseFunction (lh, data, offset, size, func, parseStruct);
		if (offset == 0) {
			break;
		}
	}
	lh->depth--;
	return offset;
}

static int storeLuaFunction(RLuaHeader *lh, LuaFunction *function) {
	if (!lh->functionList) {
		lh->functionList = r_list_newf ((RListFree)lua_func_free);
		if (!lh->functionList) {
			return 0;
		}
	}
	r_list_append (lh->functionList, function);
	return 1;
}

static LuaFunction *findLuaFunction(RLuaHeader *lh, ut64 addr) {
	if (!lh->functionList) {
		return NULL;
	}
	LuaFunction *function = NULL;
	RListIter *iter = NULL;
	r_list_foreach (lh->functionList, iter, function) {
		R_LOG_DEBUG ("Search 0x%" PFMT64x, function->offset);
		if (function->offset == addr) {
			return function;
		}
	}
	return NULL;
}

static ut64 parseCode(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, ParseStruct *parseStruct) {
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	ut64 length = parseInt (data + offset);
	offset += lh->intSize;

	ut64 code_size = 0;
	if (r_mul_overflow (length, (ut64)lh->instructionSize, &code_size) || !has_bytes (size, offset, code_size)) {
		return 0;
	}
	R_LOG_DEBUG ("Function has %" PFMT64x " Instructions", length);

	return offset + code_size;
}

static ut64 parseStringR(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, char **str_ptr, ut64 *str_len, ParseStruct *parseStruct) {
	if (!has_bytes (size, offset, 1)) {
		R_LOG_DEBUG ("Prevented oobread");
		return 0;
	}
	ut64 functionNameSize = data[offset];
	offset += 1;
	if (functionNameSize == 0xFF) {
		if (!has_bytes (size, offset, lh->sizeSize)) {
			return 0;
		}
		functionNameSize = parseSize (data + offset);
		offset += lh->sizeSize;
	}
	if (functionNameSize != 0) {
		functionNameSize--;
		if (!has_bytes (size, offset, functionNameSize)) {
			return 0;
		}
		if (str_ptr) {
			*str_ptr = r_str_ndup ((char *) (data + offset), functionNameSize);
		}
		if (str_len) {
			*str_len = functionNameSize;
		}
		if (parseStruct && parseStruct->onString) {
			parseStruct->onString (data, offset, functionNameSize, parseStruct);
		}
		R_LOG_DEBUG ("String %.*s", (int)R_MIN (functionNameSize, (ut64)ST32_MAX), data + offset);
		offset += functionNameSize;
	}
	return offset;
}

static ut64 parseString(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, ParseStruct *parseStruct) {
	return parseStringR (lh, data, offset, size, 0, 0, parseStruct);
}

static ut64 parseConstants(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, ParseStruct *parseStruct) {
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	ut64 length = parseInt (data + offset);
	offset += lh->intSize;
	R_LOG_DEBUG ("Function has %" PFMT64x " Constants", length);

	int i;
	for (i = 0; i < length; i++) {
		if (!has_bytes (size, offset, 1)) {
			return 0;
		}
		R_LOG_DEBUG ("%d: ", i);
		ut8 type = data[offset];
		offset += 1;
		switch (type) {
		case 0: // Nil
			R_LOG_DEBUG ("Nil");
			break;
		case 1: // Boolean
			if (!has_bytes (size, offset, 1)) {
				return 0;
			}
			R_LOG_DEBUG ("Boolean %d", data[offset]);
			offset += 1;
			break;
		case (3 | (0 << 4)): // Number
			if (!has_bytes (size, offset, lh->luaNumberSize)) {
				return 0;
			}
		{
#ifdef LUA_DEBUG
			ut64 num = parseLuaNumber (data + offset);
			R_LOG_DEBUG ("Number %f", *((double *)&num));
#endif
			offset += lh->luaNumberSize;
			}
			break;
		case (3 | (1 << 4)): // Integer
			if (!has_bytes (size, offset, lh->luaIntSize)) {
				return 0;
			}
			R_LOG_DEBUG ("Integer %" PFMT64x, parseLuaInt (data + offset));
			offset += lh->luaIntSize;
			break;
		case (4 | (0 << 4)): // Short String
		case (4 | (1 << 4)): // Long String
			offset = parseString (lh, data, offset, size, parseStruct);
			if (!offset) {
				return 0;
			}
			break;
		default:
			R_LOG_DEBUG ("Invalid");
			return 0;
		}
	}
	return offset;
}

static ut64 parseDebug(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, ParseStruct *parseStruct) {
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	ut64 length = parseInt (data + offset);
	offset += lh->intSize;

	if (length != 0) {
		R_LOG_DEBUG ("Instruction-Line Mappings %" PFMT64x, length);
		ut64 lines_size = 0;
		if (r_mul_overflow ((ut64)lh->intSize, length, &lines_size) || !has_bytes (size, offset, lines_size)) {
			return 0;
		}
		int i;
		for (i = 0; i < length; i++) {
			R_LOG_DEBUG ("Instruction %d Line %" PFMT64x, i, parseInt (data + offset));
			offset += lh->intSize;
		}
	}
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	length = parseInt (data + offset);
	offset += lh->intSize;
	if (length != 0) {
		R_LOG_DEBUG ("LiveRanges: %" PFMT64x, length);
		int i;
		for (i = 0; i < length; i++) {
			R_LOG_DEBUG ("LiveRange %d:", i);
			offset = parseString (lh, data, offset, size, parseStruct);
			if (offset == 0) {
				return 0;
			}
			if (!has_bytes (size, offset, lh->intSize * 2)) {
				return 0;
			}
	#ifdef LUA_DEBUG
			ut64 num1 = parseInt (data + offset);
	#endif
			offset += lh->intSize;
	#ifdef LUA_DEBUG
			ut64 num2 = parseInt (data + offset);
	#endif
			offset += lh->intSize;
		}
	}
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	length = parseInt (data + offset);
	offset += lh->intSize;
	if (length != 0) {
		R_LOG_DEBUG ("Up-Values: %" PFMT64x, length);
		int i;
		for (i = 0; i < length; i++) {
			R_LOG_DEBUG ("Up-Value %d:", i);
			offset = parseString (lh, data, offset, size, parseStruct);
			if (offset == 0) {
				return 0;
			}
		}
	}
	return offset;
}

static inline ut64 buf_parse_int(RBuffer *buf, int size, bool le) {
	switch (size) {
	case 2:
		return le? r_buf_read_le16 (buf): r_buf_read_be16 (buf);
	case 4:
		return le? r_buf_read_le32 (buf): r_buf_read_be32 (buf);
	case 8:
		return le? r_buf_read_le64 (buf): r_buf_read_be64 (buf);
	default:
		return UT64_MAX;
	}
}

static ut64 parseUpvalues(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, ParseStruct *parseStruct) {
	if (!has_bytes (size, offset, lh->intSize)) {
		return 0;
	}
	ut64 length = parseInt (data + offset);
	offset += lh->intSize;

	R_LOG_DEBUG ("Function has %" PFMT64x " Upvalues", length);

	if (length > (size - offset) / 2) {
		return 0;
	}
	int i;
	for (i = 0; i < length; i++) {
		R_LOG_DEBUG ("%d: inStack: %d id: %d", i, data[offset + 0], data[offset + 1]);
		offset += 2;
	}
	return offset;
}

static inline double buf_parse_num(RLuaHeader *lh, RBuffer *buf) {
	double ret = 0;
	ut64 num = buf_parse_int (buf, lh->luaNumberSize, lh->isLe);
	memcpy (&ret, &num, R_MIN (64, R_MIN (sizeof (double), lh->luaNumberSize)));
	return ret;
}

static ut64 add_symbol(RLuaHeader *lh, RBuffer *buf, char *name, ut64 start, const char *type) {
	ut64 end = r_buf_tell (buf); // end of field that was just parsed from bf
	if (end > start) {
		RBinSymbol *sym = RVecRBinSymbol_emplace_back (&lh->symbols);
		memset (sym, 0, sizeof (*sym));
		sym->name = r_bin_name_new (name);
		if (sym->name) {
			sym->vaddr = sym->paddr = start;
			sym->attr.size = end - start;
			sym->type = type;
			return end;
		}
		RVecRBinSymbol_pop_back (&lh->symbols);
	}
	// Caller shouldn't stop parsing lua just for a missing symbol. But we return location to save a r_buf_tell
	return end;
}

static inline bool lua53_check_header_data(RBuffer *buf) {
	const ut8 lua_data[] = "\x19\x93\r\n\x1a\n";
	ut8 tmp[6];
	r_buf_read (buf, tmp, sizeof (tmp));
	return memcmp (tmp, lua_data, sizeof (tmp)) == 0;
}

static inline bool is_valid_num_size(int size) {
	if (size == 2 || size == 4 || size == 8) {
		return true;
	}
	return false;
}

bool check_header(RBuffer *b) {
	return r_buf_read_be32 (b) == 0x1b4c7561? true: false; // "\x1bLua"
}

#define GETVALIDSIZE(x, symname) \
	{ \
		lh->x = r_buf_read8 (buf); \
		if (!is_valid_num_size (lh->x)) { \
			R_LOG_WARN ("Invalid size 0x%x for " #x " at offset: 0x%" PFMT64x, lh->x, where); \
			goto bad_header_ret; \
		} \
		where = add_symbol (lh, buf, symname, where, "NUM"); \
	}

// this function expects buf to be pointing to correct location
RLuaHeader *r_lua_load_header(RBuffer *buf) {
	ut64 start = r_buf_tell (buf);
	RLuaHeader *lh = lua_header_new (); // TODO use this when removing global
	if (!lh || !check_header (buf)) {
		R_LOG_DEBUG ("Bad lua magic at offset 0x%" PFMT64x, start);
		goto bad_header_ret;
	}
	ut64 where = add_symbol (lh, buf, "lua-header", start, "NOTYPE");

	// version
	lh->ver = r_buf_read8 (buf);
	if (lh->ver != 0x53) {
#if 0
		int mj = lh->ver >> 4;
		int mn = lh->ver & 0xf;
		R_LOG_DEBUG ("[0x%08" PFMT64x "] Reported lua version  %d.%d (0x%x) not supported", where, mj, mn, lh->ver);
#endif
		goto bad_header_ret; // TODO support more versions
	}
	where = add_symbol (lh, buf, "lua-version", where, "NOTYPE");

	// format
	lh->format = r_buf_read8 (buf);
	if (lh->format != 0) {
		R_LOG_WARN ("[0x%08" PFMT64x "]Unexpected Lua format 0x%x", where, lh->format);
	}
	where = add_symbol (lh, buf, "lua-format", where, "NOTYPE");

	// header data check
	if (!lua53_check_header_data (buf)) {
		R_LOG_DEBUG ("[0x%08" PFMT64x "] Bad Lua Data", where);
		goto bad_header_ret;
	}
	where = add_symbol (lh, buf, "stringterminators", where, "NOTYPE");

	GETVALIDSIZE (intSize, "int-size");
	GETVALIDSIZE (sizeSize, "size-size");
	GETVALIDSIZE (instructionSize, "instruction-size");
	GETVALIDSIZE (luaIntSize, "lua-int-size");
	GETVALIDSIZE (luaNumberSize, "lua-number-size");

	// check-int
	ut64 first_try = buf_parse_int (buf, lh->luaIntSize, lh->isLe);
	if (first_try != 0x5678) {
		lh->isLe = !lh->isLe;
		r_buf_seek (buf, where, R_BUF_SET);
		ut64 second_try = buf_parse_int (buf, lh->luaIntSize, lh->isLe);
		if (second_try != 0x5678) {
			R_LOG_DEBUG ("[0x%08" PFMT64x "] Can't parse lua num of size %u ([0x%" PFMT64x ", 0x%" PFMT64x " != 0x5678])", where, lh->intSize, first_try, second_try);
			goto bad_header_ret;
		}
	}
	where = add_symbol (lh, buf, "check-int", where, "NUM");

	// check numbers
	double num = buf_parse_num (lh, buf);
	if (num != 370.5) {
		R_LOG_DEBUG ("[0x%08" PFMT64x "] Lua test number failed (%lf != 370.5)", where, num);
		goto bad_header_ret;
	}
	where = add_symbol (lh, buf, "check-number", where, "NUM");

	// upvalues
	lh->upValues = r_buf_read8 (buf);
	where = add_symbol (lh, buf, "upvalues", where, "NUM");

	lh->headerSize = where - start;
	return lh;

bad_header_ret:
	lua_header_free (lh);
	return 0;
}

ut64 lua53parseFunction(RLuaHeader *lh, const ut8 *data, ut64 offset, const ut64 size, LuaFunction *parent_func, ParseStruct *parseStruct) {
	R_LOG_DEBUG ("Function 0x%" PFMT64x, offset);
	LuaFunction *function = findLuaFunction (lh, offset);
	if (function) { // if a function object was cached
		R_LOG_DEBUG ("Found cached Functione: 0x%" PFMT64x, function->offset);

		if (parseStruct != NULL && parseStruct->onString != NULL) {
			parseConstants (lh, data, function->const_offset, size, parseStruct);
		}

		parseProtos (lh, data, function->protos_offset, size, function, parseStruct);

		if (parseStruct != NULL && parseStruct->onString != NULL) {
			parseDebug (lh, data, function->debug_offset, size, parseStruct);
		}

		if (parseStruct != NULL && parseStruct->onFunction != NULL) {
			parseStruct->onFunction (lh, function, parseStruct);
		}
		return offset + function->size;
	}
	ut64 baseoffset = offset;
	function = R_NEW0 (LuaFunction);
	function->parent_func = parent_func;
	function->offset = offset;
	offset = parseStringR (lh, data, offset, size, &function->name_ptr, &function->name_size, parseStruct);
	if (!offset || !has_bytes (size, offset, lh->intSize * 2 + 3)) {
		goto beach;
	}

	function->lineDefined = parseInt (data + offset);
	R_LOG_DEBUG ("Line Defined: %" PFMT64x, function->lineDefined);
	function->lastLineDefined = parseInt (data + offset + lh->intSize);
	R_LOG_DEBUG ("Last Line Defined: %" PFMT64x, function->lastLineDefined);
	offset += lh->intSize * 2;
	function->numParams = data[offset + 0];
	R_LOG_DEBUG ("Param Count: %d", function->numParams);
	function->isVarArg = data[offset + 1];
	R_LOG_DEBUG ("Is VarArgs: %d", function->isVarArg);
	function->maxStackSize = data[offset + 2];
	R_LOG_DEBUG ("Max Stack Size: %d", function->maxStackSize);
	offset += 3;

	function->code_offset = offset;
	offset = parseCode (lh, data, offset, size, parseStruct);
	if (offset == 0) {
		goto beach;
	}
	function->code_size = parseInt (data + function->code_offset);
	function->const_offset = offset;
	offset = parseConstants (lh, data, offset, size, parseStruct);
	if (offset == 0) {
		goto beach;
	}
	function->const_size = parseInt (data + function->const_offset);
	function->upvalue_offset = offset;
	offset = parseUpvalues (lh, data, offset, size, parseStruct);
	if (offset == 0) {
		goto beach;
	}
	function->upvalue_size = parseInt (data + function->upvalue_offset);
	function->protos_offset = offset;
	offset = parseProtos (lh, data, offset, size, function, parseStruct);
	if (offset == 0) {
		goto beach;
	}
	function->protos_size = parseInt (data + function->protos_offset);
	function->debug_offset = offset;
	offset = parseDebug (lh, data, offset, size, parseStruct);
	if (offset == 0) {
		goto beach;
	}

	function->size = offset - baseoffset;
	if (parseStruct && parseStruct->onFunction) {
		parseStruct->onFunction (lh, function, parseStruct);
	}
	if (!storeLuaFunction (lh, function)) {
		lua_func_free (function);
	}
	return offset;
beach:
	lua_func_free (function);
	return 0;
}
