FST building/traversal in place

include/DiscBase.hpp

@@ -4,6 +4,7 @@
 #include <vector>
 #include <memory>
 #include <string>
+#include <stdio.h>
 #include "Util.hpp"
 #include "IDiscIO.hpp"
 #include "IFileIO.hpp"
@@ -35,22 +36,9 @@ public:
         }
     };
 
-    struct FSTNode
-    {
-        uint32_t typeAndNameOff;
-        uint32_t offset;
-        uint32_t length;
-        FSTNode(IDiscIO::IReadStream& s)
-        {
-            s.read(this, 12);
-            typeAndNameOff = SBig(typeAndNameOff);
-            offset = SBig(offset);
-            length = SBig(length);
-        }
-    };
-
     struct IPartReadStream
     {
+        virtual void seek(size_t offset, int whence=SEEK_SET)=0;
         virtual size_t read(void* buf, size_t length)=0;
     };
 
@@ -59,6 +47,18 @@ public:
         virtual size_t write(void* buf, size_t length)=0;
     };
 
+    class FSTNode
+    {
+        uint32_t typeAndNameOffset;
+        uint32_t offset;
+        uint32_t length;
+    public:
+        inline bool isDir() const {return SBig(typeAndNameOffset) >> 24;}
+        inline uint32_t getNameOffset() const {return SBig(typeAndNameOffset) & 0xffffff;}
+        inline uint32_t getOffset() const {return SBig(offset);}
+        inline uint32_t getLength() const {return SBig(length);}
+    };
+
     class IPartition
     {
     public:
@@ -77,18 +77,27 @@ public:
                 NODE_DIRECTORY
             };
         private:
+            friend class IPartition;
             const IPartition& m_parent;
             Kind m_kind;
-            std::string m_name;
 
             std::unique_ptr<IFileIO> m_hddFile;
             size_t m_discOffset;
             size_t m_discLength;
+            std::string m_name;
+
+            std::vector<Node>::const_iterator m_childrenBegin;
+            std::vector<Node>::const_iterator m_childrenEnd;
 
         public:
-            Node(const IPartition& parent, bool isDir, const char* name)
+            Node(const IPartition& parent, const FSTNode& node, const char* name)
-            : m_parent(parent), m_kind(isDir ? NODE_DIRECTORY : NODE_FILE), m_name(name) {}
+            : m_parent(parent),
+              m_kind(node.isDir() ? NODE_DIRECTORY : NODE_FILE),
+              m_discOffset(node.getOffset()),
+              m_discLength(node.getLength()),
+              m_name(name) {}
             inline Kind getKind() const {return m_kind;}
+            inline const std::string& getName() const {return m_name;}
             std::unique_ptr<IPartReadStream> beginReadStream() const
             {
                 if (m_kind != NODE_FILE)
@@ -98,10 +107,37 @@ public:
                 }
                 return m_parent.beginReadStream(m_discOffset);
             }
+            inline std::vector<Node>::const_iterator rawBegin() const {return m_childrenBegin;}
+            inline std::vector<Node>::const_iterator rawEnd() const {return m_childrenEnd;}
+
+            class DirectoryIterator : std::iterator<std::forward_iterator_tag, const Node>
+            {
+                friend class Node;
+                std::vector<Node>::const_iterator m_it;
+                DirectoryIterator(std::vector<Node>::const_iterator&& it)
+                : m_it(std::move(it)) {}
+            public:
+                inline bool operator!=(const DirectoryIterator& other) {return m_it != other.m_it;}
+                inline DirectoryIterator& operator++()
+                {
+                    if (m_it->m_kind == NODE_DIRECTORY)
+                        m_it = m_it->rawEnd();
+                    else
+                        ++m_it;
+                    return *this;
+                }
+                inline const Node& operator*() {return *m_it;}
+            };
+            inline DirectoryIterator begin() const {return DirectoryIterator(rawBegin());}
+            inline DirectoryIterator end() const {return DirectoryIterator(rawEnd());}
         };
     protected:
-        std::vector<Node> m_files;
+        uint32_t m_dolOff;
-        void parseFST();
+        uint32_t m_fstOff;
+        uint32_t m_fstSz;
+        uint32_t m_apploaderOff;
+        std::vector<Node> m_nodes;
+        void parseFST(IPartReadStream& s);
 
         const DiscBase& m_parent;
         Kind m_kind;
