Skip to content

Commit

Permalink
Do more checking when opening vpks
Browse files Browse the repository at this point in the history
  • Loading branch information
@SCell555
SCell555 committed Feb 4, 2021
1 parent 7458cea commit f3d4530
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 75 deletions.
2 changes: 1 addition & 1 deletion 7zTypes.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ typedef int WRes;


#ifndef RINOK
#define RINOK(x) { int __result__ = (x); if (__result__ != 0) return __result__; }
#define RINOK(x) if (int __result__ = (x); __result__ != 0) return __result__
#endif

typedef unsigned char Byte;
Expand Down
127 changes: 61 additions & 66 deletions Handler.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ class CHandler :
CMyComPtr<IInStream> basePak;
CObjectVector<CMyComPtr<IInStream>> paks;
UInt64 size = 0;
bool missingPak = false;
};

STDMETHODIMP CHandler::Open( IInStream* inStream, const UInt64* maxCheckStartPosition, IArchiveOpenCallback* callback ) MY_NO_THROW_DECL_ONLY
Expand All @@ -60,16 +59,16 @@ STDMETHODIMP CHandler::Open( IInStream* inStream, const UInt64* maxCheckStartPos
if ( largestId != -1 )
{
CMyComPtr<IArchiveOpenVolumeCallback> volCallback;
callback->QueryInterface( IID_IArchiveOpenVolumeCallback, (void**)&volCallback );
RINOK( callback->QueryInterface( IID_IArchiveOpenVolumeCallback, (void**)&volCallback ) );
if ( !volCallback )
return S_OK;
return E_FAIL;

UString name;
{
NWindows::NCOM::CPropVariant prop;
RINOK( volCallback->GetProperty( kpidName, &prop ) );
if ( prop.vt != VT_BSTR )
return S_OK;
return E_FAIL;
name = prop.bstrVal;
}

Expand All @@ -85,8 +84,8 @@ STDMETHODIMP CHandler::Open( IInStream* inStream, const UInt64* maxCheckStartPos
swprintf_s( base, L"%s%03d.vpk", name.GetBuf(), i );

IInStream* s;
if ( volCallback->GetStream( base, &s ) == S_FALSE || s == nullptr )
missingPak = true;
if ( volCallback->GetStream( base, &s ) != S_OK || s == nullptr )
return ERROR_FILE_NOT_FOUND;
paks.Add( s );
}
}
Expand Down Expand Up @@ -225,31 +224,26 @@ STDMETHODIMP CHandler::Extract( const UInt32* indices, UInt32 numItems, Int32 te
totalSize += f[indices[i]].second.fileLength;
}

extractCallback->SetTotal( totalSize );

CLocalProgress* lps = new CLocalProgress;
CMyComPtr<ICompressProgressInfo> progress = lps;
lps->Init( extractCallback, false );
RINOK( extractCallback->SetTotal( totalSize ) );

NCompress::CCopyCoder* copyCoderSpec = new NCompress::CCopyCoder();
CMyComPtr<ICompressCoder> copyCoder = copyCoderSpec;
CMyComPtr<CLocalProgress> progress = new CLocalProgress;
progress->Init( extractCallback, false );

CLimitedSequentialOutStream* outStreamSpec = new CLimitedSequentialOutStream;
CMyComPtr<ISequentialOutStream> outStream( outStreamSpec );
CMyComPtr<ICompressCoder> copyCoder = new NCompress::CCopyCoder;
CMyComPtr<CLimitedSequentialOutStream> outStream = new CLimitedSequentialOutStream;

const Int32 askMode = testMode ? NArchive::NExtract::NAskMode::kTest : NArchive::NExtract::NAskMode::kExtract;
UInt64 currentTotalSize = 0;
for ( UInt32 i = 0; i < numItems; i++ )
{
lps->InSize = lps->OutSize = currentTotalSize;
RINOK( lps->SetCur() );
progress->InSize = progress->OutSize = currentTotalSize;
RINOK( progress->SetCur() );
CMyComPtr<ISequentialOutStream> realOutStream;
Int32 askMode = testMode ? NArchive::NExtract::NAskMode::kTest : NArchive::NExtract::NAskMode::kExtract;
UInt32 index = allFilesMode ? i : indices[i];

RINOK( extractCallback->GetStream( index, &realOutStream, askMode ) );

const auto& item = f[index].second;

currentTotalSize += item.fileLength;

if ( !testMode && !realOutStream )
Expand All @@ -264,9 +258,9 @@ STDMETHODIMP CHandler::Extract( const UInt32* indices, UInt32 numItems, Int32 te
}

RINOK( extractCallback->PrepareOperation( askMode ) );
outStreamSpec->SetStream( realOutStream );
outStream->SetStream( realOutStream );
realOutStream.Release();
outStreamSpec->Init( item.fileLength );
outStream->Init( item.fileLength );
Int32 opRes;
CMyComPtr<ISequentialInStream> inStream;
HRESULT res = GetStream( index, &inStream );
Expand All @@ -277,9 +271,9 @@ STDMETHODIMP CHandler::Extract( const UInt32* indices, UInt32 numItems, Int32 te
else
{
RINOK( copyCoder->Code( inStream, outStream, NULL, NULL, progress ) );
opRes = outStreamSpec->IsFinishedOK() ? NArchive::NExtract::NOperationResult::kOK : NArchive::NExtract::NOperationResult::kDataError;
opRes = outStream->IsFinishedOK() ? NArchive::NExtract::NOperationResult::kOK : NArchive::NExtract::NOperationResult::kDataError;
}
outStreamSpec->ReleaseStream();
outStream->ReleaseStream();
RINOK( extractCallback->SetOperationResult( opRes ) );
}
return S_OK;
Expand All @@ -303,19 +297,18 @@ STDMETHODIMP CHandler::GetStream( UInt32 index, ISequentialInStream** stream )
if ( i.preloadLength == 0 )
return CreateLimitedInStream( i.archiveIdx == 0x7FFF ? basePak : paks[i.archiveIdx], offset, i.fileLength, stream );

CLimitedCachedInStream* limitedStreamSpec = new CLimitedCachedInStream;
CMyComPtr<IInStream> limitedStream = limitedStreamSpec;
limitedStreamSpec->SetStream( i.archiveIdx == 0x7FFF ? basePak : paks[i.archiveIdx], offset );
CMyComPtr<CLimitedCachedInStream> limitedStream = new CLimitedCachedInStream;
limitedStream->SetStream( i.archiveIdx == 0x7FFF ? basePak : paks[i.archiveIdx], offset );

auto& preload = limitedStreamSpec->Buffer;
auto& preload = limitedStream->Buffer;
preload.Alloc( i.preloadLength );

size_t size = i.preloadLength;
basePak->Seek( i.preloadOffset, STREAM_SEEK_SET, nullptr );
ReadStream( basePak, preload, &size );
RINOK( basePak->Seek( i.preloadOffset, STREAM_SEEK_SET, nullptr ) );
RINOK( ReadStream( basePak, preload, &size ) );

limitedStreamSpec->SetCache( i.preloadLength, 0 );
limitedStreamSpec->InitAndSeek( 0, i.fileLength );
limitedStream->SetCache( i.preloadLength, 0 );
RINOK( limitedStream->InitAndSeek( 0, i.fileLength ) );
*stream = limitedStream.Detach();
return S_OK;
}
Expand Down Expand Up @@ -753,21 +746,21 @@ class VpkWriter
return S_OK;
}

