//------------------------------------------------------------------------------
//	Copyright (c) 2003, Ingo Weinhold
//
//	Permission is hereby granted, free of charge, to any person obtaining a
//	copy of this software and associated documentation files (the "Software"),
//	to deal in the Software without restriction, including without limitation
//	the rights to use, copy, modify, merge, publish, distribute, sublicense,
//	and/or sell copies of the Software, and to permit persons to whom the
//	Software is furnished to do so, subject to the following conditions:
//
//	The above copyright notice and this permission notice shall be included in
//	all copies or substantial portions of the Software.
//
//	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//	AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//	LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
//	FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
//	DEALINGS IN THE SOFTWARE.
//
//	File Name:		ElfSymbolPatcher.cpp
//	Author:			Ingo Weinhold (bonefish@users.sf.net)
//	Description:	Interface declaration of classes used for patching ELF
//					symbols. Central class is ElfSymbolPatcher. It is a kind of
//					roster, managing all necessary infos (loaded images) and
//					being able to fill ElfSymbolPatchInfos with life.
//					An ElfSymbolPatchInfo represents a symbol and is able to
//					patch/restore it. An ElfSymbolPatchGroup bundles several
//					ElfSymbolPatchInfos and can update their data, e.g.
//					when images are loaded/unloaded. It uses a
//					ElfSymbolPatcher internally and provides a more convenient
//					API for the user.
//------------------------------------------------------------------------------

#include <new>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "ElfImage.h"
#include "ElfSymbolPatcher.h"

/////////////////////////////
// ElfSymbolPatchInfo::Entry
//

class ElfSymbolPatchInfo::Entry {
public:
	static	Entry*				Create(image_id image, void*** targets,
									   int32 targetCount);
			void				Delete();

			image_id			GetImage() const	{ return fImage; }

			void				Patch(void* newAddress);

private:
								Entry();
								Entry(const Entry&);
								Entry(image_id image, void*** targets,
									  int32 targetCount);
								~Entry();

private:
			image_id			fImage;
			int32				fPatchTargetCount;
			void**				fPatchTargets[1];
};

// Create
ElfSymbolPatchInfo::Entry*
ElfSymbolPatchInfo::Entry::Create(image_id image, void*** targets,
								  int32 targetCount)
{
	if (!targets || targetCount <= 0)
		return NULL;
	void* buffer = malloc(sizeof(Entry) + sizeof(void**) * (targetCount - 1));
	Entry* entry = NULL;
	if (buffer)
		entry = new(buffer) Entry(image, targets, targetCount);
	return entry;
}

// Delete
void
ElfSymbolPatchInfo::Entry::Delete()
{
	this->~Entry();
	free(this);
}

// Patch
void
ElfSymbolPatchInfo::Entry::Patch(void* newAddress)
{
//printf("ElfSymbolPatchInfo::Entry::Patch(): patching %ld addresses\n",
//fPatchTargetCount);
	for (int i = 0; i < fPatchTargetCount; i++)
		*fPatchTargets[i] = newAddress;
}

// constructor
ElfSymbolPatchInfo::Entry::Entry(image_id image, void*** targets,
								 int32 targetCount)
	: fImage(image),
	  fPatchTargetCount(targetCount)
{
	memcpy(fPatchTargets + 0, targets, targetCount * sizeof(void**));
}

// destructor
ElfSymbolPatchInfo::Entry::~Entry()
{
}


//////////////////////
// ElfSymbolPatchInfo
//

// constructor
ElfSymbolPatchInfo::ElfSymbolPatchInfo()
	: fSymbolName(),
	  fOriginalAddress(NULL),
	  fOriginalAddressImage(-1),
	  fEntries()
{
}

// destructor
ElfSymbolPatchInfo::~ElfSymbolPatchInfo()
{
	Unset();
}

// InitCheck
status_t
ElfSymbolPatchInfo::InitCheck() const
{
	return (fOriginalAddress && fSymbolName.Length() ? B_OK : B_NO_INIT);
}

// GetSymbolName
const char*
ElfSymbolPatchInfo::GetSymbolName() const
{
	return fSymbolName.String();
}

// GetOriginalAddress
void*
ElfSymbolPatchInfo::GetOriginalAddress() const
{
	return fOriginalAddress;
}

