downloadmanager.cpp

dc++（一个曾经大量使用的p2p)的源代码,dc++,开源的p2p源代码

/* 
 * Copyright (C) 2001-2003 Jacek Sieka, [e-mail]
 *
 * 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 "DownloadManager.h"
#include "ConnectionManager.h"
#include "User.h"
#include "QueueManager.h"
#include "LogManager.h"
#include "CryptoManager.h"
#include "SFVReader.h"

DownloadManager* Singleton<DownloadManager>::instance = NULL;

static const string DOWNLOAD_AREA = "Downloads";
static const string ANTI_FRAG_EXT = ".antifrag";

Download::Download(QueueItem* qi) throw() : source(qi->getCurrent()->getPath()),
	target(qi->getTarget()), tempTarget(qi->getTempTarget()), 
	comp(NULL), bytesLeft(0), rollbackBuffer(NULL), rollbackSize(0) { 
	
	setSize(qi->getSize());
	if(qi->isSet(QueueItem::FLAG_USER_LIST))
		setFlag(Download::FLAG_USER_LIST);
	if(qi->isSet(QueueItem::FLAG_RESUME))
		setFlag(Download::FLAG_RESUME);
};

void DownloadManager::onTimerSecond(u_int32_t /*aTick*/) {
	Lock l(cs);

	Download::List tickList;
	// Tick each ongoing download
	for(Download::Iter i = downloads.begin(); i != downloads.end(); ++i) {
		if((*i)->getTotal() > 0) {
			tickList.push_back(*i);
		}
	}

	if(tickList.size() > 0)
		fire(DownloadManagerListener::TICK, tickList);
}

void DownloadManager::FileMover::moveFile(const string& source, const string& target) {
	Lock l(cs);
	files.push_back(make_pair(source, target));
	if(!active) {
		active = true;
		start();
	}
}

int DownloadManager::FileMover::run() {
	for(;;) {
		FilePair next;
		{
			Lock l(cs);
			if(files.empty()) {
				active = false;
				return 0;
			}
			next = files.back();
			files.pop_back();
		}
		try {
			File::renameFile(next.first, next.second);
		} catch(const FileException&) {
			// Too bad...
		}
	}
}

void DownloadManager::removeConnection(UserConnection::Ptr aConn, bool reuse /* = false */) {
	dcassert(aConn->getDownload() == NULL);
	aConn->removeListener(this);
	ConnectionManager::getInstance()->putDownloadConnection(aConn, reuse);
}

void DownloadManager::checkDownloads(UserConnection* aConn) {

	if( ((SETTING(DOWNLOAD_SLOTS) != 0) && getDownloads() >= SETTING(DOWNLOAD_SLOTS)) ||
		((SETTING(MAX_DOWNLOAD_SPEED) != 0 && getAverageSpeed() >= (SETTING(MAX_DOWNLOAD_SPEED)*1024)) ) ) {
		
		if(!QueueManager::getInstance()->hasDownload(aConn->getUser(), QueueItem::HIGHEST)) {
			removeConnection(aConn);
			return;
		}
	}
	
	Download* d = QueueManager::getInstance()->getDownload(aConn->getUser());
	
	if(d) {
		dcassert(aConn->getDownload() == NULL);
		d->setUserConnection(aConn);
		aConn->setDownload(d);
		aConn->setState(UserConnection::STATE_FILELENGTH);

		{
			Lock l(cs);
			downloads.push_back(d);
		}
		
		if(d->isSet(Download::FLAG_RESUME)) {
			int64_t size = File::getSize(d->getTempTarget().empty() ? d->getTarget() : d->getTempTarget());
			int rollback = SETTING(ROLLBACK);
			int cutoff = max(SETTING(ROLLBACK), SETTING(BUFFER_SIZE)*1024);

			// dcassert(d->getSize() != -1);
			if( (rollback + cutoff) > min(size, d->getSize()) ) {
				d->setPos(0);
			} else {
				d->setPos(size - rollback - cutoff);
				d->setRollbackBuffer(rollback);
				d->setFlag(Download::FLAG_ROLLBACK);
			}
		} else {
			d->setPos(0);
		}
		if(d->isSet(Download::FLAG_USER_LIST) && aConn->isSet(UserConnection::FLAG_SUPPORTS_BZLIST)) {
			d->setSource("MyList.bz2");
		}

		if(BOOLSETTING(COMPRESS_TRANSFERS) && !d->isSet(Download::FLAG_USER_LIST) && 
			aConn->isSet(UserConnection::FLAG_SUPPORTS_GETZBLOCK)) {

			// This one, we'll download with a bzblock download instead...
			d->setFlag(Download::FLAG_ZDOWNLOAD);
			d->bytesLeft = d->getSize() - d->getPos();
			d->setComp(new ZDecompressor());
			aConn->getZBlock(d->getSource(), d->getPos(), d->bytesLeft);
		} else {
			aConn->get(d->getSource(), d->getPos());
		}
		return;
	}

	// Connection not needed any more, return it to the ConnectionManager...
	removeConnection(aConn, true);
}

void DownloadManager::onSending(UserConnection* aSource) {
	if(aSource->getState() != UserConnection::STATE_FILELENGTH) {
		dcdebug("DM::onFileLength Bad state, ignoring\n");
		return;
	}
	
	if(prepareFile(aSource)) {
		aSource->setDataMode();
	}
}