@@ -111,6 +147,7 @@ public:
         : m_parent(parent), m_kind(kind), m_offset(offset) {}
         inline Kind getKind() const {return m_kind;}
         virtual std::unique_ptr<IPartReadStream> beginReadStream(size_t offset=0) const=0;
+        inline const Node& getFSTRoot() const {return m_nodes[0];}
     };
 
 protected:
@@ -120,8 +157,22 @@ protected:
 public:
     DiscBase(std::unique_ptr<IDiscIO>&& dio);
     virtual bool commit()=0;
-    inline Header getHeader() const {return m_header;}
+    inline const Header& getHeader() const {return m_header;}
     inline const IDiscIO& getDiscIO() const {return *m_discIO.get();}
+    inline const IPartition* getDataPartition() const
+    {
+        for (const std::unique_ptr<IPartition>& part : m_partitions)
+            if (part->getKind() == IPartition::PART_DATA)
+                return part.get();
+        return nullptr;
+    }
+    inline const IPartition* getUpdatePartition() const
+    {
+        for (const std::unique_ptr<IPartition>& part : m_partitions)
+            if (part->getKind() == IPartition::PART_UPDATE)
+                return part.get();
+        return nullptr;
+    }
 };
 
 }

include/aes.hpp

@@ -1,11 +1,10 @@
-#ifndef __AES_HPP_
+#ifndef __AES_HPP__
-#define __AES_HPP_
+#define __AES_HPP__
 
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
 #include <memory>
-#include <cpuid.h>
 
 namespace NOD
 {
@@ -18,97 +17,8 @@ public:
     virtual void setKey(const uint8_t* key)=0;
 };
 
-class SoftwareAES : public IAES
-{
-protected:
-    uint8_t fbsub[256];
-    uint8_t rbsub[256];
-    uint8_t ptab[256], ltab[256];
-    uint32_t ftable[256];
-    uint32_t rtable[256];
-    uint32_t rco[30];
-
-    /* Parameter-dependent data */
-
-    int Nk, Nb, Nr;
-    uint8_t fi[24], ri[24];
-    uint32_t fkey[120];
-    uint32_t rkey[120];
-
-
-    uint8_t bmul(uint8_t x, uint8_t y);
-    uint32_t SubByte(uint32_t a);
-    uint8_t product(uint32_t x, uint32_t y);
-    uint32_t InvMixCol(uint32_t x);
-    uint8_t ByteSub(uint8_t x);
-    void gentables(void);
-    void gkey(int nb, int nk, const uint8_t* key);
-    void _encrypt(uint8_t* buff);
-    void _decrypt(uint8_t* buff);
-
-public:
-    void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len);
-    void decrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len);
-    void setKey(const uint8_t* key);
-};
-
-#if __AES__
-
-#include <wmmintrin.h>
-
-class NiAES : public IAES
-{
-    __m128i m_ekey[11];
-    __m128i m_dkey[11];
-public:
-    void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len)
-    {
-        __m128i feedback,data;
-        size_t i,j;
-        if (len%16)
-            len = len/16+1;
-        else
-            len /= 16;
-        feedback = _mm_loadu_si128((__m128i*)iv);
-        for (i=0 ; i<len ; i++)
-        {
-            data = _mm_loadu_si128(&((__m128i*)inbuf)[i]);
-            feedback = _mm_xor_si128(data, feedback);
-            feedback = _mm_xor_si128(feedback, m_ekey[0]);
-            for(j=1 ; j<10 ; j++)
-                feedback = _mm_aesenc_si128(feedback, m_ekey[j]);
-            feedback = _mm_aesenclast_si128(feedback, m_ekey[j]);
-            _mm_storeu_si128(&((__m128i*)outbuf)[i], feedback);
-        }
-    }
-    void decrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len)
-    {
-        __m128i data,feedback,last_in;
-        size_t i,j;
-        if (len%16)
-            len = len/16+1;
-        else
-            len /= 16;
-        feedback = _mm_loadu_si128((__m128i*)iv);
-        for (i=0 ; i<len ; i++)
-        {
-            last_in=_mm_loadu_si128(&((__m128i*)inbuf)[i]);
-            data = _mm_xor_si128(last_in, m_dkey[0]);
-            for(j=1 ; j<10 ; j++)
-                data = _mm_aesdec_si128(data, m_dkey[j]);
-            data = _mm_aesdeclast_si128(data, m_dkey[j]);
-            data = _mm_xor_si128(data, feedback);
-            _mm_storeu_si128(&((__m128i*)outbuf)[i], data);
-            feedback = last_in;
-        }
-    }
-    void setKey(const uint8_t* key);
-};
-
-#endif
-
 std::unique_ptr<IAES> NewAES();
 
 }
 
