Single-pass image creation refactor

include/NOD/DiscBase.hpp

@@ -76,7 +76,8 @@ struct Header
             m_gcnMagic = 0xC2339F3D;
     }
 
-    void write(IFileIO::IWriteStream& ws) const
+    template <class WriteStream>
+    void write(WriteStream& ws) const
     {
         Header hs(*this);
         hs.m_wiiMagic = SBig(hs.m_wiiMagic);
@@ -316,9 +317,9 @@ public:
         std::vector<FSTNode> m_buildNodes;
         std::vector<std::string> m_buildNames;
         size_t m_buildNameOff = 0;
-        virtual uint64_t userAllocate(uint64_t reqSz)=0;
+        virtual uint64_t userAllocate(uint64_t reqSz, IPartWriteStream& ws)=0;
         virtual uint32_t packOffset(uint64_t offset) const=0;
-        void recursiveBuildNodes(bool system, const SystemChar* dirIn, uint64_t dolInode);
+        void recursiveBuildNodes(IPartWriteStream& ws, bool system, const SystemChar* dirIn, uint64_t dolInode);
         void recursiveBuildFST(const SystemChar* dirIn, uint64_t dolInode,
                                std::function<void(void)> incParents);
         void addBuildName(const SystemString& str)
@@ -342,22 +343,25 @@ public:
         {
             memcpy(m_gameID, gameID, 6);
         }
-        bool buildFromDirectory(const SystemChar* dirIn, const SystemChar* dolIn,
+        virtual std::unique_ptr<IPartWriteStream> beginWriteStream(uint64_t offset)=0;
+        bool buildFromDirectory(IPartWriteStream& ws,
+                                const SystemChar* dirIn, const SystemChar* dolIn,
                                 const SystemChar* apploaderIn);
 
         const char* getGameID() const {return m_gameID;}
         const std::string& getGameTitle() const {return m_gameTitle;}
     };
 protected:
+    const SystemChar* m_outPath;
     std::unique_ptr<IFileIO> m_fileIO;
     std::vector<std::unique_ptr<PartitionBuilderBase>> m_partitions;
 public:
     std::function<void(size_t idx, const SystemString&, size_t)> m_progressCB;
     size_t m_progressIdx = 0;
     virtual ~DiscBuilderBase() {}
-    DiscBuilderBase(std::unique_ptr<IFileIO>&& fio,
+    DiscBuilderBase(const SystemChar* outPath,
                     std::function<void(size_t idx, const SystemString&, size_t)> progressCB)
-    : m_fileIO(std::move(fio)), m_progressCB(progressCB) {}
+    : m_fileIO(std::move(NewFileIO(outPath))), m_outPath(outPath), m_progressCB(progressCB) {}
 
     IFileIO& getFileIO() {return *m_fileIO;}
 };

include/NOD/DiscWii.hpp

@@ -15,13 +15,14 @@ public:
 
 class DiscBuilderWii : public DiscBuilderBase
 {
-    const SystemChar* m_outPath;
     bool m_dualLayer;
 public:
     DiscBuilderWii(const SystemChar* outPath, const char gameID[6], const char* gameTitle, bool dualLayer,
                    std::function<void(size_t, const SystemString&, size_t)> progressCB);
-    bool buildFromDirectory(const SystemChar* dirIn, const SystemChar* dolIn,
+    bool buildFromDirectory(const SystemChar* dirIn,
-                            const SystemChar* apploaderIn, const SystemChar* partHeadIn);
+                            const SystemChar* dolIn,
+                            const SystemChar* apploaderIn,
+                            const SystemChar* partHeadIn);
 };
 
 }

include/NOD/IDiscIO.hpp

@@ -47,9 +47,9 @@ struct IPartReadStream
 struct IPartWriteStream
 {
     virtual ~IPartWriteStream() {}
-    virtual void seek(int64_t offset, int whence=SEEK_SET)=0;
+    virtual void close()=0;
     virtual uint64_t position() const=0;
-    virtual uint64_t write(void* buf, uint64_t length)=0;
+    virtual uint64_t write(const void* buf, uint64_t length)=0;
 };
 
 #if NOD_ATHENA
@@ -72,24 +72,6 @@ public:
     inline atUint64 readUBytesToBuf(void* buf, atUint64 sz) {m_rs->read(buf, sz); return sz;}
 };
 
-class AthenaPartWriteStream : public Athena::io::IStreamWriter
-{
-    std::unique_ptr<IPartWriteStream> m_ws;
-public:
-    AthenaPartWriteStream(std::unique_ptr<IPartWriteStream>&& ws) : m_ws(std::move(ws)) {}
-
-    inline void seek(atInt64 off, Athena::SeekOrigin origin)
-    {
-        if (origin == Athena::Begin)
-            m_ws->seek(off, SEEK_SET);
-        else if (origin == Athena::Current)
-            m_ws->seek(off, SEEK_CUR);
-    }
-    inline atUint64 position()    const {return m_ws->position();}
-    inline atUint64 length()      const {return 0;}
-    inline void writeUBytes(const atUint8* buf, atUint64 len) {m_ws->write((void*)buf, len);}
-};
-
 #endif
 
 }

include/NOD/Util.hpp

@@ -16,6 +16,9 @@
 #include <sys/file.h>
 #include <unistd.h>
 #include <errno.h>
+#include <sys/param.h>
+#include <sys/statvfs.h>
+
 #endif
 #include <sys/stat.h>
 
@@ -261,6 +264,21 @@ static inline int FSeek(FILE* fp, int64_t offset, int whence)
 #endif
 }
 
+static inline bool CheckFreeSpace(const SystemChar* path, size_t reqSz)
+{
+#if _WIN32
+    ULARGE_INTEGER freeBytes;
+    if (!GetDiskFreeSpaceExW(path, &freeBytes, nullptr, nullptr))
+        LogModule.report(LogVisor::FatalError, "GetDiskFreeSpaceExW %s: %d", path, GetLastError());
+    return reqSz < freeBytes;
+#else
+    struct statvfs svfs;
+    if (statvfs(path, &svfs))
+        LogModule.report(LogVisor::FatalError, "statvfs %s: %s", path, strerror(errno));
+    return reqSz < svfs.f_bsize * svfs.f_bfree;
+#endif
+}
+
 }
 
 #endif // __NOD_UTIL_HPP__