bool DownloadManager::prepareFile(UserConnection* aSource, int64_t newSize /* = -1 */) {
	Download* d = aSource->getDownload();
	dcassert(d != NULL);

	if(newSize != -1)
		d->setSize(newSize);

	dcassert(d->getSize() != -1);

	string target = d->getTempTarget().empty() ? d->getTarget() : d->getTempTarget();
	Util::ensureDirectory(target);
	if(d->isSet(Download::FLAG_USER_LIST) && aSource->isSet(UserConnection::FLAG_SUPPORTS_BZLIST)) {
		target.replace(target.size() - 5, 5, "bz2");
	}
	File* file;
	try {
		// Let's check if we can find this file in a any .SFV...
		int trunc = d->isSet(Download::FLAG_RESUME) ? 0 : File::TRUNCATE;
		bool sfvcheck = BOOLSETTING(SFV_CHECK) && (d->getPos() == 0) && (SFVReader(d->getTarget()).hasCRC());
		
		if(BOOLSETTING(ANTI_FRAG) && !d->isSet(Download::FLAG_USER_LIST)) {
			// Anti-frag file...First, remove any old attempt that might have existed
			// and rename any partial file alread downloaded...
			string atarget = target + ANTI_FRAG_EXT;
			try {
				File::deleteFile(atarget);
				File::renameFile(target, atarget);
			} catch(const FileException& e) {
				dcdebug("AntiFrag: %s\n", e.getError().c_str());
			}
			file = new SizedFile(d->getSize(), atarget, File::RW, File::OPEN | File::CREATE | trunc, sfvcheck);

			d->setFlag(Download::FLAG_ANTI_FRAG);
		} else {
			file = new BufferedFile(target, File::RW, File::OPEN | File::CREATE | trunc, sfvcheck);			
		}

		file->setPos(d->getPos());
		
	} catch(const FileException& e) {
		fire(DownloadManagerListener::FAILED, d, STRING(COULD_NOT_OPEN_TARGET_FILE) + e.getError());
		aSource->setDownload(NULL);
		removeDownload(d);
		removeConnection(aSource);
		return false;
	}

	dcassert(d->getPos() != -1);
	d->setFile(file);
	
	if(d->getSize() <= d->getPos()) {
		aSource->setDownload(NULL);
		removeDownload(d, true);
		removeConnection(aSource);
		return false;
	} else {
		d->setStart(GET_TICK());
		aSource->setState(UserConnection::STATE_DONE);
		
		fire(DownloadManagerListener::STARTING, d);
		
	}
	
	return true;
}	

void DownloadManager::onFileLength(UserConnection* aSource, const string& aFileLength) {

	if(aSource->getState() != UserConnection::STATE_FILELENGTH) {
		dcdebug("DM::onFileLength Bad state, ignoring\n");
		return;
	}

	int64_t fileLength = Util::toInt64(aFileLength);
	if(prepareFile(aSource, fileLength)) {
		aSource->setDataMode(aSource->getDownload()->getSize() - aSource->getDownload()->getPos());
		aSource->startSend();
	}
}

bool DownloadManager::checkRollback(Download* d, const u_int8_t* aData, int aLen) throw(FileException) {
	dcassert(d->getRollbackBuffer());
	
	if(d->getTotal() + aLen >= d->getRollbackSize()) {
		AutoArray<u_int8_t> buf(d->getRollbackSize());
		int len = d->getRollbackSize() - (int)d->getTotal();
		dcassert(len > 0);
		dcassert(len <= d->getRollbackSize());
		memcpy(d->getRollbackBuffer() + d->getTotal(), aData, len);
		
		d->getFile()->read((u_int8_t*)buf, d->getRollbackSize());
		
		int cmp = memcmp(d->getRollbackBuffer(), buf, d->getRollbackSize());
		
		d->unsetFlag(Download::FLAG_ROLLBACK);
		d->setRollbackBuffer(0);
		
		if(cmp != 0) {
			return false;
		}
		if(!d->isSet(Download::FLAG_ANTI_FRAG))
			d->getFile()->setEOF();
		// Write the rest...the file pointer should have been moved to the correct position by now...
		d->getFile()->write(aData+len, aLen - len);
	} else {
		memcpy(d->getRollbackBuffer() + d->getTotal(), aData, aLen);
	}
	
	return true;
}

void DownloadManager::onData(UserConnection* aSource, const u_int8_t* aData, int aLen) {
	Download* d = aSource->getDownload();
	dcassert(d != NULL);

	if(d->isSet(Download::FLAG_ZDOWNLOAD)) {
		// Oops, this is a bit more work...
		dcassert(d->getComp() != NULL);
		dcassert(d->bytesLeft > 0);
		int l = aLen;
		while(l > 0) {
			const u_int8_t* data = aData + (aLen - l);
			try {
				u_int32_t b = d->getComp()->decompress(data, l);
				if(!handleData(aSource, d->getComp()->getOutbuf(), b))
					break;
				d->bytesLeft -= b;
				if(d->bytesLeft == 0) {
					aSource->setLineMode();
					handleEndData(aSource);
					if(l != 0) {
						// Uhm, this client must be sending junk data after the compressed block...
						aSource->disconnect();
						return;
					}
				}
			} catch(const CryptoException&) {
				// Oops, decompression error...could happen for many reasons
				// but the most probable is that we received bad data...
				// We remove whatever we managed to download in this sessions
				// as it might be bad all of it...
				try {
					d->getFile()->movePos(-d->getTotal());
					d->getFile()->setEOF();
				} catch(const FileException&) {
					// Ignore...
				}

				fire(DownloadManagerListener::FAILED, d, STRING(DECOMPRESSION_ERROR));

				aSource->setDownload(NULL);
				removeDownload(d);
				removeConnection(aSource);
				return;
			}
		}
	} else {
		handleData(aSource, aData, aLen);
	}
}