/*
Copyright (C) 2006 Evan Teran
                   eteran@alum.rit.edu

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

#include "DialogHeap.h"
#include "DebuggerCoreInterface.h"
#include "PluginAPI.h"
#include "MemoryRegions.h"
#include "SymbolManager.h"
#include "Util.h"
#include "Configuration.h"

#include <QHeaderView>
#include <QMessageBox>
#include <QVector>
#include <QString>
#include <QFileInfo>
#include <QSortFilterProxyModel>
#include <QtDebug>

#include <boost/bind.hpp>

#include <algorithm>

#define PREV_INUSE 0x1
#define IS_MMAPPED 0x2
#define NON_MAIN_ARENA 0x4
#define SIZE_BITS (PREV_INUSE|IS_MMAPPED|NON_MAIN_ARENA)

#define next_chunk(p, c) ((p) + ((c).chunksize()))
#define prev_chunk(p, c) ((p) - ((c).prev_size)))


struct malloc_chunk {
	ulong	prev_size;	/* Size of previous chunk (if free).  */
	ulong	size;		/* Size in bytes, including overhead. */

	struct malloc_chunk* fd;	/* double links -- used only if free. */
	struct malloc_chunk* bk;

	ulong chunksize() const	{ return size & ~(SIZE_BITS); }
	bool prev_inuse() const	{ return size & PREV_INUSE; }
};


//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
DialogHeap::DialogHeap(PluginAPI &api, QWidget *parent) : QDialog(parent), m_API(api) {
	ui.setupUi(this);
	
	m_Model = new ResultViewModel(this);
	ui.tableView->setModel(m_Model);
	
	ui.tableView->verticalHeader()->hide();
	ui.tableView->horizontalHeader()->setStretchLastSection(true);
	ui.tableView->horizontalHeader()->setResizeMode(QHeaderView::ResizeToContents);
}

//------------------------------------------------------------------------------
// Name: showEvent(QShowEvent *)
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::showEvent(QShowEvent *) {
	m_Model->clearResults();
	ui.progressBar->setValue(0);
}

//------------------------------------------------------------------------------
// Name: on_resultTable_cellDoubleClicked(int, int)
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::on_tableView_doubleClicked(const QModelIndex & index) {

	// NOTE: remember that if we use a sort filter, we need to map the indexes
	// to get at the data we need

	const Result *const item = static_cast<Result *>(index.internalPointer());
	
	if(item != 0) {
		m_API.dumpDataRange(item->block, item->block + item->size, false);
	}
}