void write_pak( ISequentialOutStream* outStream, IArchiveUpdateCallback* callback )
HRESULT write_pak( ISequentialOutStream* outStream, IArchiveUpdateCallback* callback )
{
CLocalProgress* lps = new CLocalProgress;
CMyComPtr<ICompressProgressInfo> progress = lps;
lps->Init( callback, true );

NCompress::CCopyCoder* copyCoderSpec = new NCompress::CCopyCoder();
CMyComPtr<ICompressCoder> copyCoder = copyCoderSpec;
CMyComPtr<ICompressCoder> copyCoder = new NCompress::CCopyCoder;

CMyComPtr<IOutStream> stream_;
outStream->QueryInterface( IID_IOutStream, (void**)&stream_ );
RINOK( outStream->QueryInterface( IID_IOutStream, (void**)&stream_ ) );

CCacheOutStream* stream = new CCacheOutStream();
stream->Allocate();
stream->Init( outStream, stream_ );
if ( !stream->Allocate() )
return E_OUTOFMEMORY;
RINOK( stream->Init( outStream, stream_ ) );

UInt32 size = 0, treeSize = 0;
FilesByExt sorted_files;
Expand All @@ -777,38 +770,38 @@ class VpkWriter
treeSize += static_cast<UInt32>( dirs.size() + 1 ); // add null after each all files in each directory + null after last dir
++treeSize; // null after last ext

write_header( stream, size, treeSize );
RINOK( write_header( stream, size, treeSize ) );

UInt32 currentOffset = 0;
for ( auto& [ext, dirs] : sorted_files )
{
write( stream, ext );
RINOK( write( stream, ext ) );
for ( auto& [dir, files] : dirs )
{
write( stream, dir );
RINOK( write( stream, dir ) );
for ( auto& [file, data] : files )
{
UInt64 pos;
stream->Seek( 0, STREAM_SEEK_CUR, &pos );
RINOK( stream->Seek( 0, STREAM_SEEK_CUR, &pos ) );
lps->InSize = lps->OutSize = pos;
lps->SetCur();
write( stream, file );
write<UInt32>( stream, calc_crc( *data ) );
write<UInt16>( stream, 0 );
write<UInt16>( stream, 0x7FFF );
write<UInt32>( stream, currentOffset );
write<UInt32>( stream, data->size );
write<UInt16>( stream, 0xFFFF );
RINOK( lps->SetCur() );
RINOK( write( stream, file ) );
RINOK( write<UInt32>( stream, calc_crc( *data ) ) );
RINOK( write<UInt16>( stream, 0 ) );
RINOK( write<UInt16>( stream, 0x7FFF ) );
RINOK( write<UInt32>( stream, currentOffset ) );
RINOK( write<UInt32>( stream, data->size ) );
RINOK( write<UInt16>( stream, 0xFFFF ) );
currentOffset += data->size;

}
write<std::string>( stream, {} );
RINOK( write<std::string>( stream, {} ) );
}
write<std::string>( stream, {} );
RINOK( write<std::string>( stream, {} ) );
}
write<std::string>( stream, {} );
RINOK( write<std::string>( stream, {} ) );

stream->Seek( treeSize + sizeof( libvpk::meta::VPKHeader ), STREAM_SEEK_SET, nullptr );
RINOK( stream->Seek( treeSize + sizeof( libvpk::meta::VPKHeader ), STREAM_SEEK_SET, nullptr ) );

for ( auto& [ext, dirs] : sorted_files )
{
Expand All @@ -817,41 +810,43 @@ class VpkWriter
for ( auto& [file, data] : files )
{
UInt64 pos;
stream->Seek( 0, STREAM_SEEK_CUR, &pos );
RINOK( stream->Seek( 0, STREAM_SEEK_CUR, &pos ) );
lps->InSize = lps->OutSize = pos;
lps->SetCur();
copyCoder->Code( data->stream, stream, NULL, NULL, progress );
RINOK( lps->SetCur() );
RINOK( copyCoder->Code( data->stream, stream, NULL, NULL, progress ) );
data->stream->Release();
}
}
}

