hashmanager.cpp

软件是使用VC70

/* 
 * Copyright (C) 2001-2005 Jacek Sieka, arnetheduck on gmail point com
 *
 * 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., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#include "stdinc.h"
#include "DCPlusPlus.h"

#include "HashManager.h"
#include "ResourceManager.h"
#include "SimpleXML.h"
#include "LogManager.h"
#include "File.h"

#define HASH_FILE_VERSION_STRING "2"
static const u_int32_t HASH_FILE_VERSION=2;

bool HashManager::checkTTH(const string& aFileName, int64_t aSize, u_int32_t aTimeStamp) {
	Lock l(cs);
	if(!store.checkTTH(aFileName, aSize, aTimeStamp)) {
		hasher.hashFile(aFileName, aSize);
		return false;
	}
	return true;
}

const TTHValue& HashManager::getTTH(const string& aFileName, int64_t aSize) throw(HashException) {
	Lock l(cs);
	const TTHValue* tth = store.getTTH(aFileName);
	if(tth == NULL){
		hasher.hashFile(aFileName, aSize);
		throw HashException(Util::emptyString);
	}
	return *tth;
}

bool HashManager::getTree(const TTHValue& root, TigerTree& tt) {
	Lock l(cs);
	return store.getTree(root, tt);
}

void HashManager::hashDone(const string& aFileName, u_int32_t aTimeStamp, const TigerTree& tth, int64_t speed) {
	const TTHValue* root = NULL;
	{
		Lock l(cs);
		store.addFile(aFileName, aTimeStamp, tth, true);
		root = store.getTTH(aFileName);
	}

	if(root != NULL) {
		fire(HashManagerListener::TTHDone(), aFileName, *root);
	}

	string fn = aFileName;
	if(count(fn.begin(), fn.end(), PATH_SEPARATOR) >= 2) {
		string::size_type i = fn.rfind(PATH_SEPARATOR);
		i = fn.rfind(PATH_SEPARATOR, i-1);
		fn.erase(0, i);
		fn.insert(0, "...");
	}
	if(speed > 0) {
		LogManager::getInstance()->message(STRING(HASHING_FINISHED) + fn + " (" + Util::formatBytes(speed) + "/s)");
	} else if(speed <= 0) {
		LogManager::getInstance()->message(STRING(HASHING_FINISHED) + fn);
	}
}

void HashManager::HashStore::addFile(const string& aFileName, u_int32_t aTimeStamp, const TigerTree& tth, bool aUsed) {
	if(!addTree(tth))
		return;

	string fname = Text::toLower(Util::getFileName(aFileName));
	string fpath = Text::toLower(Util::getFilePath(aFileName));

	FileInfoList& fileList = fileIndex[fpath];

	FileInfoIter j = find_if(fileList.begin(), fileList.end(), FileInfo::StringComp(fname));
	if(j != fileList.end()) {
		fileList.erase(j);
	}

	fileList.push_back(FileInfo(fname, tth.getRoot(), aTimeStamp, aUsed));
	dirty = true;
}

bool HashManager::HashStore::addTree(const TigerTree& tth) {
	try {
		if(treeIndex.find(tth.getRoot()) == treeIndex.end()) {
			int64_t index = addLeaves(tth.getLeaves());
			if(index == STORE_FAILED)
				return false;

			dcassert((index != SMALL_TREE) || ((tth.getLeaves().size() == 1) && (tth.getRoot() == tth.getLeaves()[0])));

			treeIndex.insert(make_pair(tth.getRoot(), TreeInfo(tth.getFileSize(), index, tth.getBlockSize())));
		}
	} catch(const Exception& ) {
		return false;
	}
	dirty = true;
	return true;
}

int64_t HashManager::HashStore::addLeaves(const TigerTree::MerkleList& leaves) {
	if(leaves.size() == 1)
		return SMALL_TREE;

	File f(dataFile, File::RW, File::OPEN);
	f.setPos(0);
	int64_t pos = 0;
	size_t n = sizeof(pos);
	if(f.read(&pos, n) != sizeof(pos))
		return STORE_FAILED;

	// Check if we should grow the file, we grow by a meg at a time...
	int64_t datsz = f.getSize();
	if((pos + (int64_t)(leaves.size() * TTHValue::SIZE)) >= datsz) {
		f.setPos(datsz + 1024*1024);
		f.setEOF();
	}
	f.setPos(pos);
	dcassert(leaves.size() > 0);
	f.write(leaves[0].data, (leaves.size() * TTHValue::SIZE));
	int64_t p2 = f.getPos();
	f.setPos(0);
	f.write(&p2, sizeof(p2));
	return pos;
}

bool HashManager::HashStore::getTree(const TTHValue& root, TigerTree& tth) {
	TreeIter i = treeIndex.find(root);
	if(i == treeIndex.end())
		return false;
	
	TreeInfo& ti = i->second;

	if(ti.getIndex() == SMALL_TREE) {
		tth = TigerTree(ti.getSize(), ti.getBlockSize(), i->first);
		return true;
	}

	try {
		File f(dataFile, File::READ, File::OPEN);

		f.setPos(ti.getIndex());
		size_t datalen = TigerTree::calcBlocks(ti.getSize(), ti.getBlockSize()) * TTHValue::SIZE;
		AutoArray<u_int8_t> buf(datalen);
		f.read((u_int8_t*)buf, datalen);
		tth = TigerTree(ti.getSize(), ti.getBlockSize(), buf);
		if(!(tth.getRoot() == i->first))
			return false;
	} catch(const FileException& ) {
		return false;
	}
	return true;
}

bool HashManager::HashStore::checkTTH(const string& aFileName, int64_t aSize, u_int32_t aTimeStamp) {
	string fname = Text::toLower(Util::getFileName(aFileName));
	string fpath = Text::toLower(Util::getFilePath(aFileName));
	DirIter i = fileIndex.find(fpath);
	if(i != fileIndex.end()) {
		FileInfoIter j = find_if(i->second.begin(), i->second.end(), FileInfo::StringComp(fname));
		if(j != i->second.end()) {
			FileInfo& fi = *j;
			TreeIter ti = treeIndex.find(fi.getRoot());
			if(ti == treeIndex.end() || ti->second.getSize() != aSize || fi.getTimeStamp() != aTimeStamp) {
				i->second.erase(j);
				dirty = true;
				return false;
			}
			return true;
		}
	} 
	return false;
}

const TTHValue* HashManager::HashStore::getTTH(const string& aFileName) {
	string fname = Text::toLower(Util::getFileName(aFileName));
	string fpath = Text::toLower(Util::getFilePath(aFileName));

	DirIter i = fileIndex.find(fpath);
	if(i != fileIndex.end()) {
		FileInfoIter j = find_if(i->second.begin(), i->second.end(), FileInfo::StringComp(fname));
		if(j != i->second.end()) {
			j->setUsed(true);
			return &(j->getRoot());
		}
	}
	return NULL;
}

void HashManager::HashStore::rebuild() {
	string tmpName = dataFile + ".tmp";
	try {
		File::renameFile(dataFile, tmpName);
		File src(tmpName, File::READ, File::OPEN);
		createDataFile(dataFile);
		File tgt(dataFile, File::WRITE, File::OPEN);
		tgt.setEndPos(0);

		TreeMap newIndex;
		for(DirIter i = fileIndex.begin(); i != fileIndex.end(); ++i) {
			for(FileInfoIter j = i->second.begin(); j != i->second.end(); ++j) {
				if(!j->getUsed())
					continue;

				TreeIter k = treeIndex.find(j->getRoot());
				if(k != treeIndex.end()) {
					newIndex[j->getRoot()] = k->second;
				}
			}
		}
		
		treeIndex = newIndex;

		vector<u_int8_t> buf;
		for(TreeIter i = treeIndex.begin(); i != treeIndex.end(); ++i) {
			src.setPos(i->second.getIndex());
			size_t datalen = TigerTree::calcBlocks(i->second.getSize(), i->second.getBlockSize()) * TTHValue::SIZE;
			buf.resize(datalen);
			src.read(&buf[0], datalen);
			tgt.write(&buf[0], datalen);
		}

		int64_t pos = tgt.getPos();
		tgt.setPos(0);
		tgt.write(&pos, sizeof(pos));
		// Set size to the nearest mb boundary
		tgt.setSize(((tgt.getSize() + 1024*1024 - 1) / 1024*1024) * 1024*1024);
	} catch(const FileException&) {
		try {
			File::renameFile(tmpName, dataFile);
		} catch(const FileException&) {
			// Too bad...
		}
	}

}

#define LITERAL(x) x, sizeof(x)-1
#define CHECKESCAPE(n) SimpleXML::escape(n, tmp, true)

void HashManager::HashStore::save() {
	if(dirty) {
		try {
			File ff(indexFile + ".tmp", File::WRITE, File::CREATE | File::TRUNCATE);
			BufferedOutputStream<false> f(&ff);

			string tmp;
			string b32tmp;

			f.write(SimpleXML::utf8Header);
			f.write(LITERAL("<HashStore Version=\"" HASH_FILE_VERSION_STRING "\">\r\n"));

			f.write(LITERAL("\t<Trees>\r\n"));

			for(TreeIter i = treeIndex.begin(); i != treeIndex.end(); ++i) {
				const TreeInfo& ti = i->second;
				f.write(LITERAL("\t\t<Hash Type=\"TTH\" Index=\""));
				f.write(Util::toString(ti.getIndex()));
				f.write(LITERAL("\" BlockSize=\""));
				f.write(Util::toString(ti.getBlockSize()));
				f.write(LITERAL("\" Size=\""));
				f.write(Util::toString(ti.getSize()));
				f.write(LITERAL("\" Root=\""));
				b32tmp.clear();
				f.write(i->first.toBase32(b32tmp));
				f.write(LITERAL("\"/>\r\n"));
			}

			f.write(LITERAL("\t</Trees>\r\n\t<Files>\r\n"));

			for(DirIter i = fileIndex.begin(); i != fileIndex.end(); ++i) {
				const string& dir = i->first;
				for(FileInfoIter j = i->second.begin(); j != i->second.end(); ++j) {
					const FileInfo& fi = *j;
					f.write(LITERAL("\t\t<File Name=\""));
					f.write(CHECKESCAPE(dir + fi.getFileName()));
					f.write(LITERAL("\" TimeStamp=\""));
					f.write(Util::toString(fi.getTimeStamp()));
					f.write(LITERAL("\" Root=\""));
					b32tmp.clear();
					f.write(fi.getRoot().toBase32(b32tmp));
					f.write(LITERAL("\"/>\r\n"));
				}
			}
			f.write(LITERAL("\t</Files>\r\n</HashStore>"));
			f.flush();
			ff.close();
			File::deleteFile(indexFile);
			File::renameFile(indexFile + ".tmp", indexFile);

			dirty = false;
		} catch(const FileException&) {
			// Too bad...
		}
	}
}

class HashLoader : public SimpleXMLReader::CallBack {
public:
	HashLoader(HashManager::HashStore& s) : store(s), size(0), timeStamp(0), version(HASH_FILE_VERSION), inTrees(false), inFiles(false), inHashStore(false) { };
	virtual void startTag(const string& name, StringPairList& attribs, bool simple);
	virtual void endTag(const string& name, const string& data);
	
private:
	HashManager::HashStore& store;

	string file;
	int64_t size;
	u_int32_t timeStamp;
	int version;

	bool inTrees;
	bool inFiles;
	bool inHashStore;
};

void HashManager::HashStore::load() {
	try {
		HashLoader l(*this);
		SimpleXMLReader(&l).fromXML(File(indexFile, File::READ, File::OPEN).read());
	} catch(const Exception&) {
		// ...
	}