//------------------------------------------------------------------------------
// Name: getLibraryNames(QString &libcName, QString &ldName) const
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::getLibraryNames(QString &libcName, QString &ldName) const {
	const QStringList libs = m_API.loadedLibraries();

	foreach(const QString & s, libs) {
		if(!ldName.isEmpty() && !libcName.isEmpty()) {
			break;
		}

		const QFileInfo fileinfo(s);

		// this tries its best to cover all possible libc library versioning 
		// possibilities we need to find out if this is 100% accurate, so far 
		// seems correct based on my system
		
		if(fileinfo.completeBaseName().startsWith("libc-")) {
			libcName = fileinfo.completeBaseName() + "." + fileinfo.suffix();
			qDebug() << "[Heap Analyzer] libc library appears to be:" << libcName;
			continue;
		}
		
		if(fileinfo.completeBaseName().startsWith("libc.so")) {
			libcName = fileinfo.completeBaseName() + "." + fileinfo.suffix();
			qDebug() << "[Heap Analyzer] libc library appears to be:" << libcName;
			continue;
		}

		if(fileinfo.completeBaseName().startsWith("ld-")) {
			ldName = fileinfo.completeBaseName() + "." + fileinfo.suffix();
			qDebug() << "[Heap Analyzer] ld library appears to be:" << libcName;
			continue;
		}
	}
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::processPotentialPointer(const QSet<edb::address_t> &targets, Result &result) {
	if(result.data.isEmpty()) {
		edb::address_t pointer;
		if(result.size > sizeof(edb::address_t)) {
			if(m_API.debuggerCore->readBytes(result.block + sizeof(struct malloc_chunk *) * 2, &pointer, sizeof(pointer))) {

				if(targets.contains(pointer)) {
				#if QT_POINTER_SIZE == 4
					result.data = QString("dword ptr [%1]").arg(pointer, EDB_MAX_HEX, 16, QChar('0'));
				#elif QT_POINTER_SIZE == 8
					result.data = QString("qword ptr [%1]").arg(pointer, EDB_MAX_HEX, 16, QChar('0'));
				#endif
				}
			}
		}
	}
}

//------------------------------------------------------------------------------
// Name: detectPointers()
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::detectPointers() {
	
	qDebug("[Heap Analyzer] detecting pointers in heap blocks");
	
	QSet<edb::address_t> targets;
	QVector<Result> &results = m_Model->results();
	
	// the potential targets
	qDebug("[Heap Analyzer] collecting possible targets addresses");
	foreach(const Result &result, results) {
		targets.insert(result.block + sizeof(struct malloc_chunk *) * 2);
	}
	
	std::for_each(results.begin(), results.end(), boost::bind(&DialogHeap::processPotentialPointer, this, targets, _1));

	m_Model->update();
}

//------------------------------------------------------------------------------
// Name: collectBlocks(edb::address_t start_address, edb::address_t end_address)
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::collectBlocks(edb::address_t start_address, edb::address_t end_address) {
	m_Model->clearResults();
	
	const int minStringLength = m_API.config().minStringLength;
	
	if(start_address != edb::nullptr && end_address != edb::nullptr) {
#ifdef Q_OS_LINUX
		malloc_chunk currentChunk;
		malloc_chunk nextChunk;
		edb::address_t currentChunkAddress = start_address;

		m_Model->setUpdatesEnabled(false);

		const edb::address_t how_many = end_address - start_address;
		while(currentChunkAddress != end_address) {
			// read in the current chunk..
			m_API.debuggerCore->readBytes(currentChunkAddress, &currentChunk, sizeof(currentChunk));

			// figure out the address of the next chunk
			const edb::address_t nextChunkAddress = next_chunk(currentChunkAddress, currentChunk);

			// is this the last chunk (if so, it's the 'top')
			if(nextChunkAddress == end_address) {

				const Result r = {
					currentChunkAddress, 
					currentChunk.chunksize(), 
					tr("Top"),
					QString() };
				
				m_Model->addResult(r);

			} else {

				// make sure we aren't following a broken heap...
				if(nextChunkAddress > end_address || nextChunkAddress < start_address) {
					break;
				}
				
				QString data;

				// read in the next chunk
				m_API.debuggerCore->readBytes(nextChunkAddress, &nextChunk, sizeof(nextChunk));

				// if this block is a container for an ascii string, display it...
				// there is a lot of room for improvement here, but it's a start
				QString asciiData;
				QString utf16Data;
				int asciisz;
				int utf16sz;
				if(m_API.getASCIIStringAtAddress(
						currentChunkAddress + sizeof(struct malloc_chunk *) * 2, 
						asciiData, 
						minStringLength,
						currentChunk.chunksize(),
						asciisz)) {
						
					data = asciiData;
				} else if(m_API.getUTF16StringAtAddress(
						currentChunkAddress + sizeof(struct malloc_chunk *) * 2, 
						utf16Data,
						minStringLength,
						currentChunk.chunksize(),
						utf16sz)) {
					data = utf16Data;
				}

				const Result r = { 
					currentChunkAddress, 
					currentChunk.chunksize() + sizeof(unsigned int), 
					nextChunk.prev_inuse() ? tr("Busy") : tr("Free"),
					data };
				
				m_Model->addResult(r);	
			}

			// avoif self referencing blocks
			if(currentChunkAddress == nextChunkAddress) {
				break;
			}

			currentChunkAddress = nextChunkAddress;

			ui.progressBar->setValue(CalculatePercentage(currentChunkAddress - start_address, how_many));
		}
		
		m_Model->setUpdatesEnabled(true);
		detectPointers();
		
#else
	#error "Unsupported Platform"
#endif
	}
}

//------------------------------------------------------------------------------
// Name: doFind()
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::doFind() {
	// get both the libc and ld symbols of __curbrk
	// this will be the 'before/after libc' addresses
	
	// we should get them from the regions we see
	const SymbolManager::Symbol *s;

	edb::address_t start_address	= edb::nullptr;
	edb::address_t end_address		= edb::nullptr;

	QString libcName;
	QString ldName;

	getLibraryNames(libcName, ldName);
#ifdef Q_OS_LINUX
	s = m_API.symbolManager().findSymbol(libcName + "::__curbrk");
	if(s != 0) {
		end_address = s->address;
	} else {
		QMessageBox::information(this, tr("__curbrk symbol not found in libc"), tr("Could not find the symbol for <strong>__curbrk</strong> in your libc, perhaps you need to regenerate your symbols?"));
		qDebug("[Heap Analyzer] __curbrk symbol not found in libc");
		return;
	}
	
	s = m_API.symbolManager().findSymbol(ldName + "::__curbrk");
	if(s != 0) {
		start_address = s->address;
	} else {
	
		qDebug("[Heap Analyzer] __curbrk symbol not found in ld, falling back on heuristic! This may or may not work.");
		edb::address_t test_addr;
#if defined(EDB_X86_64)
		start_address = end_address - 0x7b8;
		edb::address_t heap_symbol = start_address - (sizeof(edb::address_t) + sizeof(edb::address_t) + sizeof(edb::address_t) + sizeof(edb::address_t));
#else
		start_address = end_address - 0x4f8;
		edb::address_t heap_symbol = start_address - (sizeof(edb::address_t) + sizeof(edb::address_t) + sizeof(edb::address_t) + sizeof(unsigned int));
#endif			
		m_API.debuggerCore->readBytes(heap_symbol, &test_addr, sizeof(test_addr));
		
		if(test_addr != m_API.debuggerCore->pageSize()) {
			QMessageBox::information(this, tr("__curbrk symbol not found in ld"), tr("Could not find the symbol for <strong>__curbrk</strong> in your ld, perhaps you need to regenerate your symbols?"));
			qDebug("[Heap Analyzer] __curbrk symbol not found in ld");
			return;
		}
	}

	qDebug("[Heap Analyzer] heap start symbol : " EDB_FMT_PTR, start_address);
	qDebug("[Heap Analyzer] heap end symbol   : " EDB_FMT_PTR, end_address);

	// read the contents of those symbols
	m_API.debuggerCore->readBytes(end_address, &end_address, sizeof(end_address));
	m_API.debuggerCore->readBytes(start_address, &start_address, sizeof(start_address));
#else
	#error "Unsupported Platform"
#endif

	qDebug("[Heap Analyzer] heap start : " EDB_FMT_PTR, start_address);
	qDebug("[Heap Analyzer] heap end   : " EDB_FMT_PTR, end_address);

	collectBlocks(start_address, end_address);
}

//------------------------------------------------------------------------------
// Name: on_btnFind_clicked()
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::on_btnFind_clicked() {
	ui.btnFind->setEnabled(false);
	ui.progressBar->setValue(0);
	doFind();
	ui.progressBar->setValue(100);
	ui.btnFind->setEnabled(true);

}
