diff --git a/include/NOD/DiscBase.hpp b/include/NOD/DiscBase.hpp index 1dafece..43e1537 100644 --- a/include/NOD/DiscBase.hpp +++ b/include/NOD/DiscBase.hpp @@ -356,13 +356,15 @@ protected: const SystemChar* m_outPath; std::unique_ptr m_fileIO; std::vector> m_partitions; + int64_t m_discCapacity; public: std::function m_progressCB; size_t m_progressIdx = 0; virtual ~DiscBuilderBase() {} - DiscBuilderBase(const SystemChar* outPath, + DiscBuilderBase(const SystemChar* outPath, int64_t discCapacity, std::function progressCB) - : m_fileIO(std::move(NewFileIO(outPath))), m_outPath(outPath), m_progressCB(progressCB) {} + : m_outPath(outPath), m_fileIO(std::move(NewFileIO(outPath, discCapacity))), + m_discCapacity(discCapacity), m_progressCB(progressCB) {} IFileIO& getFileIO() {return *m_fileIO;} }; diff --git a/include/NOD/IFileIO.hpp b/include/NOD/IFileIO.hpp index 143d58c..21ebd71 100644 --- a/include/NOD/IFileIO.hpp +++ b/include/NOD/IFileIO.hpp @@ -37,9 +37,8 @@ public: virtual std::unique_ptr beginReadStream(uint64_t offset) const=0; }; -std::unique_ptr NewFileIO(const SystemString& path); -std::unique_ptr NewFileIO(const SystemChar* path); -std::unique_ptr NewMemIO(void* buf, uint64_t size); +std::unique_ptr NewFileIO(const SystemString& path, int64_t maxWriteSize=-1); +std::unique_ptr NewFileIO(const SystemChar* path, int64_t maxWriteSize=-1); } diff --git a/lib/DiscGCN.cpp b/lib/DiscGCN.cpp index 1b47fc9..7e62779 100644 --- a/lib/DiscGCN.cpp +++ b/lib/DiscGCN.cpp @@ -249,7 +249,7 @@ bool DiscBuilderGCN::buildFromDirectory(const SystemChar* dirIn, const SystemCha DiscBuilderGCN::DiscBuilderGCN(const SystemChar* outPath, const char gameID[6], const char* gameTitle, uint32_t fstMemoryAddr, std::function progressCB) -: DiscBuilderBase(outPath, progressCB) +: DiscBuilderBase(outPath, 0x57058000, progressCB) { PartitionBuilderGCN* partBuilder = new PartitionBuilderGCN(*this, PartitionBuilderBase::Kind::Data, gameID, gameTitle, fstMemoryAddr); diff --git a/lib/DiscWii.cpp b/lib/DiscWii.cpp index 43e80a5..6de2490 100644 --- a/lib/DiscWii.cpp +++ b/lib/DiscWii.cpp @@ -850,24 +850,22 @@ public: bool DiscBuilderWii::buildFromDirectory(const SystemChar* dirIn, const SystemChar* dolIn, const SystemChar* apploaderIn, const SystemChar* partHeadIn) { - size_t DISC_CAPACITY = m_dualLayer ? 0x1FB4E0000 : 0x118240000; - PartitionBuilderWii& pb = static_cast(*m_partitions[0]); uint64_t filledSz = pb.m_baseOffset; m_fileIO->beginWriteStream(); - if (!CheckFreeSpace(m_outPath, DISC_CAPACITY)) + if (!CheckFreeSpace(m_outPath, m_discCapacity)) { LogModule.report(LogVisor::Error, _S("not enough free disk space for %s"), m_outPath); return false; } ++m_progressIdx; m_progressCB(m_progressIdx, _S("Preallocating image"), -1); - m_fileIO->beginWriteStream(DISC_CAPACITY - 1)->write("", 1); + m_fileIO->beginWriteStream(m_discCapacity - 1)->write("", 1); /* Assemble image */ filledSz = pb.buildFromDirectory(dirIn, dolIn, apploaderIn, partHeadIn); - if (filledSz >= DISC_CAPACITY) + if (filledSz >= m_discCapacity) { LogModule.report(LogVisor::FatalError, "data partition exceeds disc capacity"); return false; @@ -910,7 +908,7 @@ bool DiscBuilderWii::buildFromDirectory(const SystemChar* dirIn, const SystemCha ws = m_fileIO->beginWriteStream(filledSz); uint8_t fillBuf[512]; memset(fillBuf, 0xff, 512); - for (size_t i=DISC_CAPACITY-filledSz ; i>0 ;) + for (size_t i=m_discCapacity-filledSz ; i>0 ;) { if (i >= 512) { @@ -927,7 +925,7 @@ bool DiscBuilderWii::buildFromDirectory(const SystemChar* dirIn, const SystemCha DiscBuilderWii::DiscBuilderWii(const SystemChar* outPath, const char gameID[6], const char* gameTitle, bool dualLayer, std::function progressCB) -: DiscBuilderBase(outPath, progressCB), m_dualLayer(dualLayer) +: DiscBuilderBase(outPath, dualLayer ? 0x1FB4E0000 : 0x118240000, progressCB), m_dualLayer(dualLayer) { PartitionBuilderWii* partBuilder = new PartitionBuilderWii(*this, PartitionBuilderBase::Kind::Data, gameID, gameTitle, 0x200000); diff --git a/lib/FileIOFILE.cpp b/lib/FileIOFILE.cpp index 77f56cd..e3a4531 100644 --- a/lib/FileIOFILE.cpp +++ b/lib/FileIOFILE.cpp @@ -1,5 +1,6 @@ #include #include +#include #include "NOD/Util.hpp" #include "NOD/IFileIO.hpp" @@ -9,11 +10,12 @@ namespace NOD class FileIOFILE : public IFileIO { SystemString m_path; + int64_t m_maxWriteSize; public: - FileIOFILE(const SystemString& path) - : m_path(path) {} - FileIOFILE(const SystemChar* path) - : m_path(path) {} + FileIOFILE(const SystemString& path, int64_t maxWriteSize) + : m_path(path), m_maxWriteSize(maxWriteSize) {} + FileIOFILE(const SystemChar* path, int64_t maxWriteSize) + : m_path(path), m_maxWriteSize(maxWriteSize) {} bool exists() { @@ -38,13 +40,16 @@ public: struct WriteStream : public IFileIO::IWriteStream { FILE* fp; - WriteStream(const SystemString& path) + int64_t m_maxWriteSize; + WriteStream(const SystemString& path, int64_t maxWriteSize) + : m_maxWriteSize(maxWriteSize) { fp = fopen(path.c_str(), "wb"); if (!fp) LogModule.report(LogVisor::FatalError, _S("unable to open '%s' for writing"), path.c_str()); } - WriteStream(const SystemString& path, uint64_t offset) + WriteStream(const SystemString& path, uint64_t offset, int64_t maxWriteSize) + : m_maxWriteSize(maxWriteSize) { fp = fopen(path.c_str(), "ab"); if (!fp) @@ -64,6 +69,8 @@ public: } uint64_t write(const void* buf, uint64_t length) { + if (FTell(fp) + length > m_maxWriteSize) + LogModule.report(LogVisor::FatalError, _S("write operation exceeds file's %" PRIi64 "-byte limit"), m_maxWriteSize); return fwrite(buf, 1, length, fp); } uint64_t copyFromDisc(IPartReadStream& discio, uint64_t length) @@ -92,11 +99,11 @@ public: }; std::unique_ptr beginWriteStream() const { - return std::unique_ptr(new WriteStream(m_path)); + return std::unique_ptr(new WriteStream(m_path, m_maxWriteSize)); } std::unique_ptr beginWriteStream(uint64_t offset) const { - return std::unique_ptr(new WriteStream(m_path, offset)); + return std::unique_ptr(new WriteStream(m_path, offset, m_maxWriteSize)); } struct ReadStream : public IFileIO::IReadStream @@ -162,14 +169,14 @@ public: } }; -std::unique_ptr NewFileIO(const SystemString& path) +std::unique_ptr NewFileIO(const SystemString& path, int64_t maxWriteSize) { - return std::unique_ptr(new FileIOFILE(path)); + return std::unique_ptr(new FileIOFILE(path, maxWriteSize)); } -std::unique_ptr NewFileIO(const SystemChar* path) +std::unique_ptr NewFileIO(const SystemChar* path, int64_t maxWriteSize) { - return std::unique_ptr(new FileIOFILE(path)); + return std::unique_ptr(new FileIOFILE(path, maxWriteSize)); } } diff --git a/lib/FileIOWin32.cpp b/lib/FileIOWin32.cpp index 0fa6f02..e140fde 100644 --- a/lib/FileIOWin32.cpp +++ b/lib/FileIOWin32.cpp @@ -1,5 +1,6 @@ #include #include +#include #include "NOD/Util.hpp" #include "NOD/IFileIO.hpp" @@ -9,11 +10,12 @@ namespace NOD class FileIOWin32 : public IFileIO { SystemString m_path; + int64_t m_maxWriteSize; public: - FileIOWin32(const SystemString& path) - : m_path(path) {} - FileIOWin32(const SystemChar* path) - : m_path(path) {} + FileIOWin32(const SystemString& path, int64_t maxWriteSize) + : m_path(path), m_maxWriteSize(maxWriteSize) {} + FileIOWin32(const SystemChar* path, int64_t maxWriteSize) + : m_path(path), m_maxWriteSize(maxWriteSize) {} bool exists() { @@ -44,14 +46,17 @@ public: struct WriteStream : public IFileIO::IWriteStream { HANDLE fp; - WriteStream(const SystemString& path) + int64_t m_maxWriteSize; + WriteStream(const SystemString& path, int64_t maxWriteSize) + : m_maxWriteSize(maxWriteSize) { fp = CreateFileW(path.c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, nullptr, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); if (fp == INVALID_HANDLE_VALUE) LogModule.report(LogVisor::FatalError, _S("unable to open '%s' for writing"), path.c_str()); } - WriteStream(const SystemString& path, uint64_t offset) + WriteStream(const SystemString& path, uint64_t offset, int64_t maxWriteSize) + : m_maxWriteSize(maxWriteSize) { fp = CreateFileW(path.c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, nullptr, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); @@ -67,6 +72,8 @@ public: } uint64_t write(const void* buf, uint64_t length) { + if (FTell(fp) + length > m_maxWriteSize) + LogModule.report(LogVisor::FatalError, _S("write operation exceeds file's %" PRIi64 "-byte limit"), m_maxWriteSize); DWORD ret = 0; WriteFile(fp, buf, length, &ret, nullptr); return ret; @@ -97,11 +104,11 @@ public: }; std::unique_ptr beginWriteStream() const { - return std::unique_ptr(new WriteStream(m_path)); + return std::unique_ptr(new WriteStream(m_path, m_maxWriteSize)); } std::unique_ptr beginWriteStream(uint64_t offset) const { - return std::unique_ptr(new WriteStream(m_path, offset)); + return std::unique_ptr(new WriteStream(m_path, offset, m_maxWriteSize)); } struct ReadStream : public IFileIO::IReadStream