stream->FlushCache();
RINOK( stream->FlushCache() );
stream->Release();

return S_OK;
}

private:
template <typename T>
void write( IOutStream* stream, const T& data )
HRESULT write( IOutStream* stream, const T& data )
{
stream->Write( &data, sizeof( T ), nullptr );
return stream->Write( &data, sizeof( T ), nullptr );
}

template <>
void write<std::string>( IOutStream* stream, const std::string& string )
HRESULT write<std::string>( IOutStream* stream, const std::string& string )
{
stream->Write( string.c_str(), static_cast<UInt32>( string.size() + 1 ), nullptr );
return stream->Write( string.c_str(), static_cast<UInt32>( string.size() + 1 ), nullptr );
}

static void write_header( IOutStream* stream, UInt32 size, UInt32 treeSize )
static HRESULT write_header( IOutStream* stream, UInt32 size, UInt32 treeSize )
{
libvpk::meta::VPKHeader header;
header.signature = libvpk::meta::VPKHeader::ValidSignature;
header.version = 2;
header.treeSize = static_cast<Int32>( treeSize );
header.fileDataSectionSize = static_cast<Int32>( size );

stream->Write( &header, sizeof( header ), nullptr );
return stream->Write( &header, sizeof( header ), nullptr );
}

struct Dir
Expand Down Expand Up @@ -970,19 +965,19 @@ STDMETHODIMP CHandler::UpdateItems( ISequentialOutStream* outStream, UInt32 numI
UInt64 size = prop.uhVal.QuadPart;

CMyComPtr<ISequentialInStream> fileInStream;
callback->GetStream( i, &fileInStream );
RINOK( callback->GetStream( i, &fileInStream ) );
if ( !fileInStream )
return E_FAIL;

IInStream* inStream;
fileInStream->QueryInterface( IID_IInStream, (void**)&inStream );
RINOK( fileInStream->QueryInterface( IID_IInStream, (void**)&inStream ) );
if ( !inStream )
return E_FAIL;

RINOK( writer.addItem( UnicodeStringToMultiByte( name ), static_cast<UInt32>( size ), inStream ) );
}