// GetOriginalAddressImage
image_id
ElfSymbolPatchInfo::GetOriginalAddressImage() const
{
	return fOriginalAddressImage;
}

// Patch
status_t
ElfSymbolPatchInfo::Patch(void* newAddress)
{
//printf("ElfSymbolPatchInfo::Patch(): patching %ld images\n",
//fEntries.CountItems());
	status_t error = InitCheck();
	if (error == B_OK) {
		for (int i = 0; Entry* entry = EntryAt(i); i++)
			entry->Patch(newAddress);
	}
	return error;
}

// Restore
status_t
ElfSymbolPatchInfo::Restore()
{
	return Patch(fOriginalAddress);
}

// Unset
void
ElfSymbolPatchInfo::Unset()
{
	for (int i = 0; Entry* entry = EntryAt(i); i++)
		entry->Delete();
	fEntries.MakeEmpty();
	fSymbolName.SetTo("");
	fOriginalAddress = NULL;
	fOriginalAddressImage = -1;
}

// SetSymbolName
status_t
ElfSymbolPatchInfo::SetSymbolName(const char* name)
{
	fSymbolName.SetTo(name);
	if (name && fSymbolName != name)
		return B_NO_MEMORY;
	return B_OK;
}

// SetOriginalAddress
void
ElfSymbolPatchInfo::SetOriginalAddress(void* address, image_id image)
{
	fOriginalAddress = address;
	fOriginalAddressImage = image;
}

// CreateEntry
status_t
ElfSymbolPatchInfo::CreateEntry(image_id image, BList* targets)
{
	if (!targets || targets->CountItems() == 0)
		return B_BAD_VALUE;
	Entry* entry = Entry::Create(image, (void***)targets->Items(),
								 targets->CountItems());
	if (!entry)
		return B_NO_MEMORY;
	if (!fEntries.AddItem(entry)) {
		entry->Delete();
		return B_NO_MEMORY;
	}
	return B_OK;
}

// DeleteEntry
bool
ElfSymbolPatchInfo::DeleteEntry(image_id image)
{
	for (int i = 0; Entry* entry = EntryAt(i); i++) {
		if (entry->GetImage() == image) {
			fEntries.RemoveItem(i);
			entry->Delete();
			return true;
		}
	}
	return false;
}

// EntryAt
ElfSymbolPatchInfo::Entry*
ElfSymbolPatchInfo::EntryAt(int32 index)
{
	return (Entry*)fEntries.ItemAt(index);
}

// EntryFor
ElfSymbolPatchInfo::Entry*
ElfSymbolPatchInfo::EntryFor(image_id image)
{
	for (int i = 0; Entry* entry = EntryAt(i); i++) {
		if (entry->GetImage() == image)
			return entry;
	}
	return NULL;
}


/////////////////
// UpdateAdapter
//

// constructor
ElfSymbolPatcher::UpdateAdapter::UpdateAdapter()
{
}

// destructor
ElfSymbolPatcher::UpdateAdapter::~UpdateAdapter()
{
}

// ImageAdded
void
ElfSymbolPatcher::UpdateAdapter::ImageAdded(ElfImage* image)
{
}

// ImageRemoved
void
ElfSymbolPatcher::UpdateAdapter::ImageRemoved(ElfImage* image)
{
}


////////////////////
// ElfSymbolPatcher
//

// constructor
ElfSymbolPatcher::ElfSymbolPatcher()
	: fImages(),
	  fInitStatus(B_NO_INIT)
{
	fInitStatus = _Init();
	if (fInitStatus != B_OK)
		_Cleanup();
}

// destructor
ElfSymbolPatcher::~ElfSymbolPatcher()
{
	_Cleanup();
}

// InitCheck
status_t
ElfSymbolPatcher::InitCheck() const
{
	return fInitStatus;
}