-#endif //__AES_HPP_
+#endif //__AES_HPP__

lib/DiscBase.cpp

@@ -8,12 +8,39 @@ DiscBase::DiscBase(std::unique_ptr<IDiscIO>&& dio)
 {
 }
 
-void DiscBase::IPartition::parseFST()
+void DiscBase::IPartition::parseFST(IPartReadStream& s)
 {
-    char buf[1024];
+    std::unique_ptr<uint8_t[]> fst(new uint8_t[m_fstSz]);
-    std::unique_ptr<IPartReadStream> s = beginReadStream();
+    s.seek(m_fstOff);
-    s->read(buf, 1024);
+    s.read(fst.get(), m_fstSz);
-    printf("");
+
+    const FSTNode* nodes = (FSTNode*)fst.get();
+
+    /* Root node indicates the count of all contained nodes */
+    uint32_t nodeCount = nodes[0].getLength();
+    const char* names = (char*)fst.get() + 12 * nodeCount;
+    m_nodes.clear();
+    m_nodes.reserve(nodeCount);
+
+    /* Construct nodes */
+    for (uint32_t n=0 ; n<nodeCount ; ++n)
+    {
+        const FSTNode& node = nodes[n];
+        m_nodes.emplace_back(*this, node, n ? names + node.getNameOffset() : "<root>");
+    }
+
+    /* Setup dir-child iterators */
+    for (std::vector<Node>::iterator it=m_nodes.begin();
+         it != m_nodes.end();
+         ++it)
+    {
+        Node& node = *it;
+        if (node.m_kind == Node::NODE_DIRECTORY)
+        {
+            node.m_childrenBegin = it + 1;
+            node.m_childrenEnd = m_nodes.begin() + node.m_discLength;
+        }
+    }
 }
 
 }

lib/DiscGCN.cpp

@@ -3,10 +3,50 @@
 namespace NOD
 {
 
+class PartitionGCN : public DiscBase::IPartition
+{
+public:
+    PartitionGCN(const DiscGCN& parent, Kind kind, size_t offset)
+    : IPartition(parent, kind, offset)
+    {
+        /* GCN-specific header reads */
+        std::unique_ptr<DiscBase::IPartReadStream> s = beginReadStream(0x420);
+        uint32_t vals[3];
+        s->read(vals, 12);
+        m_dolOff = SBig(vals[0]);
+        m_fstOff = SBig(vals[1]);
+        m_fstSz = SBig(vals[2]);
+
+        /* Yay files!! */
+        parseFST(*s.get());
+    }
+
+    class PartReadStream : public DiscBase::IPartReadStream
+    {
+        const PartitionGCN& m_parent;
+        std::unique_ptr<IDiscIO::IReadStream> m_dio;
+
+    public:
+        PartReadStream(const PartitionGCN& parent, size_t offset)
+        : m_parent(parent)
+        {m_dio = m_parent.m_parent.getDiscIO().beginReadStream(offset);}
+        void seek(size_t offset, int whence)
+        {m_dio->seek(offset, whence);}
+        size_t read(void* buf, size_t length)
+        {return m_dio->read(buf, length);}
+    };
+
+    std::unique_ptr<DiscBase::IPartReadStream> beginReadStream(size_t offset) const
+    {
+        return std::unique_ptr<DiscBase::IPartReadStream>(new PartReadStream(*this, offset));
+    }
+};
+
 DiscGCN::DiscGCN(std::unique_ptr<IDiscIO>&& dio)
 : DiscBase(std::move(dio))
 {
-
+    /* One lone partition for GCN */
+    m_partitions.emplace_back(new PartitionGCN(*this, IPartition::PART_DATA, 0));
 }
 
 bool DiscGCN::commit()

lib/DiscWii.cpp

@@ -213,7 +213,16 @@ public:
         aes->setKey(COMMON_KEY);
         aes->decrypt(iv, m_ticket.encKey, m_decKey, 16);
 
-        parseFST();
+        /* Wii-specific header reads (now using title key to decrypt) */
+        std::unique_ptr<DiscBase::IPartReadStream> ds = beginReadStream(0x420);
+        uint32_t vals[3];
+        ds->read(vals, 12);
+        m_dolOff = SBig(vals[0]) << 2;
+        m_fstOff = SBig(vals[1]) << 2;
+        m_fstSz = SBig(vals[2]) << 2;
+
+        /* Yay files!! */
+        parseFST(*ds.get());
     }
 
     class PartReadStream : public DiscBase::IPartReadStream
