dtb-parser.cpp

#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <span>
#include <cstdint>
#include <type_traits>
#include <cassert>
#include <cstring>

namespace arch {

template<typename...>
constexpr bool dependent_false = false;

enum class endian {
    little = __ORDER_LITTLE_ENDIAN__,
    big    = __ORDER_BIG_ENDIAN__,
    native = __BYTE_ORDER__
};

inline constexpr endian little_endian = endian::little;
inline constexpr endian big_endian = endian::big;
inline constexpr endian native_endian = endian::native;

static_assert(endian::native == endian::little || endian::native == endian::big,
		"only little and big endian are supported");

template<typename T>
inline T bswap(T val) {
	static_assert(std::is_integral_v<T>, "T must be an integral type");
	if constexpr (sizeof(T) == 1) {
		return val;
	} else if constexpr (sizeof(T) == 2) {
		return __builtin_bswap16(val);
	} else if constexpr (sizeof(T) == 4) {
		return __builtin_bswap32(val);
	} else if constexpr (sizeof(T) == 8) {
		return __builtin_bswap64(val);
	} else {
		static_assert(dependent_false<T>, "unsupported swap size");
	}
}

template<endian NewEndian, endian OldEndian = endian::native, typename T>
inline T convert_endian(T native) {
	static_assert(std::is_integral_v<T>, "T must be an integral type");
	if constexpr (NewEndian != OldEndian) {
		return bswap(native);
	} else {
		return native;
	}
}

template<typename R, typename B, endian E>
struct basic_storage {
	using rep_type = R;
	using bits_type = B;

	basic_storage() = default;

	constexpr basic_storage(R r)
	: _embedded{convert_endian<E, endian::native, B>(static_cast<B>(r))} { }

	R load() {
		return static_cast<R>(convert_endian<endian::native, E, B>(_embedded));
	}

	void store(R r) {
		_embedded = convert_endian<E, endian::native, B>(static_cast<B>(r));
	}

private:
	B _embedded;
};

template<typename T, endian E>
using scalar_storage = basic_storage<T, T, E>;

} // namespace arch

struct DtbHeader {
	arch::scalar_storage<uint32_t, arch::big_endian> magic;
	arch::scalar_storage<uint32_t, arch::big_endian> totalsize;
	arch::scalar_storage<uint32_t, arch::big_endian> off_dt_struct;
	arch::scalar_storage<uint32_t, arch::big_endian> off_dt_strings;
	arch::scalar_storage<uint32_t, arch::big_endian> off_mem_rsvmap;
	arch::scalar_storage<uint32_t, arch::big_endian> version;
	arch::scalar_storage<uint32_t, arch::big_endian> last_comp_version;
	arch::scalar_storage<uint32_t, arch::big_endian> boot_cpuid_phys;
	arch::scalar_storage<uint32_t, arch::big_endian> size_dt_strings;
	arch::scalar_storage<uint32_t, arch::big_endian> size_dt_struct;
};

struct DtbReserveEntry {
	arch::scalar_storage<uint64_t, arch::big_endian> address;
	arch::scalar_storage<uint64_t, arch::big_endian> size;
};

enum class Tag : uint32_t {
	beginNode = 1,
	endNode,
	prop,
	nop,
	end = 9
};

struct DeviceTreeNode;

template <typename T>
concept DeviceTreeWalker = requires (T t, DeviceTreeNode n) {
	{ t.push(n) } -> std::same_as<void>;
	{ t.pop() } -> std::same_as<void>;
};

struct DeviceTree {
	DeviceTree(uint8_t *data)
	: data_{data} {
		DtbHeader header;
		memcpy(&header, data, sizeof(header));
		assert(header.magic.load() == 0xd00dfeed);
		assert(header.last_comp_version.load() == 16);

		stringsBlock_ = data + header.off_dt_strings.load();
		structureBlock_ = data + header.off_dt_struct.load();

		totalSize_ = header.totalsize.load();
	}

	size_t size() const {
		return totalSize_;
	}

	DeviceTreeNode rootNode();

	const uint8_t *stringsBlock() const {
		return stringsBlock_;
	}

private:
	uint8_t *data_;

	uint8_t *stringsBlock_;
	uint8_t *structureBlock_;

	uint32_t totalSize_;
};

struct DeviceTreeProperty {
	const char *name;
	std::span<uint8_t> data;
};

namespace detail {
	inline Tag readTag(uint8_t *&ptr) {
		Tag t;
		do {
			arch::scalar_storage<uint32_t, arch::big_endian> tag;
			memcpy(&tag, ptr, 4);
			ptr += 4;
			t = static_cast<Tag>(tag.load());
		} while (t == Tag::nop);

		assert(t != Tag::end);

		return t;
	}

	inline const char *readStringInline(uint8_t *&ptr) {
		auto str = reinterpret_cast<char *>(ptr);
		ptr += (strlen(str) + 4) & ~3;

		return str;
	}

	inline const char *readString(DeviceTree *tree, uint8_t *&ptr) {
		arch::scalar_storage<uint32_t, arch::big_endian> strOff;
		memcpy(&strOff, ptr, 4);
		ptr += 4;

		return reinterpret_cast<const char *>(tree->stringsBlock()) + strOff.load();
	}

	inline uint32_t readLength(uint8_t *&ptr) {
		arch::scalar_storage<uint32_t, arch::big_endian> len;
		memcpy(&len, ptr, 4);
		ptr += 4;

		return len.load();
	}

