/*
 * dtb_grow.c - add properties to a node in a devicetree that has
 * room, the relocated copy from dtb_reloc.c. the insert point is
 * the node's FDT_END_NODE token, everything after it moves up by
 * the inserted size, the header totalsize tracks it.
 *
 * the insert is safe when the node sits at the end of the struct
 * block, which is the common shape, /chosen is created last by
 * firmware and the tail behind it is two end tokens and the
 * block end. the strings block sits after the grow room and
 * never moves.
 *
 * Copyright (C) 2026 Bradley Morgan <brads@mainlining.org>
 */

#include <string.h>
#include <endian.h>
#include <boot.h>
#include <dtb_patch.h>

#define FDT_BEGIN_NODE	1
#define FDT_END_NODE	2
#define FDT_PROP	3
#define FDT_NOP		4
#define FDT_END		9

static uint32_t be32(const void *p)
{
	const uint8_t *b = p;

	return ((uint32_t)b[0] << 24) | ((uint32_t)b[1] << 16) |
	       ((uint32_t)b[2] << 8) | (uint32_t)b[3];
}

static void put_be32(void *p, uint32_t v)
{
	uint8_t *b = p;

	b[0] = (uint8_t)(v >> 24);
	b[1] = (uint8_t)(v >> 16);
	b[2] = (uint8_t)(v >> 8);
	b[3] = (uint8_t)v;
}

static int name_eq(const char *a, const char *b)
{
	while (*a && *a != '@') {
		if (*a != *b)
			return 0;
		a++;
		b++;
	}
	return *b == '\0' || *b == '@';
}

/*
 * insert one property into /chosen before its end token. value is
 * copied as raw cells, len the byte count. name lands in the free
 * space after the strings block. returns 0 or -1.
 */
int tb_dtb_add_chosen_prop(uintptr_t dtb, const char *name,
			   const void *val, size_t len)
{
	uint8_t *basep = (uint8_t *)dtb;
	uint32_t off_struct = be32(basep + 8);
	uint32_t off_strings = be32(basep + 12);
	uint32_t totalsize = be32(basep + 4);
	uint8_t *p = basep + off_struct;
	uint8_t *ins;
	size_t name_len = strlen(name) + 1;
	size_t prop_size;
	int depth = 0;
	int in_chosen = 0;

	if (be32(basep) != 0xd00dfeed)
		return -1;

	/* find the chosen node's end token, one level under the root */
	while (p < basep + totalsize) {
		uint32_t token = be32(p);

		if (token == FDT_BEGIN_NODE) {
			char *n = (char *)(p + 4);
			size_t nlen = strlen(n) + 1;

			depth++;
			if (depth == 2 && name_eq(n, "chosen"))
				in_chosen = 1;
			p += 4 + ((nlen + 3) & ~3);
		} else if (token == FDT_END_NODE) {
			if (in_chosen && depth == 2) {
				ins = p;
				break;
			}
			depth--;
			p += 4;
		} else if (token == FDT_PROP) {
			uint32_t plen = be32(p + 4);

			p += 12 + ((plen + 3) & ~3);
		} else if (token == FDT_NOP) {
			p += 4;
		} else if (token == FDT_END) {
			break;
		} else {
			return -1;
		}
	}

	if (!ins)
		return -2;

	ins = p;

	/*
	 * the insert: the strings block moves up by prop_size so the
	 * struct block can grow into its old place, the struct tail
	 * after chosen moves up by prop_size, the new name lands at
	 * the end of the moved strings block, and totalsize covers
	 * both. prop name offsets are strings relative so they keep
	 * resolving after the move.
	 */
	{
		prop_size = 12 + ((len + 3) & ~3);
		size_t strings_len = (size_t)be32(basep + 32);

		/* strings block up by prop_size */
		for (size_t i = strings_len; i > 0; i--)
			basep[off_strings + prop_size + i - 1] =
				basep[off_strings + i - 1];

		/* struct tail after the insert point up by prop_size */
		{
			size_t tail = (size_t)(basep + off_strings - ins);

			for (size_t i = tail; i > 0; i--)
				ins[i + prop_size - 1] = ins[i - 1];
		}

		/* the prop token, name offset = old strings length */
		put_be32(ins, FDT_PROP);
		put_be32(ins + 4, (uint32_t)len);
		put_be32(ins + 8, (uint32_t)strings_len);
		for (size_t i = 0; i < len; i++)
			ins[12 + i] = ((const uint8_t *)val)[i];
		for (size_t i = len; i < ((len + 3) & ~3); i++)
			ins[12 + i] = 0;

		/* the name at the end of the moved strings block */
		for (size_t i = 0; i < name_len; i++)
			basep[off_strings + prop_size + strings_len + i] =
				name[i];

		/*
		 * size_dt_struct bounds the token walk, libfdt
		 * rejects anything past it as BADSTRUCTURE. it grows
		 * by the prop size here, the strings size by the name
		 * length, totalsize by both.
		 */
		put_be32(basep + 4, totalsize + (uint32_t)prop_size +
			 (uint32_t)name_len);
		put_be32(basep + 12, off_strings + (uint32_t)prop_size);
		put_be32(basep + 36, be32(basep + 36) + (uint32_t)prop_size);
		/*
		 * size_dt_strings must grow too, libfdt validates
		 * name offsets against it and rejects the whole tree
		 * when the new names sit past the declared end. the
		 * kernel's early parser is the same libfdt, a stale
		 * field there means no memory node and a page table
		 * panic before the first print.
		 */
		put_be32(basep + 32, (uint32_t)strings_len +
			 (uint32_t)name_len);
	}

	return 0;
}