// Update
status_t
ElfSymbolPatcher::Update(UpdateAdapter* updateAdapter)
{
	if (InitCheck() != B_OK)
		return B_NO_INIT;
	// remove obsolete images
	int32 count = fImages.CountItems();
	for (int i = count - 1; i >= 0; i--) {
		ElfImage* image = _ImageAt(i);
		image_info info;
		if (get_image_info(image->GetID(), &info) != B_OK) {
			if (updateAdapter)
				updateAdapter->ImageRemoved(image);
			fImages.RemoveItem(i);
			delete image;
		}
	}
	// add new images
	status_t error = B_OK;
	image_info info;
	int32 cookie = 0;
	while (get_next_image_info(0, &cookie, &info) == B_OK) {
		ElfImage* image = _ImageForID(info.id);
		if (image)
			continue;
		image = new(std::nothrow) ElfImage;
		if (!image)
			return B_NO_MEMORY;
		if (!fImages.AddItem(image)) {
			delete image;
			return B_NO_MEMORY;
		}
		error = image->SetTo(info.id);
		if (updateAdapter)
			updateAdapter->ImageAdded(image);
	}
	return error;
}

// Unload
void
ElfSymbolPatcher::Unload()
{
	for (int i = 0; ElfImage* image = _ImageAt(i); i++)
		image->Unload();
}

// GetSymbolPatchInfo
status_t
ElfSymbolPatcher::GetSymbolPatchInfo(const char* symbolName,
									 ElfSymbolPatchInfo* info)
{
	// check parameters and intialization
	if (!symbolName || !info)
		return B_BAD_VALUE;
	if (InitCheck() != B_OK)
		return B_NO_INIT;
	// set the symbol name
	info->Unset();
	status_t error = info->SetSymbolName(symbolName);
	if (error != B_OK)
		return error;
	for (int i = 0; ElfImage* image = _ImageAt(i); i++) {
//printf("searching in image: %ld\n", image->GetID());
		// get the symbol's relocations
		BList patchTargets;
		error = image->GetSymbolRelocations(symbolName, &patchTargets);
		if (error != B_OK)
			break;
		if (patchTargets.CountItems() > 0) {
			error = info->CreateEntry(image->GetID(), &patchTargets);
			if (error != B_OK)
				break;
		}
		// get the symbol's address
		void* address = NULL;
		if (image->FindSymbol(symbolName, &address) == B_OK && address) {
			if (info->GetOriginalAddress()) {
				// A symbol with that name lives in at least two images.
				// Better bail out.
// TODO: That doesn't work so well (on gcc 4). E.g. the libsupc++ symbols might
// appear in several images.
//printf("Found the symbol in more than one image!\n");
				error = B_ERROR;
				break;
			} else
				info->SetOriginalAddress(address, image->GetID());
		}
	}
	// set the symbol address
	if (!info->GetOriginalAddress())
{
//printf("Symbol not found in any image!\n");
		error = B_ERROR;
}
	// cleanup on error
	if (error != B_OK)
		info->Unset();
	return error;
}

// UpdateSymbolPatchInfo
status_t
ElfSymbolPatcher::UpdateSymbolPatchInfo(ElfSymbolPatchInfo* info,
										ElfImage* image)
{
	if (!info || !image || !info->GetSymbolName())
		return B_BAD_VALUE;
	// get the symbol's relocations
	BList patchTargets;
	status_t error
		= image->GetSymbolRelocations(info->GetSymbolName(), &patchTargets);
	if (error == B_OK)
		error = info->CreateEntry(image->GetID(), &patchTargets);
	return error;
}

// _Init
status_t
ElfSymbolPatcher::_Init()
{
	status_t error = B_OK;
	image_info info;
	int32 cookie = 0;
	while (get_next_image_info(0, &cookie, &info) == B_OK) {
		ElfImage* image = new(std::nothrow) ElfImage;
		if (!image)
			return B_NO_MEMORY;
		if (!fImages.AddItem(image)) {
			delete image;
			return B_NO_MEMORY;
		}
		error = image->SetTo(info.id);
	}
	return error;
}

// _Cleanup
void
ElfSymbolPatcher::_Cleanup()
{
	for (int i = 0; ElfImage* image = _ImageAt(i); i++)
		delete image;
	fImages.MakeEmpty();
}

// _ImageAt
ElfImage*
ElfSymbolPatcher::_ImageAt(int32 index) const
{
	return (ElfImage*)fImages.ItemAt(index);
}

// _ImageForID
ElfImage*
ElfSymbolPatcher::_ImageForID(image_id id) const
{
	for (int i = 0; ElfImage* image = _ImageAt(i); i++) {
		if (image->GetID() == id)
			return image;
	}
	return NULL;
}


///////////////////////
// ElfSymbolPatchGroup
//

