storage.cpp

这是一个嵌入式linux系统下的BT下载工具包

/***************************************************************************
 *   Copyright (C) 2005-2006 Gao Xianchao                                  *
 *                 2007 Gao Xianchao gnap_an linux_lyb ahlongxp            *
 *                                                                         *
 *   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.             *
 ***************************************************************************/

/*
 * Author:	gxc
 * Create data:	2005-10-16 17:09
 */
 
#include <assert.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/unistd.h>
#include <fcntl.h>
#include "Storage.h"
#include "log.h"
#include "utils.h"
#include "sha1.h"

CStorage::CStorage()
{
}

CStorage::~CStorage()
{
}

void CStorage::setBTTask(IBTTask* task)
{
	_task = task;
}

IBTTask* CStorage::getBTTask()
{
	return _task;
}

bool CStorage::openFiles()
{
	unsigned int count = _task->getTorrentFile()->getFileCount();
	for(unsigned int i=0; i<count; ++i)
	{
		TFileInfo fileInfo = _task->getTorrentFile()->getFileInfo(i);
		
		if(!openFile(i, fileInfo))
		{
			return false;
		}
	}	
	
	return true;
}

void CStorage::saveBitset()
{	
	FILE* stream = fopen(_psfPath.c_str(), "w");
	if(stream == NULL)
	{
		return;
	}
	
	unsigned int count = _task->getTorrentFile()->getPieceCount();
	for(unsigned int i=0; i<count; ++i)
	{
		char c;
		if(_bitSet.isSet(i))
		{
			c = '1';
		}
		else
		{
			c = '0';
		}
		
		fputc(c, stream);			
	}
	fclose(stream);		
}

void CStorage::checkFiles()
{
	FILE* stream = fopen(_psfPath.c_str(), "w");
	if(stream == NULL)
	{
		return;
	}
	
	unsigned int count = _task->getTorrentFile()->getPieceCount();
	for(unsigned int i=0; i<count; ++i)
	{
		std::string piece = readPiece(i);
		
		char infoHash[20];
		sha1_block((unsigned char*)piece.c_str(), piece.length(), (unsigned char*)infoHash);
		std::string hash;
		hash.append(infoHash, 20);
		
		char c;
		if(hash == _task->getTorrentFile()->getPieceHash(i))
		{
			_bitSet.set(i, true);
			c = '1';
		}
		else
		{
			_bitSet.set(i, false);
			c = '0';
		}
		
		fputc(c, stream);			
	}
	fclose(stream);
}

bool CStorage::needCheckFiles()
{
	bool recheck = true;
	
	int psf = open(_psfPath.c_str(), O_RDWR | O_CREAT, S_IRUSR | S_IWUSR | S_IRWXG | S_IRWXO);
	if(psf != -1)
	{		
		struct stat st;
		fstat(psf, &st);
		
		if(st.st_size == 0)
		{
			ftruncate(psf, 1);
		}
		
		bool fileModified = false;
		
		TFileHandleList::iterator iter = _fileHandleList.begin();
		for(; iter!=_fileHandleList.end(); ++iter)
		{
			if(iter->baned)
			{
				continue;
			}
			
			if(difftime(iter->mtime, st.st_mtime) > 0)
			{
				fileModified = true;
				break;
			}	
		}
	
		if(!fileModified)
		{
			// load birset from .psf file
			FILE* stream = fdopen(psf, "r");
			unsigned int index = 0;
			for(;;)
			{
				int ret = fgetc(stream);
				if(ret == EOF)
				{
					break;
				}
				
				if((char)ret == '1')
				{
					_bitSet.set(index, true);
				}
				else
				{
					_bitSet.set(index, false);
				}
				
				if(index == _task->getTorrentFile()->getPieceCount() -1)
				{
					recheck = false;
					break;
				}
				
				index++;
			}
			fclose(stream);
		}	
		close(psf);
	}
	
	return recheck;
}

void CStorage::makeRequestRange()
{
	//LOG_DEBUG("piece count = "<<_task->getTorrentFile()->getPieceCount());
	unsigned int nextBegin = 0;
	bool rangeListDoen = false;
	while(!rangeListDoen)
	{
		TIntPair pair;
		pair.begin = nextBegin;
		pair.end = nextBegin + _task->getTorrentFile()->getPieceCount()/4;
		
		if(pair.end == 0
			|| pair.end >= _task->getTorrentFile()->getPieceCount())
		{
			pair.end = _task->getTorrentFile()->getPieceCount()-1;
			rangeListDoen = true;
		}
		
		nextBegin = pair.end+1;
		//LOG_DEBUG("make range "<<pair.begin<<"-"<<pair.end);
		_rangeList.push_back(pair);
	}	
}