@@ -240,6 +249,17 @@ public:
             size_t block = m_offset / 0x7c00;
             m_dio = m_parent.m_parent.getDiscIO().beginReadStream(m_baseOffset + block * 0x8000);
         }
+        void seek(size_t offset, int whence)
+        {
+            if (whence == SEEK_SET)
+                m_offset = offset;
+            else if (whence == SEEK_CUR)
+                m_offset += offset;
+            else
+                return;
+            size_t block = m_offset / 0x7c00;
+            m_dio->seek(m_baseOffset + block * 0x8000);
+        }
         size_t read(void* buf, size_t length)
         {
             size_t cacheOffset = m_offset % 0x7c00;

lib/NODLib.cpp

@@ -52,6 +52,7 @@ std::unique_ptr<DiscBase> OpenDiscFromImage(const char* path, bool& isWii)
 
     if (!discIO)
     {
+        fclose(fp);
         throw std::runtime_error("'" + std::string(path) + "' is not a valid image");
         return std::unique_ptr<DiscBase>();
     }

lib/aes.cpp

@@ -12,6 +12,7 @@
 #include <stdio.h>
 //#include <stdlib.h>
 #include <string.h>
+#include <cpuid.h>
 
 namespace NOD
 {
@@ -30,13 +31,13 @@ namespace NOD
 
 static const uint8_t InCo[4] = {0xB, 0xD, 0x9, 0xE}; /* Inverse Coefficients */
 
-static uint32_t pack(const uint8_t* b)
+static inline uint32_t pack(const uint8_t* b)
 {
     /* pack bytes into a 32-bit Word */
     return ((uint32_t)b[3] << 24) | ((uint32_t)b[2] << 16) | ((uint32_t)b[1] << 8) | (uint32_t)b[0];
 }
 
-static void unpack(uint32_t a, uint8_t* b)
+static inline void unpack(uint32_t a, uint8_t* b)
 {
     /* unpack bytes from a word */
     b[0] = (uint8_t)a;
@@ -45,7 +46,7 @@ static void unpack(uint32_t a, uint8_t* b)
     b[3] = (uint8_t)(a >> 24);
 }
 
-static uint8_t xtime(uint8_t a)
+static inline uint8_t xtime(uint8_t a)
 {
     uint8_t b;
 
@@ -57,6 +58,40 @@ static uint8_t xtime(uint8_t a)
     return a;
 }
 
+class SoftwareAES : public IAES
+{
+protected:
+    uint8_t fbsub[256];
+    uint8_t rbsub[256];
+    uint8_t ptab[256], ltab[256];
+    uint32_t ftable[256];
+    uint32_t rtable[256];
+    uint32_t rco[30];
+
+    /* Parameter-dependent data */
+
+    int Nk, Nb, Nr;
+    uint8_t fi[24], ri[24];
+    uint32_t fkey[120];
+    uint32_t rkey[120];
+
+
+    uint8_t bmul(uint8_t x, uint8_t y);
+    uint32_t SubByte(uint32_t a);
+    uint8_t product(uint32_t x, uint32_t y);
+    uint32_t InvMixCol(uint32_t x);
+    uint8_t ByteSub(uint8_t x);
+    void gentables(void);
+    void gkey(int nb, int nk, const uint8_t* key);
+    void _encrypt(uint8_t* buff);
+    void _decrypt(uint8_t* buff);
+
+public:
+    void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len);
+    void decrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len);
+    void setKey(const uint8_t* key);
+};
+
 uint8_t SoftwareAES::bmul(uint8_t x, uint8_t y)
 {
     /* x.y= AntiLog(Log(x) + Log(y)) */
@@ -454,6 +489,55 @@ void SoftwareAES::encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, si
 
 #if __AES__
 
+#include <wmmintrin.h>
+
+class NiAES : public IAES
+{
+    __m128i m_ekey[11];
+    __m128i m_dkey[11];
+public:
+    void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len)
+    {
+        __m128i feedback,data;
+        size_t i,j;
+        if (len%16)
+            len = len/16+1;
+        else
+            len /= 16;
+        feedback = _mm_loadu_si128((__m128i*)iv);
+        for (i=0 ; i<len ; i++)
+        {
+            data = _mm_loadu_si128(&((__m128i*)inbuf)[i]);
+            feedback = _mm_xor_si128(data, feedback);
+            feedback = _mm_xor_si128(feedback, m_ekey[0]);
+            for (j=1 ; j<10 ; j++)
+                feedback = _mm_aesenc_si128(feedback, m_ekey[j]);
+            feedback = _mm_aesenclast_si128(feedback, m_ekey[j]);
+            _mm_storeu_si128(&((__m128i*)outbuf)[i], feedback);
+        }
+    }
+    void decrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len)
+    {
+        __m128i data,feedback,last_in;
+        size_t i,j;
+        if (len%16)
+            len = len/16+1;
+        else
+            len /= 16;
+        feedback = _mm_loadu_si128((__m128i*)iv);
+        for (i=0 ; i<len ; i++)
+        {
+            last_in=_mm_loadu_si128(&((__m128i*)inbuf)[i]);
+            data = _mm_xor_si128(last_in, m_dkey[0]);
+            for (j=1 ; j<10 ; j++)
+                data = _mm_aesdec_si128(data, m_dkey[j]);
+            data = _mm_aesdeclast_si128(data, m_dkey[j]);
+            data = _mm_xor_si128(data, feedback);
+            _mm_storeu_si128(&((__m128i*)outbuf)[i], data);
+            feedback = last_in;
+        }
+    }
+
     static inline __m128i AES_128_ASSIST (__m128i temp1, __m128i temp2)
     {
         __m128i temp3;
@@ -468,7 +552,7 @@ static inline __m128i AES_128_ASSIST (__m128i temp1, __m128i temp2)
         return temp1;
     }
 
