File-size limiting

include/NOD/DiscBase.hpp

@@ -356,13 +356,15 @@ protected:
     const SystemChar* m_outPath;
     std::unique_ptr<IFileIO> m_fileIO;
     std::vector<std::unique_ptr<PartitionBuilderBase>> m_partitions;
+    int64_t m_discCapacity;
 public:
     std::function<void(size_t idx, const SystemString&, size_t)> m_progressCB;
     size_t m_progressIdx = 0;
     virtual ~DiscBuilderBase() {}
-    DiscBuilderBase(const SystemChar* outPath,
+    DiscBuilderBase(const SystemChar* outPath, int64_t discCapacity,
                     std::function<void(size_t idx, const SystemString&, size_t)> 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;}
 };

include/NOD/IFileIO.hpp

@@ -37,9 +37,8 @@ public:
     virtual std::unique_ptr<IReadStream> beginReadStream(uint64_t offset) const=0;
 };
 
-std::unique_ptr<IFileIO> NewFileIO(const SystemString& path);
+std::unique_ptr<IFileIO> NewFileIO(const SystemString& path, int64_t maxWriteSize=-1);
-std::unique_ptr<IFileIO> NewFileIO(const SystemChar* path);
+std::unique_ptr<IFileIO> NewFileIO(const SystemChar* path, int64_t maxWriteSize=-1);
-std::unique_ptr<IFileIO> NewMemIO(void* buf, uint64_t size);
 
 }
 

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<void(size_t, const SystemString&, size_t)> progressCB)
-: DiscBuilderBase(outPath, progressCB)
+: DiscBuilderBase(outPath, 0x57058000, progressCB)
 {
     PartitionBuilderGCN* partBuilder = new PartitionBuilderGCN(*this, PartitionBuilderBase::Kind::Data,
                                                                gameID, gameTitle, fstMemoryAddr);

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<PartitionBuilderWii&>(*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<void(size_t, const SystemString&, size_t)> 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);

lib/FileIOFILE.cpp

@@ -1,5 +1,6 @@
 #include <stdio.h>
 #include <stdlib.h>
+#include <inttypes.h>
 #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)
+    FileIOFILE(const SystemString& path, int64_t maxWriteSize)
-    : m_path(path) {}
+    : m_path(path), m_maxWriteSize(maxWriteSize) {}
-    FileIOFILE(const SystemChar* path)
+    FileIOFILE(const SystemChar* path, int64_t maxWriteSize)
-    : m_path(path) {}
+    : 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<IWriteStream> beginWriteStream() const
     {
-        return std::unique_ptr<IWriteStream>(new WriteStream(m_path));
+        return std::unique_ptr<IWriteStream>(new WriteStream(m_path, m_maxWriteSize));
     }
     std::unique_ptr<IWriteStream> beginWriteStream(uint64_t offset) const
     {
-        return std::unique_ptr<IWriteStream>(new WriteStream(m_path, offset));
+        return std::unique_ptr<IWriteStream>(new WriteStream(m_path, offset, m_maxWriteSize));
     }
 
     struct ReadStream : public IFileIO::IReadStream
@@ -162,14 +169,14 @@ public:
     }
 };
 
-std::unique_ptr<IFileIO> NewFileIO(const SystemString& path)
+std::unique_ptr<IFileIO> NewFileIO(const SystemString& path, int64_t maxWriteSize)
 {
-    return std::unique_ptr<IFileIO>(new FileIOFILE(path));
+    return std::unique_ptr<IFileIO>(new FileIOFILE(path, maxWriteSize));
 }
 
-std::unique_ptr<IFileIO> NewFileIO(const SystemChar* path)
+std::unique_ptr<IFileIO> NewFileIO(const SystemChar* path, int64_t maxWriteSize)
 {
-    return std::unique_ptr<IFileIO>(new FileIOFILE(path));
+    return std::unique_ptr<IFileIO>(new FileIOFILE(path, maxWriteSize));
 }
 
 }

lib/FileIOWin32.cpp

@@ -1,5 +1,6 @@
 #include <stdio.h>
 #include <stdlib.h>
+#include <inttypes.h>
 #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)
+    FileIOWin32(const SystemString& path, int64_t maxWriteSize)
-    : m_path(path) {}
+    : m_path(path), m_maxWriteSize(maxWriteSize) {}
-    FileIOWin32(const SystemChar* path)
+    FileIOWin32(const SystemChar* path, int64_t maxWriteSize)
-    : m_path(path) {}
+    : 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<IWriteStream> beginWriteStream() const
     {
-        return std::unique_ptr<IWriteStream>(new WriteStream(m_path));
+        return std::unique_ptr<IWriteStream>(new WriteStream(m_path, m_maxWriteSize));
     }
     std::unique_ptr<IWriteStream> beginWriteStream(uint64_t offset) const
     {
-        return std::unique_ptr<IWriteStream>(new WriteStream(m_path, offset));
+        return std::unique_ptr<IWriteStream>(new WriteStream(m_path, offset, m_maxWriteSize));
     }
 
     struct ReadStream : public IFileIO::IReadStream