// constructor
ElfSymbolPatchGroup::ElfSymbolPatchGroup(ElfSymbolPatcher* patcher)
	: fPatcher(patcher),
	  fPatchInfos(),
	  fOwnsPatcher(false),
	  fPatched(false)
{
	// create a patcher if none has been supplied
	if (!fPatcher) {
		fPatcher = new(std::nothrow) ElfSymbolPatcher;
		if (fPatcher) {
			if (fPatcher->InitCheck() == B_OK)
				fOwnsPatcher = true;
			else {
				delete fPatcher;
				fPatcher = NULL;
			}
		}
	}
}

// destructor
ElfSymbolPatchGroup::~ElfSymbolPatchGroup()
{
	RemoveAllPatches();
	if (fPatcher && fOwnsPatcher)
		delete fPatcher;
}

// AddPatch
status_t
ElfSymbolPatchGroup::AddPatch(const char* symbolName, void* newAddress,
							  void** originalAddress)
{
	// check initialization and parameters
	if (!fPatcher)
		return B_NO_INIT;
	if (!symbolName || !originalAddress)
		return B_BAD_VALUE;
	// allocate patch info
	PatchInfo* patchInfo = new(std::nothrow) PatchInfo;
	if (!patchInfo)
		return B_NO_MEMORY;
	// init and add the patch info
	status_t error = fPatcher->GetSymbolPatchInfo(symbolName, patchInfo);
	if (error == B_OK) {
		if (fPatchInfos.AddItem(patchInfo)) {
			patchInfo->fNewAddress = newAddress;
			*originalAddress = patchInfo->GetOriginalAddress();
		} else
			error = B_NO_MEMORY;
	}
	// cleanup on failure
	if (error != B_OK && patchInfo) {
		fPatchInfos.RemoveItem(patchInfo);
		delete patchInfo;
	}
	return error;
}

// RemoveAllPatches
void
ElfSymbolPatchGroup::RemoveAllPatches()
{
	for (int i = 0; PatchInfo* info = (PatchInfo*)fPatchInfos.ItemAt(i); i++)
		delete info;
	fPatchInfos.MakeEmpty();
	fPatched = false;
}

// Patch
status_t
ElfSymbolPatchGroup::Patch()
{
//printf("ElfSymbolPatchGroup::Patch(): patching %ld symbols\n",
//fPatchInfos.CountItems());
	if (!fPatcher)
		return B_NO_INIT;
	if (fPatched)
		return B_OK;
	for (int i = 0; PatchInfo* info = (PatchInfo*)fPatchInfos.ItemAt(i); i++)
		info->Patch(info->fNewAddress);
	fPatched = true;
	return B_OK;
}

// Restore
status_t
ElfSymbolPatchGroup::Restore()
{
	if (!fPatcher)
		return B_NO_INIT;
	if (!fPatched)
		return B_OK;
	for (int i = 0; PatchInfo* info = (PatchInfo*)fPatchInfos.ItemAt(i); i++)
		info->Restore();
	fPatched = false;
	return B_OK;
}

// Update
status_t
ElfSymbolPatchGroup::Update()
{
	if (!fPatcher)
		return B_NO_INIT;
	return fPatcher->Update(this);
}

// ImageAdded
void
ElfSymbolPatchGroup::ImageAdded(ElfImage* image)
{
	for (int i = 0; PatchInfo* info = (PatchInfo*)fPatchInfos.ItemAt(i); i++) {
		fPatcher->UpdateSymbolPatchInfo(info, image);
		if (fPatched) {
			ElfSymbolPatchInfo::Entry* entry = info->EntryFor(image->GetID());
			if (entry)
				info->Patch(info->fNewAddress);
		}
	}
}

// ImageRemoved
void
ElfSymbolPatchGroup::ImageRemoved(ElfImage* image)
{
	int32 count = fPatchInfos.CountItems();
	for (int i = count - 1; i >= 0; i--) {
		PatchInfo* info = (PatchInfo*)fPatchInfos.ItemAt(i);
		if (info->GetOriginalAddressImage() == image->GetID()) {
			// the image the symbol lives in, has been removed: remove the
			// complete patch info
			fPatchInfos.RemoveItem(i);
			delete info;
		} else
			info->DeleteEntry(image->GetID());
	}
}