FST building/traversal in place

This commit is contained in:
Jack Andersen 2015-06-29 14:07:46 -10:00
parent 3f31ad21a0
commit ec8db1694b
8 changed files with 336 additions and 182 deletions

View File

@ -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)
: m_parent(parent), m_kind(isDir ? NODE_DIRECTORY : NODE_FILE), m_name(name) {}
Node(const IPartition& parent, const FSTNode& node, const char* 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;
void parseFST();
uint32_t m_dolOff;
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;
}
};
}

View File

@ -1,11 +1,10 @@
#ifndef __AES_HPP_
#define __AES_HPP_
#ifndef __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__

View File

@ -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<IPartReadStream> s = beginReadStream();
s->read(buf, 1024);
printf("");
std::unique_ptr<uint8_t[]> fst(new uint8_t[m_fstSz]);
s.seek(m_fstOff);
s.read(fst.get(), m_fstSz);
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;
}
}
}
}

View File

@ -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()

View File

@ -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;

View File

@ -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>();
}

View File

@ -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,76 +489,127 @@ void SoftwareAES::encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, si
#if __AES__
static inline __m128i AES_128_ASSIST (__m128i temp1, __m128i temp2)
#include <wmmintrin.h>
class NiAES : public IAES
{
__m128i temp3;
temp2 = _mm_shuffle_epi32 (temp2 ,0xff);
temp3 = _mm_slli_si128 (temp1, 0x4);
temp1 = _mm_xor_si128 (temp1, temp3);
temp3 = _mm_slli_si128 (temp3, 0x4);
temp1 = _mm_xor_si128 (temp1, temp3);
temp3 = _mm_slli_si128 (temp3, 0x4);
temp1 = _mm_xor_si128 (temp1, temp3);
temp1 = _mm_xor_si128 (temp1, temp2);
return temp1;
}
__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 NiAES::setKey(const uint8_t* key)
{
__m128i temp1, temp2;
static inline __m128i AES_128_ASSIST (__m128i temp1, __m128i temp2)
{
__m128i temp3;
temp2 = _mm_shuffle_epi32 (temp2 ,0xff);
temp3 = _mm_slli_si128 (temp1, 0x4);
temp1 = _mm_xor_si128 (temp1, temp3);
temp3 = _mm_slli_si128 (temp3, 0x4);
temp1 = _mm_xor_si128 (temp1, temp3);
temp3 = _mm_slli_si128 (temp3, 0x4);
temp1 = _mm_xor_si128 (temp1, temp3);
temp1 = _mm_xor_si128 (temp1, temp2);
return temp1;
}
temp1 = _mm_loadu_si128((__m128i*)key);
m_ekey[0] = temp1;
m_dkey[10] = temp1;
temp2 = _mm_aeskeygenassist_si128 (temp1,0x1);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[1] = temp1;
m_dkey[9] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x2);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[2] = temp1;
m_dkey[8] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x4);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[3] = temp1;
m_dkey[7] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x8);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[4] = temp1;
m_dkey[6] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x10);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[5] = temp1;
m_dkey[5] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x20);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[6] = temp1;
m_dkey[4] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x40);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[7] = temp1;
m_dkey[3] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x80);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[8] = temp1;
m_dkey[2] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x1b);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[9] = temp1;
m_dkey[1] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128 (temp1,0x36);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[10] = temp1;
m_dkey[0] = temp1;
void setKey(const uint8_t* key)
{
__m128i temp1, temp2;
}
temp1 = _mm_loadu_si128((__m128i*)key);
m_ekey[0] = temp1;
m_dkey[10] = temp1;
temp2 = _mm_aeskeygenassist_si128(temp1, 0x1);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[1] = temp1;
m_dkey[9] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x2);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[2] = temp1;
m_dkey[8] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x4);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[3] = temp1;
m_dkey[7] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x8);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[4] = temp1;
m_dkey[6] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x10);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[5] = temp1;
m_dkey[5] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x20);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[6] = temp1;
m_dkey[4] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x40);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[7] = temp1;
m_dkey[3] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x80);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[8] = temp1;
m_dkey[2] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x1b);
temp1 = AES_128_ASSIST(temp1, temp2);
m_ekey[9] = temp1;
m_dkey[1] = _mm_aesimc_si128(temp1);
temp2 = _mm_aeskeygenassist_si128(temp1, 0x36);
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
}
}

View File

@ -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;
}