/*
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 "Debugger.h"
#include "MemoryRegions.h"
#include "SymbolManager.h"
#include "Configuration.h"
#include "ResultViewModel.h"
#include "Util.h"

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

#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 {
	unsigned int	prev_size;	/* Size of previous chunk (if free).  */
	unsigned int	size;		/* Size in bytes, including overhead. */

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

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


//------------------------------------------------------------------------------
// Name: DialogHeap(QWidget *parent)
// Desc: 
//------------------------------------------------------------------------------
DialogHeap::DialogHeap(QWidget *parent) : QDialog(parent) {
	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) {
		edb::v1::dumpDataRange(item->block, item->block + item->size);
	}
}

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

	foreach(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();
			edb::v1::log("[Heap Analyzer] libc library appears to be: %s", qPrintable(libcName));
			continue;
		}
		
		if(fileinfo.completeBaseName().startsWith("libc.so")) {
			libcName = fileinfo.completeBaseName() + "." + fileinfo.suffix();
			edb::v1::log("[Heap Analyzer] libc library appears to be: %s", qPrintable(libcName));
			continue;
		}

		if(fileinfo.completeBaseName().startsWith("ld-")) {
			ldName = fileinfo.completeBaseName() + "." + fileinfo.suffix();
			edb::v1::log("[Heap Analyzer] ld library appears to be: %s", qPrintable(ldName));
			continue;
		}
	}
}

//------------------------------------------------------------------------------
// Name: detectPointers()
// Desc: 
//------------------------------------------------------------------------------
void DialogHeap::detectPointers() {
	
	QVector<edb::address_t> targets;
	QVector<Result> &results = m_Model->results();
	
	// the potential targets
	foreach(Result result, results) {
		const edb::address_t addr = result.block + sizeof(struct malloc_chunk *) * 2;
		targets.append(addr);
	}
	
	// look for potential pointers
	for(QVector<Result>::iterator it = results.begin(); it != results.end(); ++it) {
		if(it->data.isEmpty()) {
		
			edb::address_t pointer;
			if(it->size > sizeof(edb::address_t)) {
				if(edb::v1::debuggerBase->readBytes(it->block + sizeof(struct malloc_chunk *) * 2, &pointer, sizeof(pointer))) {
					
					if(targets.contains(pointer)) {
					#if QT_POINTER_SIZE == 4
						it->data.sprintf("dword ptr [" EDB_FMT_PTR "]", pointer);
					#elif QT_POINTER_SIZE == 8
						it->data.sprintf("qword ptr [" EDB_FMT_PTR "]", pointer);
					#endif
					}
				}
			}
		}	
	}

	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();
	
	if(start_address != 0x00000000 && end_address != 0x00000000) {
#ifdef Q_OS_LINUX
		malloc_chunk currentChunk;
		malloc_chunk nextChunk;
		edb::address_t currentChunkAddress = start_address;

		const edb::address_t how_many = end_address - start_address;
		while(currentChunkAddress != end_address) {
			// read in the current chunk..
			edb::v1::debuggerBase->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(), 
					"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
				edb::v1::debuggerBase->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;
				int asciisz;
				if(edb::v1::getASCIIStringAtAddress(
						currentChunkAddress + sizeof(struct malloc_chunk *) * 2, 
						asciiData, 
						edb::v1::config().minAsciiStringLength,
						currentChunk.chunksize(),
						asciisz)) {
						
					data = asciiData;
				}

				const Result r = { 
					currentChunkAddress, 
					currentChunk.chunksize() + sizeof(unsigned int), 
					nextChunk.prev_inuse() ? "Busy" : "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));
		}
		
		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	= 0;
	edb::address_t end_address		= 0;

	QString libcName;
	QString ldName;

	getLibraryNames(libcName, ldName);
#ifdef Q_OS_LINUX
	s = edb::v1::symbolManager().findSymbol(libcName + "::__curbrk");
	if(s != 0) {
		end_address = s->address;
	} else {
		QMessageBox::information(this, "__curbrk symbol not found in libc", "Could not find the symbol for <strong>__curbrk</strong> in your libc, perhaps you need to regenerate your symbols?");
		edb::v1::log("[Heap Analyzer] __curbrk symbol not found in libc");
		return;
	}


	s = edb::v1::symbolManager().findSymbol(ldName + "::__curbrk");
	if(s != 0) {
		start_address = s->address;
	} else {
	
		edb::v1::log("[Heap Analyzer] __curbrk symbol not found in ld, falling back on heuristic! This may or may not work.");
		start_address = end_address - 0x4f8;
		
		edb::address_t test_addr;
		edb::address_t heap_symbol = start_address - 4 * sizeof(void *);
					
		edb::v1::debuggerBase->readBytes(heap_symbol, &test_addr, sizeof(test_addr));
		
		if(test_addr != edb::v1::debuggerBase->pageSize()) {
			QMessageBox::information(this, "__curbrk symbol not found in ld", "Could not find the symbol for <strong>__curbrk</strong> in your ld, perhaps you need to regenerate your symbols?");
			edb::v1::log("[Heap Analyzer] __curbrk symbol not found in ld");
			return;
		}
	}

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

	// read the contents of those symbols
	edb::v1::debuggerBase->readBytes(end_address, &end_address, sizeof(end_address));
	edb::v1::debuggerBase->readBytes(start_address, &start_address, sizeof(start_address));
#else
	#error "Unsupported Platform"
#endif	
	edb::v1::log("[Heap Analyzer] heap start: " EDB_FMT_PTR, start_address);
	edb::v1::log("[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);

}