	inline std::span<uint8_t> readPropData(uint8_t *&ptr, uint32_t len) {
		auto dataPtr = ptr;
		ptr += (len + 3) & ~3;

		return {dataPtr, len};
	}

	inline void skipProp(uint8_t *&ptr) {
		auto len = readLength(ptr);
		ptr += 4; // skip name
		ptr += (len + 3) & ~3; // skip data
	}
} // namespace detail

struct DeviceTreeNode {
	DeviceTreeNode(DeviceTree *tree, uint8_t *base)
	: tree_{tree}, base_{base}, nodeOff_{}, propOff_{}, name_{},
			properties_{nullptr, nullptr, nullptr} {
		uint8_t *tmp = base;

		auto tag = detail::readTag(tmp);
		assert(tag == Tag::beginNode);
		name_ = detail::readStringInline(tmp);

		propOff_ = tmp;
		nodeOff_ = findNodeOff_();

		properties_ = {tree, propOff_, nodeOff_};
	}

	bool operator==(const DeviceTreeNode &other) const {
		return tree_ == other.tree_ && base_ == other.base_;
	}

	template <DeviceTreeWalker T>
	void walkChildren(T &&walker) {
		uint8_t *ptr = nodeOff_;
		int depth = 0;

		while (true) {
			auto tag = detail::readTag(ptr);

			switch (tag) {
				case Tag::beginNode:
					// construct a node and push it
					depth++;
					walker.push({tree_, ptr - 4});
					(void)detail::readStringInline(ptr);
					break;

				case Tag::prop:
					// skip properties of subnodes
					// TODO: ensure nodes and properties are not interleaved?
					detail::skipProp(ptr);
					break;

				case Tag::endNode:
					// pop a node
					walker.pop();
					if (!depth--)
						return;
					break;

				default:
					std::cout << "Unknown tag " << (uint32_t)tag << " found" << std::endl;
			}
		}
	}

	const char *name() const {
		return name_;
	}

	DeviceTree *tree() const {
		return tree_;
	}

	struct PropertyRange {
		PropertyRange(DeviceTree *tree, uint8_t *begin, uint8_t *end)
		: begin_{tree, begin}, end_{tree, end} { }

		struct Iterator {
			Iterator(DeviceTree *tree, uint8_t *ptr)
			: tree_{tree}, ptr_{ptr} { }

			bool operator==(const Iterator &other) const = default;

			DeviceTreeProperty operator*() {
				auto p = ptr_;

				auto tag = detail::readTag(p);
				assert(tag == Tag::prop);

				auto len = detail::readLength(p);
				auto name = detail::readString(tree_, p);
				auto data = detail::readPropData(p, len);
				return DeviceTreeProperty{name, data};
			}

			Iterator &operator++() {
				next();
				return *this;
			}

			Iterator &operator++(int) {
				next();
				return *this;
			}
		private:
			void next() {
				ptr_ += 4; // skip tag
				detail::skipProp(ptr_);
				detail::readTag(ptr_); // skip potential nop tags
				ptr_ -= 4; // rewind back to tag
			}

			DeviceTree *tree_;
			uint8_t *ptr_;
		};

		Iterator begin() const {
			return begin_;
		}

		Iterator end() const {
			return end_;
		}

	private:
		Iterator begin_;
		Iterator end_;
	};

	PropertyRange properties() const {
		return properties_;
	}

private:
	uint8_t *findNodeOff_() {
		uint8_t *ptr = propOff_;

		Tag tag;
		while (true) {
			tag = detail::readTag(ptr);

			if (tag != Tag::prop)
				break;

			detail::skipProp(ptr);
		}

		return ptr - 4; // pointer to tag
	}

	DeviceTree *tree_;
	uint8_t *base_;

	uint8_t *nodeOff_;
	uint8_t *propOff_;
	const char *name_;

	PropertyRange properties_;
};

inline DeviceTreeNode DeviceTree::rootNode() {
	return {this, structureBlock_};
}

template <DeviceTreeWalker T>
void walkTree(DeviceTree &dt, T &&walker) {
	walker.push(dt.rootNode());
	dt.rootNode().walkChildren(static_cast<T &>(walker));
	walker.pop();
}

int main(int argc, char **argv) {
	if (argc < 2) {
		std::cerr << "Missing arg: filename" << std::endl;
		return 1;
	}

	std::ifstream f{argv[1], std::fstream::binary};

	if (!f) {
		std::cerr << "Failed to open file" << std::endl;
		return 1;
	}

	std::vector<uint8_t> data;
	std::copy(std::istreambuf_iterator<char>{f}, std::istreambuf_iterator<char>{}, std::back_inserter(data));

	DeviceTree dt{data.data()};

	struct {
		void push(DeviceTreeNode node) {
			for (int i = 0; i < depth; i++)
				std::cout << "\t";

			std::cout << " - Node \"" << node.name() << "\"" << std::endl;

			for (auto prop : node.properties()) {
				for (int i = 0; i < depth + 1; i++)
					std::cout << "\t";

				std::cout << " - Property \"" << prop.name << "\": " << prop.data.data() << std::endl;
			}

			depth++;
		}

		void pop() {
			depth--;
		}

		int depth = 0;
	} walker;

	walkTree(dt, walker);
}