1 files changed, 252 insertions, 215 deletions

diff --git a/view/kernelcache/api/kernelcacheapi.h b/view/kernelcache/api/kernelcacheapi.h
index 35c53d78..aceb30a5 100644
--- a/view/kernelcache/api/kernelcacheapi.h
+++ b/view/kernelcache/api/kernelcacheapi.h
@@ -1,274 +1,311 @@
 #pragma once
 
 #include <binaryninjaapi.h>
-#include "../core/MetadataSerializable.hpp"
-#include "../api/view/macho/machoview.h"
 #include "kernelcachecore.h"
 
-using namespace BinaryNinja;
+template<class T>
+class KCRefCountObject {
+	void AddRefInternal() { m_refs.fetch_add(1); }
 
-namespace KernelCacheAPI {
-	template<class T>
-	class KCRefCountObject {
-		void AddRefInternal() { m_refs.fetch_add(1); }
+	void ReleaseInternal() {
+		if (m_refs.fetch_sub(1) == 1)
+			delete this;
+	}
 
-		void ReleaseInternal() {
-			if (m_refs.fetch_sub(1) == 1)
-				delete this;
-		}
+public:
+	std::atomic<int> m_refs;
+	T *m_object;
 
-	public:
-		std::atomic<int> m_refs;
-		T *m_object;
+	KCRefCountObject() : m_refs(0), m_object(nullptr) {}
+
+	virtual ~KCRefCountObject() = default;
+
+	T *GetObject() const { return m_object; }
+
+	static T *GetObject(KCRefCountObject *obj) {
+		if (!obj)
+			return nullptr;
+		return obj->GetObject();
+	}
 
-		KCRefCountObject() : m_refs(0), m_object(nullptr) {}
+	void AddRef() { AddRefInternal(); }
 
-		virtual ~KCRefCountObject() {}
+	void Release() { ReleaseInternal(); }
 
-		T *GetObject() const { return m_object; }
+	void AddRefForRegistration() { AddRefInternal(); }
+};
 
-		static T *GetObject(KCRefCountObject *obj) {
-			if (!obj)
-				return nullptr;
-			return obj->GetObject();
+
+template<class T, T *(*AddObjectReference)(T *), void (*FreeObjectReference)(T *)>
+class KCCoreRefCountObject {
+	void AddRefInternal() { m_refs.fetch_add(1); }
+
+	void ReleaseInternal() {
+		if (m_refs.fetch_sub(1) == 1) {
+			if (!m_registeredRef)
+				delete this;
 		}
+	}
 
-		void AddRef() { AddRefInternal(); }
+public:
+	std::atomic<int> m_refs;
+	bool m_registeredRef = false;
+	T *m_object;
 
-		void Release() { ReleaseInternal(); }
+	KCCoreRefCountObject() : m_refs(0), m_object(nullptr) {}
 
-		void AddRefForRegistration() { AddRefInternal(); }
-	};
+	virtual ~KCCoreRefCountObject() = default;
 
+	T *GetObject() const { return m_object; }
 
-	template<class T, T *(*AddObjectReference)(T *), void (*FreeObjectReference)(T *)>
-	class KCCoreRefCountObject {
-		void AddRefInternal() { m_refs.fetch_add(1); }
+	static T *GetObject(KCCoreRefCountObject *obj) {
+		if (!obj)
+			return nullptr;
+		return obj->GetObject();
+	}
 
-		void ReleaseInternal() {
-			if (m_refs.fetch_sub(1) == 1) {
-				if (!m_registeredRef)
-					delete this;
-			}
-		}
+	void AddRef() {
+		if (m_object && (m_refs != 0))
+			AddObjectReference(m_object);
+		AddRefInternal();
+	}
 
-	public:
-		std::atomic<int> m_refs;
-		bool m_registeredRef = false;
-		T *m_object;
+	void Release() {
+		if (m_object)
+			FreeObjectReference(m_object);
+		ReleaseInternal();
+	}
+
+	void AddRefForRegistration() { m_registeredRef = true; }
 
-		KCCoreRefCountObject() : m_refs(0), m_object(nullptr) {}
+	void ReleaseForRegistration() {
+		m_object = nullptr;
+		m_registeredRef = false;
+		if (m_refs == 0)
+			delete this;
+	}
+};
 
-		virtual ~KCCoreRefCountObject() {}
+template <class T>
+class KCRef
+{
+	T* m_obj;
+#ifdef BN_REF_COUNT_DEBUG
+	void* m_assignmentTrace = nullptr;
+#endif
 
-		T *GetObject() const { return m_object; }
+public:
+	KCRef() : m_obj(NULL) {}
 
-		static T *GetObject(KCCoreRefCountObject *obj) {
-			if (!obj)
-				return nullptr;
-			return obj->GetObject();
+	KCRef(T* obj) : m_obj(obj)
+	{
+		if (m_obj)
+		{
+			m_obj->AddRef();
+#ifdef BN_REF_COUNT_DEBUG
+			m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
 		}
+	}
 
-		void AddRef() {
-			if (m_object && (m_refs != 0))
-				AddObjectReference(m_object);
-			AddRefInternal();
+	KCRef(const KCRef<T>& obj) : m_obj(obj.m_obj)
+	{
+		if (m_obj)
+		{
+			m_obj->AddRef();
+#ifdef BN_REF_COUNT_DEBUG
+			m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
 		}
+	}
 
-		void Release() {
-			if (m_object)
-				FreeObjectReference(m_object);
-			ReleaseInternal();
+	KCRef(KCRef<T>&& other) : m_obj(other.m_obj)
+	{
+		other.m_obj = 0;
+#ifdef BN_REF_COUNT_DEBUG
+		m_assignmentTrace = other.m_assignmentTrace;
+#endif
+	}
+
+	~KCRef()
+	{
+		if (m_obj)
+		{
+			m_obj->Release();
+#ifdef BN_REF_COUNT_DEBUG
+			BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+#endif
 		}
+	}
 
-		void AddRefForRegistration() { m_registeredRef = true; }
+	KCRef<T>& operator=(const BinaryNinja::Ref<T>& obj)
+	{
+#ifdef BN_REF_COUNT_DEBUG
+		if (m_obj)
+			BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+		if (obj.m_obj)
+			m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
+		T* oldObj = m_obj;
+		m_obj = obj.m_obj;
+		if (m_obj)
+			m_obj->AddRef();
+		if (oldObj)
+			oldObj->Release();
+		return *this;
+	}
 
-		void ReleaseForRegistration() {
-			m_object = nullptr;
-			m_registeredRef = false;
-			if (m_refs == 0)
-				delete this;
+	KCRef<T>& operator=(KCRef<T>&& other)
+	{
+		if (m_obj)
+		{
+#ifdef BN_REF_COUNT_DEBUG
+			BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+#endif
+			m_obj->Release();
 		}
-	};
+		m_obj = other.m_obj;
+		other.m_obj = 0;
+#ifdef BN_REF_COUNT_DEBUG
+		m_assignmentTrace = other.m_assignmentTrace;
+#endif
+		return *this;
+	}
 
-	struct KCMemoryRegion {
-		uint64_t vmAddress;
-		uint64_t size;
-		std::string prettyName;
-	};
+	KCRef<T>& operator=(T* obj)
+	{
+#ifdef BN_REF_COUNT_DEBUG
+		if (m_obj)
+			BNUnregisterObjectRefDebugTrace(typeid(T).name(), m_assignmentTrace);
+		if (obj)
+			m_assignmentTrace = BNRegisterObjectRefDebugTrace(typeid(T).name());
+#endif
+		T* oldObj = m_obj;
+		m_obj = obj;
+		if (m_obj)
+			m_obj->AddRef();
+		if (oldObj)
+			oldObj->Release();
+		return *this;
+	}
 
-	struct BackingCacheMapping {
-		uint64_t vmAddress;
-		uint64_t size;
-		uint64_t fileOffset;
-	};
+	operator T*() const
+	{
+		return m_obj;
+	}
 
-	struct BackingCache {
-		std::string path;
-		bool isPrimary;
-		std::vector<BackingCacheMapping> mappings;
-	};
+	T* operator->() const
+	{
+		return m_obj;
+	}
 
-	struct KCImageMemoryMapping {
-		std::string name;
-		uint64_t vmAddress;
-		uint64_t size;
-		bool loaded;
-		uint64_t rawViewOffset;
-	};
+	T& operator*() const
+	{
+		return *m_obj;
+	}
 
-	struct KCImage {
-		std::string name;
-		uint64_t headerFileAddress;
-		std::vector<KCImageMemoryMapping> mappings;
-	};
+	bool operator!() const
+	{
+		return m_obj == NULL;
+	}
 
-	struct KCSymbol {
-		uint64_t address;
-		std::string name;
-		std::string image;
-	};
+	bool operator==(const T* obj) const
+	{
+		return T::GetObject(m_obj) == T::GetObject(obj);
+	}
 
-	using namespace BinaryNinja;
-	struct KernelCacheMachOHeader : public KernelCacheCore::MetadataSerializable<KernelCacheMachOHeader> {
-		uint64_t textBase = 0;
-		uint64_t loadCommandOffset = 0;
-		mach_header_64 ident;
-		std::string identifierPrefix;
-		std::string installName;
+	bool operator==(const KCRef<T>& obj) const
+	{
+		return T::GetObject(m_obj) == T::GetObject(obj.m_obj);
+	}
 
-		std::vector<std::pair<uint64_t, bool>> entryPoints;
-		std::vector<uint64_t> m_entryPoints; //list of entrypoints
+	bool operator!=(const T* obj) const
+	{
+		return T::GetObject(m_obj) != T::GetObject(obj);
+	}
 
-		symtab_command symtab;
-		dysymtab_command dysymtab;
-		dyld_info_command dyldInfo;
-		routines_command_64 routines64;
-		function_starts_command functionStarts;
-		std::vector<section_64> moduleInitSections;
-		linkedit_data_command exportTrie;
-		linkedit_data_command chainedFixups {};
+	bool operator!=(const KCRef<T>& obj) const
+	{
+		return T::GetObject(m_obj) != T::GetObject(obj.m_obj);
+	}
 
-		uint64_t relocationBase;
-		// Section and program headers, internally use 64-bit form as it is a superset of 32-bit
-		std::vector<segment_command_64> segments; //only three types of sections __TEXT, __DATA, __IMPORT
-		segment_command_64 linkeditSegment;
-		std::vector<section_64> sections;
-		std::vector<std::string> sectionNames;
+	bool operator<(const T* obj) const
+	{
+		return T::GetObject(m_obj) < T::GetObject(obj);
+	}
 
-		std::vector<section_64> symbolStubSections;
-		std::vector<section_64> symbolPointerSections;
+	bool operator<(const KCRef<T>& obj) const
+	{
+		return T::GetObject(m_obj) < T::GetObject(obj.m_obj);
+	}
 
-		std::vector<std::string> dylibs;
+	T* GetPtr() const
+	{
+		return m_obj;
+	}
+};
 
-		build_version_command buildVersion;
-		std::vector<build_tool_version> buildToolVersions;
 
-		bool dysymPresent = false;
-		bool dyldInfoPresent = false;
-		bool exportTriePresent = false;
-		bool chainedFixupsPresent = false;
-		bool routinesPresent = false;
-		bool functionStartsPresent = false;
-		bool relocatable = false;
 
-		void Store(KernelCacheCore::SerializationContext& context) const {
-			MSS(textBase);
-			MSS(loadCommandOffset);
-			MSS_SUBCLASS(ident);
-			MSS(identifierPrefix);
-			MSS(installName);
-			MSS(entryPoints);
-			MSS(m_entryPoints);
-			MSS_SUBCLASS(symtab);
-			MSS_SUBCLASS(dysymtab);
-			MSS_SUBCLASS(dyldInfo);
-			MSS_SUBCLASS(routines64);
-			MSS_SUBCLASS(functionStarts);
-			MSS_SUBCLASS(moduleInitSections);
-			MSS_SUBCLASS(exportTrie);
-			MSS_SUBCLASS(chainedFixups);
-			MSS(relocationBase);
-			MSS_SUBCLASS(segments);
-			MSS_SUBCLASS(linkeditSegment);
-			MSS_SUBCLASS(sections);
-			MSS(sectionNames);
-			MSS_SUBCLASS(symbolStubSections);
-			MSS_SUBCLASS(symbolPointerSections);
-			MSS(dylibs);
-			MSS_SUBCLASS(buildVersion);
-			MSS_SUBCLASS(buildToolVersions);
-			MSS(dysymPresent);
-			MSS(dyldInfoPresent);
-			MSS(exportTriePresent);
-			MSS(chainedFixupsPresent);
-			MSS(routinesPresent);
-			MSS(functionStartsPresent);
-			MSS(relocatable);
-		}
+namespace KernelCacheAPI {
+	struct CacheMappingInfo
+	{
+		uint64_t vmAddress;
+		uint64_t size;
+		uint64_t fileOffset;
+	};
 
-		static KernelCacheMachOHeader Load(KernelCacheCore::DeserializationContext& context) {
-			KernelCacheMachOHeader header;
-			header.MSL(textBase);
-			header.MSL(loadCommandOffset);
-			header.MSL(ident);
-			header.MSL(identifierPrefix);
-			header.MSL(installName);
-			header.MSL(entryPoints);
-			header.MSL(m_entryPoints);
-			header.MSL(symtab);
-			header.MSL(dysymtab);
-			header.MSL(dyldInfo);
-			header.MSL(routines64);
-			header.MSL(functionStarts);
-			header.MSL(moduleInitSections);
-			header.MSL(exportTrie);
-			header.MSL(chainedFixups);
-			header.MSL(relocationBase);
-			header.MSL(segments);
-			header.MSL(linkeditSegment);
-			header.MSL(sections);
-			header.MSL(sectionNames);
-			header.MSL(symbolStubSections);
-			header.MSL(symbolPointerSections);
-			header.MSL(dylibs);
-			header.MSL(buildVersion);
-			header.MSL(buildToolVersions);
-			header.MSL(dysymPresent);
-			header.MSL(dyldInfoPresent);
-			header.MSL(exportTriePresent);
-			header.MSL(chainedFixupsPresent);
-			header.MSL(routinesPresent);
-			header.MSL(functionStartsPresent);
-			header.MSL(relocatable);
-			return header;
-		}
+	struct CacheImage
+	{
+		uint64_t headerFileAddress;
+		uint64_t headerVirtualAddress;
+		std::string name;
 	};
 
+	struct CacheEntry
+	{
+		std::string path;
+		std::string name;
+		BNKernelCacheEntryType entryType;
+		std::vector<CacheMappingInfo> mappings;
+	};
 
-	class KernelCache : public KCCoreRefCountObject<BNKernelCache, BNNewKernelCacheReference, BNFreeKernelCacheReference> {
-	public:
-		KernelCache(Ref<BinaryView> view);
+	struct CacheSymbol
+	{
+		BNSymbolType type;
+		uint64_t address;
+		std::string name;
+
+		std::pair<std::string, BinaryNinja::Ref<BinaryNinja::Type>> DemangledName(BinaryNinja::BinaryView &view) const;
+		BinaryNinja::Ref<BinaryNinja::Symbol> GetBNSymbol(BinaryNinja::BinaryView& view) const;
+	};
+
+	std::string GetSymbolTypeAsString(const BNSymbolType& type);
 
-		BNKCViewState GetState();
-		static BNKCViewLoadProgress GetLoadProgress(Ref<BinaryView> view);
-		static uint64_t FastGetImageCount(Ref<BinaryView> view);
+	class KernelCacheController : public KCCoreRefCountObject<BNKernelCacheController, BNNewKernelCacheControllerReference, BNFreeKernelCacheControllerReference> {
+	public:
+		explicit KernelCacheController(BNKernelCacheController* controller);
+		static KCRef<KernelCacheController> GetController(BinaryNinja::BinaryView& view);
 
-		bool LoadImageWithInstallName(const std::string& installName);
-		bool LoadImageContainingAddress(uint64_t addr);
-		std::vector<std::string> GetAvailableImages();
+		// Attempt to load the given image into the view.
+		//
+		// It is the callers responsibility to run linear sweep and update analysis, as you might want to add
+		// multiple images at a time.
+		bool ApplyImage(BinaryNinja::BinaryView& view, const CacheImage& image);
 
-		bool IsImageLoaded(const uint64_t address) const;
+		bool IsImageLoaded(const CacheImage& image) const;
 
-		std::vector<KCSymbol> LoadAllSymbolsAndWait();
+		std::optional<CacheImage> GetImageAt(uint64_t address) const;
+		std::optional<CacheImage> GetImageContaining(uint64_t address) const;
+		std::optional<CacheImage> GetImageWithName(const std::string& name) const;
 
-		std::string GetNameForAddress(uint64_t address);
-		std::string GetImageNameForAddress(uint64_t address);
+		std::vector<std::string> GetImageDependencies(const CacheImage& image) const;
 
-		std::vector<KCImage> GetImages();
-		std::vector<KCImage> GetLoadedImages();
+		std::optional<CacheSymbol> GetSymbolAt(uint64_t address) const;
+		std::optional<CacheSymbol> GetSymbolWithName(const std::string& name) const;
 
-		std::optional<KernelCacheMachOHeader> GetMachOHeaderForImage(const std::string& name);
-		std::optional<KernelCacheMachOHeader> GetMachOHeaderForAddress(uint64_t address);
+		std::vector<CacheImage> GetImages() const;
+		std::vector<CacheImage> GetLoadedImages() const;
+		std::vector<CacheSymbol> GetSymbols() const;
 	};
-}
\ No newline at end of file
+}