bool CStorage::start()
{
	LOG_DEBUG("CStorage starting, tick = "<<GetTickCount());
	
	_destDir = _task->getDestPath();
	
	if(_task->getTorrentFile()->IsUTF8Valid())
	{
		std::string path = _destDir + _task->getTorrentFile()->getTopDirUTF8();
		_psfPath = filename_from_utf8(path.c_str()) + "/" + filename_from_utf8(_task->getTaskName().c_str()) + ".psf";
	}
	else
	{
		_psfPath =  filename_from_utf8(_destDir.c_str()) + _task->getTorrentFile()->getTopDir() + "/" +_task->getTaskName() + ".psf";
	}
	_bitSet.alloc(_task->getTorrentFile()->getPieceCount());
	_newTask = true;
	_banedSize = 0;
	
	makeRequestRange();
	
	if(!openFiles())
	{
		return false;
	}
	
	genBanedBitSet();
	
	if(!_newTask && needCheckFiles())
	{
		checkFiles();
	}
	
	_finishedPiece = _bitSet.getSetedCount();

	LOG_DEBUG("CStorage started, tick = "<<GetTickCount());
	
	return true;
}

void CStorage::stop()
{
	saveWriteCacheToDisk();
	
	TFileHandleList::iterator iter = _fileHandleList.begin();
	for(; iter!=_fileHandleList.end(); ++iter)
	{
		if(iter->baned)
		{
			continue;
		}
		
		close(iter->handle);
	}
	
	if(finished())
	{
		unlink(_psfPath.c_str());
	}
}

void CStorage::setBanedFileList(TBanedFileList bandFileList)
{
	_bandFileList = bandFileList;
}

bool CStorage::finished()
{
	return _finishedPiece == (_bitSet.getSize() - _banedBitSet.getSetedCount());
}

std::string CStorage::getBitfield()
{
	if(_bitSet.isEmpty())
	{
		return "";
	}
	
	return _bitSet.getStream();
}

IBitSet* CStorage::getBitSet()
{
	return &_bitSet;
}

IBitSet* CStorage::getBanedBitSet()
{
	return &_banedBitSet;
}

unsigned int CStorage::getPieceLength(unsigned int pieceIndex)
{
	if(pieceIndex < _task->getTorrentFile()->getPieceCount()-1)
	{
		return _task->getTorrentFile()->getPieceLength();
	}
	
	return (unsigned int)(_task->getTorrentFile()->getTotalSize() - ((int64_t)_task->getTorrentFile()->getPieceLength())*pieceIndex);
}

unsigned int CStorage::getPieceTask(IBitSet* bitSet)
{	
	for(size_t i=0; i<_rangeList.size(); ++i)
	{
		TIntPair range = _rangeList.front();
		_rangeList.pop_front();
		_rangeList.push_back(range);		
		
		unsigned int index = getPieceTaskInRange(bitSet, range.begin, range.end);
		if(index != NONE_PIECE_INDEX)
		{
			return index;
		}
	}
	
	TIntList::iterator iter = _downloadingPieces.begin();
	for(; iter!=_downloadingPieces.end(); ++iter)
	{
		if(bitSet->isSet(*iter))
		{
			unsigned index = *iter;
			_downloadingPieces.erase(iter);
			_downloadingPieces.push_back(index);
			return index;
		}
	}
	
	return NONE_PIECE_INDEX;	
}


unsigned int CStorage::getPieceTaskInRange(IBitSet* bitSet, unsigned int beginIndex, unsigned int endIndex)
{
	//LOG_DEBUG("getPieceTaskInRange begin="<<beginIndex<<" end="<<endIndex);
	for(unsigned int i=beginIndex; i<=endIndex; ++i)
	{
		if(bitSet->isSet(i) 
			&& !_bitSet.isSet(i)
			&& !_banedBitSet.isSet(i))
		{
			TIntList::iterator iter = _downloadingPieces.begin();
			for(; iter!=_downloadingPieces.end(); ++iter)
			{
				if(*iter == i)
				{
					break;
				}
			}
			if(iter != _downloadingPieces.end())
			{
				//LOG_DEBUG("already in _downloadingPieces, index="<<i)
				continue;
			}
			
			//LOG_DEBUG("_downloadingPieces.push_back, index="<<i)
			_downloadingPieces.push_back(i);
			return i;		
		}
	}
	
	return NONE_PIECE_INDEX;	
}

void CStorage::removeDownloadingPiece(unsigned int index)
{
	TIntList::iterator iter = _downloadingPieces.begin();
	for(; iter!=_downloadingPieces.end(); ++iter)
	{
		if(*iter == index)
		{
			_downloadingPieces.erase(iter);
			break;
		}
	}	
}

void CStorage::abandonPieceTask(unsigned int pieceIndex)
{
	removeDownloadingPiece(pieceIndex);
}

void CStorage::writePiece(unsigned int pieceIndex,std::string& data)
{
	removeDownloadingPiece(pieceIndex);