-void NiAES::setKey(const uint8_t* key)
+    void setKey(const uint8_t* key)
     {
         __m128i temp1, temp2;
 
@@ -515,15 +599,17 @@ void NiAES::setKey(const uint8_t* key)
         temp1 = AES_128_ASSIST(temp1, temp2);
         m_ekey[10] = temp1;
         m_dkey[0] = temp1;
-
     }
+};
 
 #endif
 
 static int HAS_AES_NI = -1;
 
+
 std::unique_ptr<IAES> NewAES()
 {
+#if __AES__
     if (HAS_AES_NI == -1)
     {
         unsigned int a,b,c,d;
@@ -537,6 +623,9 @@ std::unique_ptr<IAES> NewAES()
         return std::unique_ptr<IAES>(new NiAES);
     else
         return std::unique_ptr<IAES>(new SoftwareAES);
+#else
+    return std::unique_ptr<IAES>(new SoftwareAES);
+#endif
 }
 
 }

main.cpp

@@ -14,6 +14,22 @@ int main(int argc, char* argv[])
     if (!disc)
         return -1;
 
+    const NOD::DiscBase::IPartition* dataPart = disc->getDataPartition();
+    if (dataPart)
+    {
+        for (const NOD::DiscBase::IPartition::Node& node : dataPart->getFSTRoot())
+        {
+            if (node.getKind() == NOD::DiscBase::IPartition::Node::NODE_FILE)
+                printf("FILE: %s\n", node.getName().c_str());
+            else if (node.getKind() == NOD::DiscBase::IPartition::Node::NODE_DIRECTORY)
+            {
+                printf("DIR: %s\n", node.getName().c_str());
+                for (const NOD::DiscBase::IPartition::Node& subnode : node)
+                    printf("SUBFILE: %s\n", subnode.getName().c_str());
+            }
+        }
+    }
+
     return 0;
 }