lib/DiscBase.cpp

@@ -131,7 +131,14 @@ bool DiscBase::IPartition::extractToDirectory(const SystemString& path,
     }
 
     /* Extract Filesystem */
-    return m_nodes[0].extractToDirectory(path, ctx);
+    SystemString fsPath = path + _S("/fsroot");
+    if (Mkdir(fsPath.c_str(), 0755) && errno != EEXIST)
+    {
+        LogModule.report(LogVisor::Error, _S("unable to mkdir '%s'"), fsPath.c_str());
+        return false;
+    }
+
+    return m_nodes[0].extractToDirectory(fsPath, ctx);
 }
 
 static uint64_t GetInode(const SystemChar* path)
@@ -157,13 +164,17 @@ static uint64_t GetInode(const SystemChar* path)
     return inode;
 }
 
-static bool IsSystemFile(const SystemString& name)
+static bool IsSystemFile(const SystemString& name, bool& isDol)
 {
+    isDol = false;
     if (name.size() < 4)
         return false;
 
     if (!StrCaseCmp((&*name.cend()) - 4, _S(".dol")))
+    {
+        isDol = true;
         return true;
+    }
     if (!StrCaseCmp((&*name.cend()) - 4, _S(".rel")))
         return true;
     if (!StrCaseCmp((&*name.cend()) - 4, _S(".rso")))
@@ -180,7 +191,23 @@ static bool IsSystemFile(const SystemString& name)
     return false;
 }
 