writer.write_pak( outStream, callback );
RINOK( writer.write_pak( outStream, callback ) );

RINOK( callback->SetOperationResult( NArchive::NExtract::NOperationResult::kOK ) );

Expand Down
20 changes: 12 additions & 8 deletions libvpk++.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,20 +88,24 @@ namespace libvpk

HRESULT open( IInStream* stream, IArchiveOpenCallback* callback )
{
UInt64 pos;
RINOK( stream->Seek( 0, STREAM_SEEK_CUR, &pos ) );
auto initialHeader = helpers::read<meta::VPKHeader1>( stream );
if ( initialHeader.signature != meta::VPKHeader1::ValidSignature )
return E_FAIL;
return S_FALSE;

if ( initialHeader.version == 1 )
m_header = initialHeader;
else if ( initialHeader.version == 2 )
{
// Return to the beginning and read a full VPK 2 ptr.
stream->Seek( 0, STREAM_SEEK_SET, nullptr );
RINOK( stream->Seek( pos, STREAM_SEEK_SET, nullptr ) );
m_header = helpers::read<meta::VPKHeader2>( stream );
if ( m_header.signature != meta::VPKHeader1::ValidSignature )
return S_FALSE;
}
else
return E_FAIL;
return S_FALSE;

const UInt64 total = m_header.treeSize;
callback->SetTotal( nullptr, &total );
Expand Down Expand Up @@ -145,8 +149,8 @@ namespace libvpk
{
UInt64 pos;
const UInt64 files = m_files.size();
stream->Seek( 0, STREAM_SEEK_CUR, &pos );
callback->SetCompleted( &files, &pos );
RINOK( stream->Seek( 0, STREAM_SEEK_CUR, &pos ) );
RINOK( callback->SetCompleted( &files, &pos ) );
}

auto name = helpers::read<std::string>( stream );
Expand Down Expand Up @@ -181,16 +185,16 @@ namespace libvpk

// Read the terminator for the file info.
if ( helpers::read<uint16_t>( stream ) != 0xFFFF )
return E_FAIL;
return S_FALSE;

UInt64 pos;
stream->Seek( 0, STREAM_SEEK_CUR, &pos );
RINOK( stream->Seek( 0, STREAM_SEEK_CUR, &pos ) );
VPKFileDesc desc{ archiveIndex, preloadBytes, static_cast<uint32_t>( pos ), offset, length, crc };

// Skip over the preload section
if ( desc.preloadLength != 0 )
{
stream->Seek( desc.preloadLength, STREAM_SEEK_CUR, nullptr );
RINOK( stream->Seek( desc.preloadLength, STREAM_SEEK_CUR, nullptr ) );
desc.fileLength += desc.preloadLength;
}

Expand Down

0 comments on commit f3d4530

Please sign in to comment.