#pragma once

#include <binaryninjaapi.h>

namespace BinaryNinja {

// This set of structs is based on the objc4 source,

// however pointers have been replaced with view_ptr_t

// Used for pointers within BinaryView, primarily to make it far more clear in typedefs

// whether the size of a field can vary between architectures.

// These should _not_ be used in sizeof or direct Read() calls.

typedef uint64_t view_ptr_t;

typedef struct {

view_ptr_t name;

view_ptr_t types;

view_ptr_t imp;

} method_t;

typedef struct {

uint32_t name;

uint32_t types;

uint32_t imp;

} method_entry_t;

typedef struct {

view_ptr_t offset;

view_ptr_t name;

view_ptr_t type;

uint32_t alignmentRaw;

uint32_t size;

} ivar_t;

typedef struct {

view_ptr_t name;

view_ptr_t attributes;

} property_t;

typedef struct {

uint32_t entsizeAndFlags;

uint32_t count;

} method_list_t;

typedef struct {

uint32_t entsizeAndFlags;

uint32_t count;

} ivar_list_t;

typedef struct {

uint32_t entsizeAndFlags;

uint32_t count;

} property_list_t;

typedef struct {

uint64_t count;

} protocol_list_t;

struct relative_list_list_entry_t {

uint64_t imageIndex: 16;

int64_t listOffset: 48;

};

typedef struct {

view_ptr_t isa;

view_ptr_t mangledName;

view_ptr_t protocols;

view_ptr_t instanceMethods;

view_ptr_t classMethods;

view_ptr_t optionalInstanceMethods;

view_ptr_t optionalClassMethods;

view_ptr_t instanceProperties;

uint32_t size;

uint32_t flags;

} protocol_t;

typedef struct {

uint32_t flags;

uint32_t instanceStart;

uint32_t instanceSize;

uint32_t reserved;

view_ptr_t ivarLayout;

view_ptr_t name;

view_ptr_t baseMethods;

view_ptr_t baseProtocols;

view_ptr_t ivars;

view_ptr_t weakIvarLayout;

view_ptr_t baseProperties;

} class_ro_t;

typedef struct {

view_ptr_t isa;

view_ptr_t super;

view_ptr_t cache;

view_ptr_t vtable;

view_ptr_t data;

} class_t;

typedef struct {

view_ptr_t name;

view_ptr_t cls;

view_ptr_t instanceMethods;

view_ptr_t classMethods;

view_ptr_t protocols;

view_ptr_t instanceProperties;

} category_t;

typedef struct {

view_ptr_t receiver;

view_ptr_t current_class;

} objc_super2;

typedef struct {

view_ptr_t imp;

view_ptr_t sel;

} message_ref_t;

struct Method {

std::string name;

std::string types;

view_ptr_t imp;

};

struct Ivar {

uint32_t offset;

std::string name;

std::string type;

uint32_t alignment;

uint32_t size;

};

struct Property {

std::string name;

std::string attributes;

};

struct ClassBase {

std::map<uint64_t, Method> methodList;

std::map<uint64_t, Ivar> ivarList;

};

struct Class {

std::string name;

ClassBase instanceClass;

ClassBase metaClass;

// Loaded by type processing

QualifiedName associatedName;

};

class Protocol {

public:

std::string name;

std::vector<QualifiedName> protocols;

ClassBase instanceMethods;

ClassBase classMethods;

ClassBase optionalInstanceMethods;

ClassBase optionalClassMethods;

};

struct QualifiedNameOrType {

BinaryNinja::Ref<BinaryNinja::Type> type = nullptr;

BinaryNinja::QualifiedName name;

size_t ptrCount = 0;

};

class ObjCReader {

public:

virtual ~ObjCReader() = default;

/*! Read from the current cursor position into buffer `dest` and advance the cursor that many bytes

\throws Exception

\param dest Address to write the read bytes to

\param len Number of bytes to write

*/

virtual void Read(void* dest, size_t len) = 0;

/*! Read a null-terminated string from the current cursor position

\throws Exception

\return the string

*/

virtual std::string ReadCString(size_t maxLength = -1) = 0;

/*! Read a uint8_t from the current cursor position and advance the cursor by 1 byte

\throws Exception

\return The read value

*/

virtual uint8_t Read8() = 0;

/*! Read a uint16_t from the current cursor position and advance the cursor by 2 bytes

\throws Exception

\return The read value

*/

virtual uint16_t Read16() = 0;

/*! Read a uint32_t from the current cursor position and advance the cursor by 4 bytes

\throws Exception

\return The read value

*/

virtual uint32_t Read32() = 0;

/*! Read a uint64_t from the current cursor position and advance the cursor by 8 bytes

\throws Exception

\return The read value

*/

virtual uint64_t Read64() = 0;

/*! Read a int8_t from the current cursor position and advance the cursor by 1 byte

\throws Exception

\return The read value

*/

virtual int8_t ReadS8() = 0;

/*! Read a int16_t from the current cursor position and advance the cursor by 2 bytes

\throws Exception

\return The read value

*/

virtual int16_t ReadS16() = 0;

/*! Read a int32_t from the current cursor position and advance the cursor by 4 bytes

\throws Exception

\return The read value

*/

virtual int32_t ReadS32() = 0;

/*! Read a int64_t from the current cursor position and advance the cursor by 8 bytes

\throws Exception

\return The read value

*/

virtual int64_t ReadS64() = 0;

/*! Read a pointer from the current cursor position and advance it that many bytes

\throws Exception

\return The value that was read

*/

virtual uint64_t ReadPointer() = 0;

/*! Get the current cursor position

\return The current cursor position

*/

virtual uint64_t GetOffset() const = 0;

/*! Set the cursor position

\param offset The new cursor position

*/

virtual void Seek(uint64_t offset) = 0;

/*! Set the cursor position, relative to the current position

\param offset Offset to the current cursor position

*/

virtual void SeekRelative(int64_t offset) = 0;

};

class ObjCProcessor {

struct Types {

Ref<Type> id;

Ref<Type> sel;

Ref<Type> BOOL;

} m_types;

struct TypeNames {

QualifiedName nsInteger;

QualifiedName nsuInteger;

QualifiedName cgFloat;

QualifiedName cfStringFlag;

QualifiedName cfString;

QualifiedName cfStringUTF16;

QualifiedName imageInfoFlags;

QualifiedName imageInfoSwiftVersion;

QualifiedName imageInfo;

QualifiedName methodEntry;

QualifiedName method;

QualifiedName methodList;

QualifiedName classRO;

QualifiedName cls;

QualifiedName category;

QualifiedName protocol;

QualifiedName protocolList;

QualifiedName ivar;

QualifiedName ivarList;

QualifiedName nsConstantArray;

QualifiedName nsConstantDictionary;

QualifiedName nsConstantDoubleNumber;

QualifiedName nsConstantFloatNumber;

QualifiedName nsConstantIntegerNumber;

QualifiedName nsConstantDate;

QualifiedName nsConstantData;

} m_typeNames;

// TODO(WeiN76LQh): this is to avoid a bug with defining a classes protocol list in the DSC plugin. Remove once fixed

bool m_skipClassBaseProtocols;

std::map<uint64_t, Class> m_classes;

std::map<uint64_t, Class> m_categories;

std::map<uint64_t, Protocol> m_protocols;

std::unordered_map<uint64_t, std::string> m_selectorCache;

std::unordered_map<uint64_t, Method> m_localMethods;

// Required for workflow_objc type heuristics, should be removed when that is no longer a thing.

std::map<uint64_t, std::string> m_selRefToName;

std::map<uint64_t, std::vector<uint64_t>> m_selRefToImplementations;

std::map<uint64_t, std::vector<uint64_t>> m_selToImplementations;

// --

uint64_t ReadPointerAccountingForRelocations(ObjCReader* reader);

std::unordered_map<uint64_t, uint64_t> m_relocationPointerRewrites;

static Ref<Metadata> SerializeMethod(uint64_t loc, const Method& method);

static Ref<Metadata> SerializeClass(uint64_t loc, const Class& cls);

Ref<Metadata> SerializeMetadata();

std::vector<QualifiedNameOrType> ParseEncodedType(const std::string& type);

void DefineObjCSymbol(BNSymbolType symbolType, QualifiedName typeName, const std::string& name, uint64_t addr, bool deferred);

void DefineObjCSymbol(BNSymbolType symbolType, Ref<Type> type, const std::string& name, uint64_t addr, bool deferred);

void ReadIvarList(ObjCReader* reader, ClassBase& cls, std::string_view name, view_ptr_t start);

void ReadMethodList(ObjCReader* reader, ClassBase& cls, std::string_view name, view_ptr_t start);

void ReadListOfMethodLists(ObjCReader* reader, ClassBase& cls, std::string_view name, view_ptr_t start);

void LoadClasses(ObjCReader* reader, Ref<Section> listSection);

void LoadCategories(ObjCReader* reader, Ref<Section> listSection);

void LoadProtocols(ObjCReader* reader, Ref<Section> listSection);

void GenerateClassTypes();

bool ApplyMethodType(Class& cls, Method& method, bool isInstanceMethod);

void ApplyMethodTypes(Class& cls);

std::optional<std::string> ClassNameForTargetOfPointerAt(ObjCReader* reader, uint64_t offset);

void ProcessCFStrings();

void ProcessNSConstantArrays();

void ProcessNSConstantDictionaries();

void ProcessNSConstantIntegerNumbers();

void ProcessNSConstantFloatingPointNumbers();

void ProcessNSConstantDatas();

void PostProcessObjCSections(ObjCReader* reader);

protected:

Ref<BinaryView> m_data;

Ref<Logger> m_logger;

virtual uint64_t GetObjCRelativeMethodBaseAddress(ObjCReader* reader);

virtual void GetRelativeMethod(ObjCReader* reader, method_t& meth);

virtual std::shared_ptr<ObjCReader> GetReader() = 0;

// Because an objective-c processor might have access to other non-view symbols that we want to retrieve.

// By default, this will just get symbol at the address in the view.

virtual Ref<Symbol> GetSymbol(uint64_t address);

virtual Ref<Section> GetSectionWithName(const char* sectionName);

public:

virtual ~ObjCProcessor() = default;

ObjCProcessor(BinaryView* data, const char* loggerName, bool skipClassBaseProtocols = false);

void ProcessObjCData();

void ProcessObjCLiterals();

void AddRelocatedPointer(uint64_t location, uint64_t rewrite);

};

}