-void DiscBuilderBase::PartitionBuilderBase::recursiveBuildNodes(bool system, const SystemChar* dirIn,
+/** Patches out pesky #001 integrity check performed by game's OSInit.
+ *  This is required for multi-DOL games, but doesn't harm functionality otherwise */
+static size_t PatchDOL(IFileIO::IReadStream& in, IPartWriteStream& out, size_t sz)
+{
+    std::unique_ptr<uint8_t[]> buf(new uint8_t[sz]);
+    sz = in.read(buf.get(), sz);
+    uint8_t* found = static_cast<uint8_t*>(memmem(buf.get(), sz,
+                            "\x3C\x03\xF8\x00\x28\x00\x00\x00\x40\x82\x00\x0C"
+                            "\x38\x60\x00\x01\x48\x00\x02\x44\x38\x61\x00\x18\x48", 25));
+    if (found)
+        found[11] = '\x04';
+    return out.write(buf.get(), sz);
+}
+
+void DiscBuilderBase::PartitionBuilderBase::recursiveBuildNodes(IPartWriteStream& ws,
+                                                                bool system,
+                                                                const SystemChar* dirIn,
                                                                 uint64_t dolInode)
 {
     DirectoryEnumerator dEnum(dirIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
@@ -188,12 +215,12 @@ void DiscBuilderBase::PartitionBuilderBase::recursiveBuildNodes(bool system, con
     {
         if (e.m_isDir)
         {
-            size_t dirNodeIdx = m_buildNodes.size();
+            recursiveBuildNodes(ws, system, e.m_path.c_str(), dolInode);
-            recursiveBuildNodes(system, e.m_path.c_str(), dolInode);
         }
         else
         {
-            bool isSys = IsSystemFile(e.m_name);
+            bool isDol;
+            bool isSys = IsSystemFile(e.m_name, isDol);
             if (system ^ isSys)
                 continue;
 
@@ -201,27 +228,31 @@ void DiscBuilderBase::PartitionBuilderBase::recursiveBuildNodes(bool system, con
                 continue;
 
             size_t fileSz = ROUND_UP_32(e.m_fileSz);
-            uint64_t fileOff = userAllocate(fileSz);
+            uint64_t fileOff = userAllocate(fileSz, ws);
             m_fileOffsetsSizes[e.m_path] = std::make_pair(fileOff, fileSz);
-            std::unique_ptr<IFileIO::IWriteStream> ws = m_parent.getFileIO().beginWriteStream(fileOff);
+            std::unique_ptr<IFileIO::IReadStream> rs = NewFileIO(e.m_path)->beginReadStream();
-            FILE* fp = Fopen(e.m_path.c_str(), _S("rb"), FileLockType::Read);
-            if (!fp)
-                LogModule.report(LogVisor::FatalError, _S("unable to open '%s' for reading"), e.m_path.c_str());
-            char buf[0x8000];
             size_t xferSz = 0;
+            if (isDol)
+            {
+                xferSz = PatchDOL(*rs, ws, e.m_fileSz);
+                m_parent.m_progressCB(++m_parent.m_progressIdx, e.m_name + _S(" [PATCHED]"), xferSz);
+            }
+            else
+            {
+                char buf[0x8000];
                 ++m_parent.m_progressIdx;
                 while (xferSz < e.m_fileSz)
                 {
-                size_t rdSz = fread(buf, 1, std::min(0x8000ul, e.m_fileSz - xferSz), fp);
+                    size_t rdSz = rs->read(buf, std::min(0x8000ul, e.m_fileSz - xferSz));
                     if (!rdSz)
                         break;
-                ws->write(buf, rdSz);
+                    ws.write(buf, rdSz);
                     xferSz += rdSz;
                     m_parent.m_progressCB(m_parent.m_progressIdx, e.m_name, xferSz);
                 }
-            fclose(fp);
+            }
             for (size_t i=0 ; i<fileSz-xferSz ; ++i)
-                ws->write("\xff", 1);
+                ws.write("\xff", 1);
         }
     }
 }
@@ -258,7 +289,8 @@ void DiscBuilderBase::PartitionBuilderBase::recursiveBuildFST(const SystemChar*
     }
 }
 
-bool DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(const SystemChar* dirIn,
+bool DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(IPartWriteStream& ws,
+                                                               const SystemChar* dirIn,
                                                                const SystemChar* dolIn,
                                                                const SystemChar* apploaderIn)
 {
@@ -266,7 +298,6 @@ bool DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(const SystemChar*
         LogModule.report(LogVisor::FatalError, "all arguments must be supplied to buildFromDirectory()");
 
     /* Clear file */
-    m_parent.getFileIO().beginWriteStream();
     ++m_parent.m_progressIdx;
     m_parent.m_progressCB(m_parent.m_progressIdx, "Preparing output image", -1);
 
@@ -274,41 +305,26 @@ bool DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(const SystemChar*
     m_buildNodes.emplace_back(true, m_buildNameOff, 0, 1);
     addBuildName(_S("<root>"));
 
-    /* Write DOL first (ensures that it's within a 32-bit offset for Wii apploaders) */
+    /* Write Boot DOL first (first thing seeked to after Apploader) */
     {
         Sstat dolStat;
         if (Stat(dolIn, &dolStat))
             LogModule.report(LogVisor::FatalError, _S("unable to stat %s"), dolIn);
         size_t fileSz = ROUND_UP_32(dolStat.st_size);
-        uint64_t fileOff = userAllocate(fileSz);
+        uint64_t fileOff = userAllocate(fileSz, ws);
         m_dolOffset = fileOff;
         m_dolSize = fileSz;
-        std::unique_ptr<IFileIO::IWriteStream> ws = m_parent.getFileIO().beginWriteStream(fileOff);
+        std::unique_ptr<IFileIO::IReadStream> rs = NewFileIO(dolIn)->beginReadStream();
-        FILE* fp = Fopen(dolIn, _S("rb"), FileLockType::Read);
+        size_t xferSz = PatchDOL(*rs, ws, dolStat.st_size);
-        if (!fp)
+        m_parent.m_progressCB(++m_parent.m_progressIdx, SystemString(dolIn) + _S(" [PATCHED]"), xferSz);
-            LogModule.report(LogVisor::FatalError, _S("unable to open '%s' for reading"), dolIn);
-        char buf[8192];
-        size_t xferSz = 0;
-        SystemString dolName(dolIn);
-        ++m_parent.m_progressIdx;
-        while (xferSz < dolStat.st_size)
-        {
-            size_t rdSz = fread(buf, 1, std::min(size_t(8192), size_t(dolStat.st_size - xferSz)), fp);
-            if (!rdSz)
-                break;
-            ws->write(buf, rdSz);
-            xferSz += rdSz;
-            m_parent.m_progressCB(m_parent.m_progressIdx, dolName, xferSz);
-        }
-        fclose(fp);
         for (size_t i=0 ; i<fileSz-xferSz ; ++i)
-            ws->write("\xff", 1);
+            ws.write("\xff", 1);
     }
 
     /* Gather files in root directory */
     uint64_t dolInode = GetInode(dolIn);
-    recursiveBuildNodes(true, dirIn, dolInode);
+    recursiveBuildNodes(ws, true, dirIn, dolInode);
-    recursiveBuildNodes(false, dirIn, dolInode);
+    recursiveBuildNodes(ws, false, dirIn, dolInode);
     recursiveBuildFST(dirIn, dolInode, [&](){m_buildNodes[0].incrementLength();});
 
     return true;

lib/DiscGCN.cpp

@@ -113,12 +113,40 @@ class PartitionBuilderGCN : public DiscBuilderBase::PartitionBuilderBase
 {
     uint64_t m_curUser = 0x57058000;
     uint32_t m_fstMemoryAddr;
+
 public:
+    class PartWriteStream : public IPartWriteStream
+    {
+        const PartitionBuilderGCN& m_parent;
+        uint64_t m_offset;
+        std::unique_ptr<IFileIO::IWriteStream> m_fio;
+
+    public:
+        PartWriteStream(const PartitionBuilderGCN& parent, uint64_t offset)
+        : m_parent(parent), m_offset(offset)
+        {
+            m_fio = m_parent.m_parent.getFileIO().beginWriteStream(offset);
+        }
+        void close() {m_fio.reset();}
+        uint64_t position() const {return m_offset;}
+        uint64_t write(const void* buf, uint64_t length)
+        {
+            uint64_t len = m_fio->write(buf, length);
+            m_offset += len;
+            return len;
+        }
+        void seek(size_t off)
+        {
+            m_offset = off;
+            m_fio = m_parent.m_parent.getFileIO().beginWriteStream(off);
+        }
+    };
+
     PartitionBuilderGCN(DiscBuilderBase& parent, Kind kind,
                         const char gameID[6], const char* gameTitle, uint32_t fstMemoryAddr)
     : DiscBuilderBase::PartitionBuilderBase(parent, kind, gameID, gameTitle), m_fstMemoryAddr(fstMemoryAddr) {}
 
-    uint64_t userAllocate(uint64_t reqSz)
+    uint64_t userAllocate(uint64_t reqSz, IPartWriteStream& ws)
     {
         m_curUser -= reqSz;
         m_curUser &= 0xfffffffffffffff0;
@@ -127,6 +155,7 @@ public:
             LogModule.report(LogVisor::FatalError, "user area low mark reached");
             return -1;
         }
+        static_cast<PartWriteStream&>(ws).seek(m_curUser);
         return m_curUser;
     }
 
@@ -135,17 +164,23 @@ public:
         return offset;
     }
 
+    std::unique_ptr<IPartWriteStream> beginWriteStream(uint64_t offset)
+    {
+        return std::make_unique<PartWriteStream>(*this, offset);
+    }
+
     bool buildFromDirectory(const SystemChar* dirIn, const SystemChar* dolIn, const SystemChar* apploaderIn)
     {
-        bool result = DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(dirIn, dolIn, apploaderIn);
+        std::unique_ptr<IPartWriteStream> ws = beginWriteStream(0);
+        bool result = DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(*ws, dirIn, dolIn, apploaderIn);
         if (!result)
             return false;
 
-        std::unique_ptr<IFileIO::IWriteStream> ws = m_parent.getFileIO().beginWriteStream(0);
+        ws = beginWriteStream(0);
         Header header(m_gameID, m_gameTitle.c_str(), false);
         header.write(*ws);
 
-        ws = m_parent.getFileIO().beginWriteStream(0x2440);
+        ws = beginWriteStream(0x2440);
         FILE* fp = Fopen(apploaderIn, _S("rb"), FileLockType::Read);
         if (!fp)
             LogModule.report(LogVisor::FatalError, _S("unable to open %s for reading"), apploaderIn);
@@ -182,7 +217,7 @@ public:
             LogModule.report(LogVisor::FatalError,
                              "FST flows into user area (one or the other is too big)");
 
-        ws = m_parent.getFileIO().beginWriteStream(0x420);
+        ws = beginWriteStream(0x420);
         uint32_t vals[7];
         vals[0] = SBig(uint32_t(m_dolOffset));
         vals[1] = SBig(uint32_t(fstOff));
@@ -200,13 +235,21 @@ public:
 bool DiscBuilderGCN::buildFromDirectory(const SystemChar* dirIn, const SystemChar* dolIn,
                                         const SystemChar* apploaderIn)
 {    
+    m_fileIO->beginWriteStream();
+
+    if (!CheckFreeSpace(m_outPath, 0x57058000))
+    {
+        LogModule.report(LogVisor::Error, _S("not enough free disk space for %s"), m_outPath);
+        return false;
+    }
+
     PartitionBuilderGCN& pb = static_cast<PartitionBuilderGCN&>(*m_partitions[0]);
     return pb.buildFromDirectory(dirIn, dolIn, apploaderIn);
 }
 
 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(std::move(NewFileIO(outPath)), progressCB)
+: DiscBuilderBase(outPath, progressCB)
 {
     PartitionBuilderGCN* partBuilder = new PartitionBuilderGCN(*this, PartitionBuilderBase::Kind::Data,
                                                                gameID, gameTitle, fstMemoryAddr);

lib/DiscWii.cpp

@@ -433,17 +433,171 @@ void DiscWii::writeOutDataPartitionHeader(const SystemChar* pathOut) const
     }
 }
 
+static const uint8_t ZEROIV[16] = {0};
+
 class PartitionBuilderWii : public DiscBuilderBase::PartitionBuilderBase
 {
+    friend class DiscBuilderWii;
+    uint64_t m_baseOffset;
+    uint64_t m_userOffset = 0;
     uint64_t m_curUser = 0x1F0000;
+    std::unique_ptr<IAES> m_aes;
+    uint8_t m_h3[4916][20] = {};
+
 public:
+    class PartWriteStream : public IPartWriteStream
+    {
+        friend class PartitionBuilderWii;
+        PartitionBuilderWii& m_parent;
+        uint64_t m_baseOffset;
+        uint64_t m_offset;
+        std::unique_ptr<IFileIO::IWriteStream> m_fio;
+        bool m_closed = false;
+
+        size_t m_curGroup = SIZE_MAX;
+        char m_buf[0x200000];
+
+        void encryptGroup(uint8_t h3Out[20])
+        {
+            sha1nfo sha;
+            uint8_t h2[8][20];
+
+            for (int s=0 ; s<8 ; ++s)
+            {
+                char* ptr1 = m_buf + s*0x40000;
+                uint8_t h1[8][20];
+
+                for (int c=0 ; c<8 ; ++c)
+                {
+                    char* ptr0 = ptr1 + c*0x8000;
+                    uint8_t h0[31][20];
+
+                    for (int j=0 ; j<31 ; ++j)
+                    {
+                        sha1_init(&sha);
+                        sha1_write(&sha, ptr0 + (j+1)*0x400, 0x400);
+                        memcpy(h0[j], sha1_result(&sha), 20);
+                    }
+
+                    sha1_init(&sha);
+                    sha1_write(&sha, (char*)h0, 0x26C);
+                    memcpy(h1[c], sha1_result(&sha), 20);
+
+                    memcpy(ptr0, h0, 0x26C);
+                    memset(ptr0+0x26C, 0, 0x014);
+                }
+
+                sha1_init(&sha);
+                sha1_write(&sha, (char*)h1, 0x0A0);
+                memcpy(h2[s], sha1_result(&sha), 20);
+
+                for (int c=0 ; c<8 ; ++c)
+                {
+                    char* ptr0 = ptr1 + c*0x8000;
+                    memcpy(ptr0+0x280, h1, 0x0A0);
+                    memset(ptr0+0x320, 0, 0x020);
+                }
+            }
+
+            sha1_init(&sha);
+            sha1_write(&sha, (char*)h2, 0x0A0);
+            memcpy(h3Out, sha1_result(&sha), 20);
+
+            for (int s=0 ; s<8 ; ++s)
+            {
+                char* ptr1 = m_buf + s*0x40000;
+                for (int c=0 ; c<8 ; ++c)
+                {
+                    char* ptr0 = ptr1 + c*0x8000;
+                    memcpy(ptr0+0x340, h2, 0x0A0);
+                    memset(ptr0+0x3E0, 0, 0x020);
+                    m_parent.m_aes->encrypt(ZEROIV, (uint8_t*)ptr0, (uint8_t*)ptr0, 0x400);
+                    m_parent.m_aes->encrypt((uint8_t*)(ptr0+0x3D0), (uint8_t*)(ptr0+0x400), (uint8_t*)(ptr0+0x400), 0x7c00);
+                }
+            }
+
+            if (m_fio->write(m_buf, 0x200000) != 0x200000)
+                LogModule.report(LogVisor::FatalError, "unable to write full disc group");
+        }
+
+    public:
+        PartWriteStream(PartitionBuilderWii& parent, uint64_t baseOffset, uint64_t offset)
+        : m_parent(parent), m_baseOffset(baseOffset), m_offset(offset)
+        {
+            if (offset % 0x1F0000)
+                LogModule.report(LogVisor::FatalError, "partition write stream MUST begin on 0x1F0000-aligned boundary");
+            size_t group = m_offset / 0x1F0000;
+            m_fio = m_parent.m_parent.getFileIO().beginWriteStream(m_baseOffset + group * 0x200000);
+            m_curGroup = group;
+        }
+        ~PartWriteStream() {close();}
+        void close()
+        {
+            if (m_closed)
+                return;
+            m_closed = true;
+            size_t rem = m_offset % 0x1F0000;
+            if (rem)
+            {
+                rem = 0x1F0000 - rem;
+                write(nullptr, rem);
+            }
+            encryptGroup(m_parent.m_h3[m_curGroup]);
+            m_fio.reset();
+        }
+        uint64_t position() const {return m_offset;}
+        uint64_t write(const void* buf, uint64_t length)
+        {
+            size_t group = m_offset / 0x1F0000;
+            size_t block = (m_offset - group * 0x1F0000) / 0x7c00;
+            size_t cacheOffset = m_offset % 0x7c00;
+            uint64_t cacheSize;
+            uint64_t rem = length;
+            const uint8_t* src = (uint8_t*)buf;
+
+            while (rem)
+            {
+                if (group != m_curGroup)
+                {
+                    encryptGroup(m_parent.m_h3[m_curGroup]);
+                    m_curGroup = group;
+                }
+
+                cacheSize = rem;
+                if (cacheSize + cacheOffset > 0x7c00)
+                    cacheSize = 0x7c00 - cacheOffset;
+
+                if (src)
+                {
+                    memcpy(m_buf + block * 0x8000 + 0x400 + cacheOffset, src, cacheSize);
+                    src += cacheSize;
+                }
+                else
+                    memset(m_buf + block * 0x8000 + 0x400 + cacheOffset, 0, cacheSize);
+
+                rem -= cacheSize;
+                cacheOffset = 0;
+                ++block;
+                if (block == 64)
+                {
+                    block = 0;
+                    ++group;
+                }
+            }
+
+            m_offset += length;
+            return length;
+        }
+    };
+
     PartitionBuilderWii(DiscBuilderBase& parent, Kind kind,
-                        const char gameID[6], const char* gameTitle)
+                        const char gameID[6], const char* gameTitle, uint64_t baseOffset)
-    : DiscBuilderBase::PartitionBuilderBase(parent, kind, gameID, gameTitle) {}
+    : DiscBuilderBase::PartitionBuilderBase(parent, kind, gameID, gameTitle),
+      m_baseOffset(baseOffset), m_aes(NewAES()) {}
 
     uint64_t getCurUserEnd() const {return m_curUser;}
 
-    uint64_t userAllocate(uint64_t reqSz)
+    uint64_t userAllocate(uint64_t reqSz, IPartWriteStream& ws)
     {
         reqSz = ROUND_UP_32(reqSz);
         if (m_curUser + reqSz >= 0x1FB450000)
@@ -452,6 +606,14 @@ public:
             return -1;
         }
         uint64_t ret = m_curUser;
+        PartWriteStream& cws = static_cast<PartWriteStream&>(ws);
+        if (cws.m_offset > ret)
+        {
+            LogModule.report(LogVisor::FatalError, "partition overwrite error");
+            return -1;
+        }
+        while (cws.m_offset < ret)
+            cws.write("\xff", 1);
         m_curUser += reqSz;
         return ret;
     }
@@ -461,102 +623,34 @@ public:
         return uint32_t(offset >> uint64_t(2));
     }
 
-    bool buildFromDirectory(const SystemChar* dirIn, const SystemChar* dolIn, const SystemChar* apploaderIn)
+    std::unique_ptr<IPartWriteStream> beginWriteStream(uint64_t offset)
     {
-        bool result = DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(dirIn, dolIn, apploaderIn);
+        return std::make_unique<PartWriteStream>(*this, m_baseOffset + m_userOffset, offset);
-        if (!result)
-            return false;
-
-        std::unique_ptr<IFileIO::IWriteStream> ws;
-
-        /* Pad out user area to nearest cleartext sector */
-        uint64_t curUserRem = m_curUser % 0x1F0000;
-        if (curUserRem)
-        {
-            ws = m_parent.getFileIO().beginWriteStream(m_curUser);
-            curUserRem = 0x1F0000 - curUserRem;
-            for (size_t i=0 ; i<curUserRem ; ++i)
-                ws->write("\xff", 1);
-            m_curUser += curUserRem;
     }
 
-        ws = m_parent.getFileIO().beginWriteStream(0);
+    uint64_t buildFromDirectory(const SystemChar* dirIn,
-        Header header(m_gameID, m_gameTitle.c_str(), true, 0, 0, 0);
+                                const SystemChar* dolIn,
-        header.write(*ws);
+                                const SystemChar* apploaderIn,
-
+                                const SystemChar* partHeadIn)
-        ws = m_parent.getFileIO().beginWriteStream(0x2440);
-        FILE* fp = Fopen(apploaderIn, _S("rb"), FileLockType::Read);
-        if (!fp)
-            LogModule.report(LogVisor::FatalError, _S("unable to open %s for reading"), apploaderIn);
-        char buf[8192];
-        size_t xferSz = 0;
-        SystemString apploaderName(apploaderIn);
-        ++m_parent.m_progressIdx;
-        while (true)
-        {
-            size_t rdSz = fread(buf, 1, 8192, fp);
-            if (!rdSz)
-                break;
-            ws->write(buf, rdSz);
-            xferSz += rdSz;
-            if (0x2440 + xferSz >= 0x1F0000)
-                LogModule.report(LogVisor::FatalError,
-                                 "apploader flows into user area (one or the other is too big)");
-            m_parent.m_progressCB(m_parent.m_progressIdx, apploaderName, xferSz);
-        }
-        fclose(fp);
-
-        size_t fstOff = ROUND_UP_32(xferSz);
-        size_t fstSz = sizeof(FSTNode) * m_buildNodes.size();
-        for (size_t i=0 ; i<fstOff-xferSz ; ++i)
-            ws->write("\xff", 1);
-        fstOff += 0x2440;
-        ws->write(m_buildNodes.data(), fstSz);
-        for (const std::string& str : m_buildNames)
-            ws->write(str.data(), str.size()+1);
-        fstSz += m_buildNameOff;
-        fstSz = ROUND_UP_32(fstSz);
-
-        if (fstOff + fstSz >= 0x1F0000)
-            LogModule.report(LogVisor::FatalError,
-                             "FST flows into user area (one or the other is too big)");
-
-        ws = m_parent.getFileIO().beginWriteStream(0x420);
-        uint32_t vals[4];
-        vals[0] = SBig(uint32_t(m_dolOffset >> uint64_t(2)));
-        vals[1] = SBig(uint32_t(fstOff >> uint64_t(2)));
-        vals[2] = SBig(uint32_t(fstSz));
-        vals[3] = SBig(uint32_t(fstSz));
-        ws->write(vals, sizeof(vals));
-
-        return true;
-    }
-
-    uint64_t cryptAndFakesign(IFileIO& out, uint64_t offset, const SystemChar* partHeadIn) const
     {
         /* Read head and validate key members */
-        FILE* fp = Fopen(partHeadIn, _S("rb"), FileLockType::Read);
+        std::unique_ptr<IFileIO> ph = NewFileIO(partHeadIn);
-        if (!fp)
-            LogModule.report(LogVisor::FatalError, _S("unable to open %s for reading"), partHeadIn);
 
         uint8_t tkey[16];
         {
-            FSeek(fp, 0x1BF, SEEK_SET);
+            if (ph->beginReadStream(0x1BF)->read(tkey, 16) != 16)
-            if (fread(tkey, 1, 16, fp) != 16)
                 LogModule.report(LogVisor::FatalError, _S("unable to read title key from %s"), partHeadIn);
         }
 
         uint8_t tkeyiv[16] = {};
         {
-            FSeek(fp, 0x1DC, SEEK_SET);
+            if (ph->beginReadStream(0x1DC)->read(tkeyiv, 8) != 8)
-            if (fread(tkeyiv, 1, 8, fp) != 8)
                 LogModule.report(LogVisor::FatalError, _S("unable to read title key IV from %s"), partHeadIn);
         }
 
         uint8_t ccIdx;
         {
-            FSeek(fp, 0x1F1, SEEK_SET);
+            if (ph->beginReadStream(0x1F1)->read(&ccIdx, 1) != 1)
-            if (fread(&ccIdx, 1, 1, fp) != 1)
                 LogModule.report(LogVisor::FatalError, _S("unable to read common key index from %s"), partHeadIn);
             if (ccIdx > 1)
                 LogModule.report(LogVisor::FatalError, _S("common key index may only be 0 or 1"));
@@ -564,8 +658,7 @@ public:
 
         uint32_t tmdSz;
         {
-            FSeek(fp, 0x2A4, SEEK_SET);
+            if (ph->beginReadStream(0x2A4)->read(&tmdSz, 4) != 4)
-            if (fread(&tmdSz, 1, 4, fp) != 4)
                 LogModule.report(LogVisor::FatalError, _S("unable to read TMD size from %s"), partHeadIn);
             tmdSz = SBig(tmdSz);
         }
@@ -573,8 +666,7 @@ public:
         uint64_t h3Off;
         {
             uint32_t h3Ptr;
-            FSeek(fp, 0x2B4, SEEK_SET);
+            if (ph->beginReadStream(0x2B4)->read(&h3Ptr, 4) != 4)
-            if (fread(&h3Ptr, 1, 4, fp) != 4)
                 LogModule.report(LogVisor::FatalError, _S("unable to read H3 pointer from %s"), partHeadIn);
             h3Off = uint64_t(SBig(h3Ptr)) << 2;
         }
@@ -582,25 +674,26 @@ public:
         uint64_t dataOff;
         {
             uint32_t dataPtr;
-            if (fread(&dataPtr, 1, 4, fp) != 4)
+            if (ph->beginReadStream(0x2B8)->read(&dataPtr, 4) != 4)
                 LogModule.report(LogVisor::FatalError, _S("unable to read data pointer from %s"), partHeadIn);
             dataOff = uint64_t(SBig(dataPtr)) << 2;
         }
+        m_userOffset = dataOff;
 
         std::unique_ptr<uint8_t[]> tmdData(new uint8_t[tmdSz]);
-        FSeek(fp, 0x2C0, SEEK_SET);
+        if (ph->beginReadStream(0x2C0)->read(tmdData.get(), tmdSz) != tmdSz)
-        if (fread(tmdData.get(), 1, tmdSz, fp) != tmdSz)
             LogModule.report(LogVisor::FatalError, _S("unable to read TMD from %s"), partHeadIn);
 
         /* Copy partition head up to H3 table */
-        std::unique_ptr<IFileIO::IWriteStream> ws = out.beginWriteStream(offset);
+        std::unique_ptr<IFileIO::IWriteStream> ws = m_parent.getFileIO().beginWriteStream(m_baseOffset);
         {
             uint64_t remCopy = h3Off;
+
             uint8_t copyBuf[8192];
-            FSeek(fp, 0, SEEK_SET);
+            std::unique_ptr<IFileIO::IReadStream> rs = ph->beginReadStream();
             while (remCopy)
             {
-                size_t rdBytes = fread(copyBuf, 1, std::min(size_t(8192), size_t(remCopy)), fp);
+                size_t rdBytes = rs->read(copyBuf, std::min(size_t(8192), size_t(remCopy)));
                 if (rdBytes)
                 {
                     ws->write(copyBuf, rdBytes);
@@ -613,111 +706,104 @@ public:
             }
         }
 
-        fclose(fp);
-
         /* Prepare crypto pass */
-        std::unique_ptr<IFileIO::IReadStream> rs = m_parent.getFileIO().beginReadStream(0);
+        m_aes->setKey(COMMON_KEYS[ccIdx]);
-        sha1nfo sha;
+        m_aes->decrypt(tkeyiv, tkey, tkey, 16);
-        std::unique_ptr<IAES> aes = NewAES();
+        m_aes->setKey(tkey);
-        aes->setKey(COMMON_KEYS[ccIdx]);
-        aes->decrypt(tkeyiv, tkey, tkey, 16);
-        aes->setKey(tkey);
-        static const uint8_t ZEROIV[16] = {0};
 
-        std::unique_ptr<char[]> cleartext(new char[0x1F0000]);
+        {
-        std::unique_ptr<char[]> ciphertext(new char[0x200000]);
+            /* Assemble partition data */
-        uint8_t h3[4916][20] = {};
+            std::unique_ptr<IPartWriteStream> cws = beginWriteStream(0x1F0000);
+            bool result = DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(*cws, dirIn, dolIn, apploaderIn);
+            if (!result)
+                return 0;
 
-        uint64_t groupCount = m_curUser / 0x1F0000;
+            /* Pad out user area to nearest cleartext sector */
-        ws = out.beginWriteStream(offset + dataOff);
+            m_curUser = cws->position();
-        SystemString cryptoName(_S("Hashing and encrypting"));
+            uint64_t curUserRem = m_curUser % 0x1F0000;
+            if (curUserRem)
+            {
+                curUserRem = 0x1F0000 - curUserRem;
+                for (size_t i=0 ; i<curUserRem ; ++i)
+                    cws->write("\xff", 1);
+                m_curUser += curUserRem;
+            }
+
+            /* Begin crypto write and add content header */
+            cws = beginWriteStream(0);
+            Header header(m_gameID, m_gameTitle.c_str(), true, 0, 0, 0);
+            header.write(*cws);
+
+            /* Get Apploader Size */
+            Sstat theStat;
+            if (Stat(apploaderIn, &theStat))
+                LogModule.report(LogVisor::FatalError, _S("unable to stat %s"), apploaderIn);
+
+            /* Compute boot table members and write */
+            size_t fstOff = 0x2440 + ROUND_UP_32(theStat.st_size);
+            size_t fstSz = sizeof(FSTNode) * m_buildNodes.size();
+            fstSz += m_buildNameOff;
+            fstSz = ROUND_UP_32(fstSz);
+
+            if (fstOff + fstSz >= 0x1F0000)
+                LogModule.report(LogVisor::FatalError,
+                                 "FST flows into user area (one or the other is too big)");
+
+            cws->write(nullptr, 0x420 - sizeof(Header));
+            uint32_t vals[4];
+            vals[0] = SBig(uint32_t(m_dolOffset >> uint64_t(2)));
+            vals[1] = SBig(uint32_t(fstOff >> uint64_t(2)));
+            vals[2] = SBig(uint32_t(fstSz));
+            vals[3] = SBig(uint32_t(fstSz));
+            cws->write(vals, 16);
+
+            /* Write Apploader */
+            cws->write(nullptr, 0x2440 - 0x430);
+            std::unique_ptr<IFileIO::IReadStream> rs = NewFileIO(apploaderIn)->beginReadStream();
+            char buf[8192];
+            size_t xferSz = 0;
+            SystemString apploaderName(apploaderIn);
             ++m_parent.m_progressIdx;
-        for (uint64_t g=0 ; g<groupCount ; ++g)
+            while (true)
             {
-            char* cleartext2 = cleartext.get();
+                size_t rdSz = rs->read(buf, 8192);
-            char* ciphertext2 = ciphertext.get();
+                if (!rdSz)
-
+                    break;
-            if (rs->read(cleartext2, 0x1F0000) != 0x1F0000)
+                cws->write(buf, rdSz);
-                LogModule.report(LogVisor::FatalError, "cleartext file too short");
+                xferSz += rdSz;
-
+                if (0x2440 + xferSz >= 0x1F0000)
-            uint8_t h2[8][20];
+                    LogModule.report(LogVisor::FatalError,
-
+                                     "apploader flows into user area (one or the other is too big)");
-            for (int s=0 ; s<8 ; ++s)
+                m_parent.m_progressCB(m_parent.m_progressIdx, apploaderName, xferSz);
-            {
-                char* cleartext1 = cleartext2 + s*0x3E000;
-                char* ciphertext1 = ciphertext2 + s*0x40000;
-                uint8_t h1[8][20];
-
-                for (int c=0 ; c<8 ; ++c)
-                {
-                    char* cleartext0 = cleartext1 + c*0x7c00;
-                    char* ciphertext0 = ciphertext1 + c*0x8000;
-                    uint8_t h0[31][20];
-
-                    for (int j=0 ; j<31 ; ++j)
-                    {
-                        sha1_init(&sha);
-                        sha1_write(&sha, cleartext0 + j*0x400, 0x400);
-                        memcpy(h0[j], sha1_result(&sha), 20);
             }
 
-                    sha1_init(&sha);
+            size_t fstOffRel = fstOff - 0x2440;
-                    sha1_write(&sha, (char*)h0, 0x26C);
+            if (xferSz > fstOffRel)
-                    memcpy(h1[c], sha1_result(&sha), 20);
+                LogModule.report(LogVisor::FatalError, "apploader unexpectedly flows into FST");
+            for (size_t i=0 ; i<fstOffRel-xferSz ; ++i)
+                cws->write("\xff", 1);
 
-                    memcpy(ciphertext0, h0, 0x26C);
+            /* Write FST */
-                    memset(ciphertext0+0x26C, 0, 0x014);
+            cws->write(m_buildNodes.data(), m_buildNodes.size() * sizeof(FSTNode));
-                }
+            for (const std::string& str : m_buildNames)
-
+                cws->write(str.data(), str.size()+1);
-                sha1_init(&sha);
-                sha1_write(&sha, (char*)h1, 0x0A0);
-                memcpy(h2[s], sha1_result(&sha), 20);
-
-                for (int c=0 ; c<8 ; ++c)
-                {
-                    char* ciphertext0 = ciphertext1 + c*0x8000;
-                    memcpy(ciphertext0+0x280, h1, 0x0A0);
-                    memset(ciphertext0+0x320, 0, 0x020);
-                }
-            }
-
-            sha1_init(&sha);
-            sha1_write(&sha, (char*)h2, 0x0A0);
-            memcpy(h3[g], sha1_result(&sha), 20);
-
-            for (int s=0 ; s<8 ; ++s)
-            {
-                char* cleartext1 = cleartext2 + s*0x3E000;
-                char* ciphertext1 = ciphertext2 + s*0x40000;
-                for (int c=0 ; c<8 ; ++c)
-                {
-                    char* cleartext0 = cleartext1 + c*0x7c00;
-                    char* ciphertext0 = ciphertext1 + c*0x8000;
-                    memcpy(ciphertext0+0x340, h2, 0x0A0);
-                    memset(ciphertext0+0x3E0, 0, 0x020);
-                    aes->encrypt(ZEROIV, (uint8_t*)ciphertext0, (uint8_t*)ciphertext0, 0x400);
-                    aes->encrypt((uint8_t*)(ciphertext0+0x3D0), (uint8_t*)cleartext0, (uint8_t*)(ciphertext0+0x400), 0x7c00);
-                }
-            }
-
-            if (ws->write(ciphertext2, 0x200000) != 0x200000)
-                LogModule.report(LogVisor::FatalError, "unable to write full disc sector");
-            m_parent.m_progressCB(m_parent.m_progressIdx, cryptoName, (g+1)*0x200000);
         }
 
         /* Write new crypto content size */
+        uint64_t groupCount = m_curUser / 0x1F0000;
         uint64_t cryptContentSize = (groupCount * 0x200000) >> uint64_t(2);
         uint32_t cryptContentSizeBig = SBig(uint32_t(cryptContentSize));
-        ws = out.beginWriteStream(offset + 0x2BC);
+        ws = m_parent.getFileIO().beginWriteStream(m_baseOffset + 0x2BC);
         ws->write(&cryptContentSizeBig, 0x4);
 
         /* Write new H3 */
-        ws = out.beginWriteStream(offset + h3Off);
+        ws = m_parent.getFileIO().beginWriteStream(m_baseOffset + h3Off);
-        ws->write(h3, 0x18000);
+        ws->write(m_h3, 0x18000);
 
         /* Compute content hash and replace in TMD */
+        sha1nfo sha;
         sha1_init(&sha);
-        sha1_write(&sha, (char*)h3, 0x18000);
+        sha1_write(&sha, (char*)m_h3, 0x18000);
         memcpy(tmdData.get() + 0x1F4, sha1_result(&sha), 20);
 
         /* Same for content size */
@@ -759,10 +845,10 @@ public:
         }
         m_parent.m_progressCB(m_parent.m_progressIdx, bfName, attempts);
 
-        ws = out.beginWriteStream(offset + 0x2C0);
+        ws = m_parent.getFileIO().beginWriteStream(m_baseOffset + 0x2C0);
         ws->write(tmdData.get(), tmdSz);
 
-        return offset + dataOff + groupCount * 0x200000;
+        return m_baseOffset + dataOff + groupCount * 0x200000;
     }
 };
 
@@ -772,19 +858,17 @@ bool DiscBuilderWii::buildFromDirectory(const SystemChar* dirIn, const SystemCha
     size_t DISC_CAPACITY = m_dualLayer ? 0x1FB4E0000 : 0x118240000;
 
     PartitionBuilderWii& pb = static_cast<PartitionBuilderWii&>(*m_partitions[0]);
-    uint64_t filledSz = 0x200000;
+    uint64_t filledSz = pb.m_baseOffset;
-    std::unique_ptr<IFileIO> imgOut = NewFileIO(m_outPath);
+    m_fileIO->beginWriteStream();
-    imgOut->beginWriteStream();
 
-    /* Assemble cleartext data partition into temporary file */
+    if (!CheckFreeSpace(m_outPath, DISC_CAPACITY))
-    SystemString clearPath(m_outPath);
+    {
-    clearPath += _S(".cleardata");
+        LogModule.report(LogVisor::Error, _S("not enough free disk space for %s"), m_outPath);
-    m_fileIO = std::move(NewFileIO(clearPath));
-    if (!pb.buildFromDirectory(dirIn, dolIn, apploaderIn))
         return false;
+    }
 
-    /* Fakesign cleartext into output file */
+    /* Assemble image */
-    filledSz = pb.cryptAndFakesign(*imgOut, filledSz, partHeadIn);
+    filledSz = pb.buildFromDirectory(dirIn, dolIn, apploaderIn, partHeadIn);
     if (filledSz >= DISC_CAPACITY)
     {
         LogModule.report(LogVisor::FatalError, "data partition exceeds disc capacity");
@@ -795,21 +879,21 @@ bool DiscBuilderWii::buildFromDirectory(const SystemChar* dirIn, const SystemCha
     m_progressCB(m_progressIdx, "Finishing Disc", -1);
 
     /* Populate disc header */
-    std::unique_ptr<IFileIO::IWriteStream> ws = imgOut->beginWriteStream(0);
+    std::unique_ptr<IFileIO::IWriteStream> ws = m_fileIO->beginWriteStream(0);
     Header header(pb.getGameID(), pb.getGameTitle().c_str(), true, 0, 0, 0);
     header.write(*ws);
 
     /* Populate partition info */
-    ws = imgOut->beginWriteStream(0x40000);
+    ws = m_fileIO->beginWriteStream(0x40000);
     uint32_t vals[2] = {SBig(uint32_t(1)), SBig(uint32_t(0x40020 >> uint64_t(2)))};
     ws->write(vals, 8);
 
-    ws = imgOut->beginWriteStream(0x40020);
+    ws = m_fileIO->beginWriteStream(0x40020);
-    vals[0] = SBig(uint32_t(0x200000 >> uint64_t(2)));
+    vals[0] = SBig(uint32_t(pb.m_baseOffset >> uint64_t(2)));
     ws->write(vals, 4);
 
     /* Populate region info */
-    ws = imgOut->beginWriteStream(0x4E000);
+    ws = m_fileIO->beginWriteStream(0x4E000);
     const char* gameID = pb.getGameID();
     if (gameID[3] == 'P')
         vals[0] = SBig(uint32_t(2));
@@ -820,27 +904,24 @@ bool DiscBuilderWii::buildFromDirectory(const SystemChar* dirIn, const SystemCha
     ws->write(vals, 4);
 
     /* Make disc unrated */
-    ws = imgOut->beginWriteStream(0x4E010);
+    ws = m_fileIO->beginWriteStream(0x4E010);
     for (int i=0 ; i<16 ; ++i)
         ws->write("\x80", 1);
 
     /* Fill image to end */
-    ws = imgOut->beginWriteStream(filledSz);
+    ws = m_fileIO->beginWriteStream(filledSz);
     for (size_t i=0 ; i<DISC_CAPACITY-filledSz ; ++i)
         ws->write("\xff", 1);
 
-    /* Delete cleartext file */
-    Unlink(clearPath.c_str());
-
     return true;
 }
 
 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(std::move(std::unique_ptr<IFileIO>()), progressCB), m_outPath(outPath), m_dualLayer(dualLayer)
+: DiscBuilderBase(outPath, progressCB), m_dualLayer(dualLayer)
 {
     PartitionBuilderWii* partBuilder = new PartitionBuilderWii(*this, PartitionBuilderBase::Kind::Data,
-                                                               gameID, gameTitle);
+                                                               gameID, gameTitle, 0x200000);
     m_partitions.emplace_back(partBuilder);
 }