Fixes and changes:

- Handles Split MFT
- Handles Windows XP formatted drives

8 files changed, 266 insertions, 28 deletions

diff --git a/src/drive.h b/src/drive.h
index bec704e..f0df9f1 100644
--- a/src/drive.h
+++ b/src/drive.h
@@ -27,6 +27,7 @@
 #include "usuals.h"
 
 const uint16 kSectorSize = 0x200;
+const uint64 kInvalidSector = 0xFFFFFFFFFFFFFFFF;
 struct DriveLock;
 
 struct PartitionInfo
diff --git a/src/misc.c b/src/misc.c
index 5ea67cd..8939b73 100644
--- a/src/misc.c
+++ b/src/misc.c
@@ -65,7 +65,7 @@ void* _refadd_dbg(void* pBuff)
 	{
 		// Increment the counter value
 		size_t* pMem = (size_t*)pBuff - 2;
-		assert(pMem[0] = kRefSig);
+		assert(pMem[0] == kRefSig);
 		pMem[1]++;
 	}
 
@@ -87,7 +87,7 @@ void _refrelease_dbg(void* pBuff)
 	{
 		// Decrement the counter value
 		size_t* pMem = (size_t*)pBuff - 2;
-		assert(pMem[0] = kRefSig);
+		assert(pMem[0] == kRefSig);
 
 		if(!--pMem[1])
 			free(pMem);
diff --git a/src/ntfs.c b/src/ntfs.c
index 28b4cca..82f6f86 100644
--- a/src/ntfs.c
+++ b/src/ntfs.c
@@ -50,11 +50,18 @@ NTFS_AttribHeader* NTFS_SearchAttribute(byte* pLocation, uint32 attrType, void*
 	return NULL;
 }
 
+byte* NTFS_GetAttributeList(NTFS_RecordHeader* pRecord)
+{
+  byte* pLocation = (byte*)pRecord;
+  ASSERT(pRecord->x_offUpdSeqArr != 0);
+  ASSERT(pRecord->x_offUpdSeqArr < 0x100);
+  pLocation += pRecord->x_offUpdSeqArr;
+  return pLocation;
+}
 
 NTFS_AttribHeader* NTFS_FindAttribute(NTFS_RecordHeader* pRecord, uint32 attrType, void* pEnd)
 {
-	byte* pLocation = (byte*)pRecord;
-	pLocation += kNTFS_RecHeaderLen;
+	byte* pLocation = NTFS_GetAttributeList(pRecord);
 
 	return NTFS_SearchAttribute(pLocation, attrType, pEnd, false);
 }
diff --git a/src/ntfsx.c b/src/ntfsx.c
index fe7d898..b2f8a9f 100644
--- a/src/ntfsx.c
+++ b/src/ntfsx.c
@@ -27,6 +27,7 @@
 #include "stdafx.h"
 #include "ntfs.h"
 #include "ntfsx.h"
+#include "scrounge.h"
 
 //////////////////////////////////////////////////////////////////////
 // Construction/Destruction
@@ -73,6 +74,14 @@ bool NTFS_Cluster::Read(PartitionInfo* pInfo, uint64 begSector, HANDLE hIn)
 	return true;
 }
 
+void NTFS_Cluster::Free()
+{ 
+  if(m_pCluster) 
+      refrelease(m_pCluster); 
+  
+  m_pCluster = NULL;
+}
+
 NTFS_Record::NTFS_Record(PartitionInfo* pInfo)
 {
 	m_pInfo	= (PartitionInfo*)refadd(pInfo);
@@ -257,3 +266,131 @@ bool NTFS_DataRun::Next()
 
 	return true;
 }
+
+NTFS_MFTMap::NTFS_MFTMap(PartitionInfo* pInfo)
+{
+	m_pInfo	= (PartitionInfo*)refadd(pInfo);
+  m_pBlocks = NULL;
+  m_count = 0;
+}
+
+NTFS_MFTMap::~NTFS_MFTMap()
+{
+	refrelease(m_pInfo);
+
+  if(m_pBlocks)
+    free(m_pBlocks);
+}
+
+bool NTFS_MFTMap::Load(NTFS_Record* pRecord, HANDLE hIn)
+{
+  bool bRet = TRUE;
+	NTFS_Attribute* pAttribData = NULL; // Data Attribute
+	NTFS_DataRun* pDataRun = NULL;		// Data runs for nonresident data
+
+  {
+    printf("[Processing MFT] ");
+
+    // TODO: Check here whether MFT has already been loaded
+  
+	  // Get the MFT's data
+	  pAttribData = pRecord->FindAttribute(kNTFS_DATA, hIn);
+	  if(!pAttribData) RET_ERROR(ERROR_NTFS_INVALID);
+
+    if(!pAttribData->GetHeader()->bNonResident)
+      RET_ERROR(ERROR_NTFS_INVALID);
+
+		pDataRun = pAttribData->GetDataRun();
+    if(!pDataRun)
+      RET_ERROR(ERROR_NTFS_INVALID);
+
+		NTFS_AttribNonResident* pNonRes = (NTFS_AttribNonResident*)pAttribData->GetHeader();
+    
+    if(m_pBlocks)
+    {
+      free(m_pBlocks);
+      m_pBlocks = NULL;
+    }
+
+    m_count = 0;
+    uint32 allocated = 0;
+
+		// Now loop through the data run
+		if(pDataRun->First())
+		{
+			do
+			{
+        if(pDataRun->m_bSparse)
+          RET_ERROR(ERROR_NTFS_INVALID);
+
+        if(m_count >= allocated)
+        {
+          allocated += 16;
+          m_pBlocks = (NTFS_Block*)realloc(m_pBlocks, allocated * sizeof(NTFS_Block));
+          if(!m_pBlocks)
+      			RET_FATAL(ERROR_NOT_ENOUGH_MEMORY);
+        }
+
+        uint64 length = pDataRun->m_numClusters * ((m_pInfo->clusterSize * kSectorSize) / kNTFS_RecordLen);
+        if(length == 0)
+          continue;
+
+        uint64 firstSector = (pDataRun->m_firstCluster * m_pInfo->clusterSize) + m_pInfo->firstSector;
+        if(firstSector >= m_pInfo->lastSector)
+          continue;
+
+        m_pBlocks[m_count].length = length;
+        m_pBlocks[m_count].firstSector = firstSector;
+        m_count++;
+			} 
+			while(pDataRun->Next());
+    }
+
+    bRet = true;
+    ::SetLastError(ERROR_SUCCESS);
+  }
+
+clean_up:
+
+  if(pAttribData)
+    delete pAttribData;
+  if(pDataRun)
+    delete pDataRun;
+
+  return bRet;
+}
+
+uint64 NTFS_MFTMap::GetLength()
+{
+  uint64 length = 0;
+
+  for(uint32 i = 0; i < m_count; i++)
+    length += m_pBlocks[i].length;
+
+  return length;
+}
+
+uint64 NTFS_MFTMap::SectorForIndex(uint64 index)
+{
+  for(uint32 i = 0; i < m_count; i++)
+  {
+    NTFS_Block* p = m_pBlocks + i;
+
+    if(index > p->length)
+    {
+      index -= p->length;
+    }
+    else
+    {
+      uint64 sector = index * (kNTFS_RecordLen / kSectorSize);
+      sector += p->firstSector;
+
+      if(sector > m_pInfo->lastSector)
+        return kInvalidSector;
+
+      return sector;
+    }
+  }
+
+  return kInvalidSector;
+}
diff --git a/src/ntfsx.h b/src/ntfsx.h
index 586b36b..836f010 100644
--- a/src/ntfsx.h
+++ b/src/ntfsx.h
@@ -64,8 +64,8 @@ public:
 
 	bool New(PartitionInfo* pInfo);
 	bool Read(PartitionInfo* pInfo, uint64 begSector, HANDLE hIn);
-	void Free()
-		{ if(m_pCluster) refrelease(m_pCluster); }
+	void Free();
+
 	
 	uint32 m_cbCluster;
 	byte* m_pCluster;
@@ -95,7 +95,8 @@ protected:
 	uint32 m_cbMem;
 };
 
-class NTFS_Record : public NTFS_Cluster
+class NTFS_Record : 
+  public NTFS_Cluster
 {
 public:
 	NTFS_Record(PartitionInfo* pInfo);
@@ -111,4 +112,28 @@ protected:
 	PartitionInfo* m_pInfo;
 };
 
+class NTFS_MFTMap
+{
+public:
+  NTFS_MFTMap(PartitionInfo* pInfo);
+  ~NTFS_MFTMap();
+
+  bool Load(NTFS_Record* pRecord, HANDLE hIn);
+
+  uint64 GetLength();
+  uint64 SectorForIndex(uint64 index);
+
+protected:
+  PartitionInfo* m_pInfo;
+
+  struct NTFS_Block
+  {
+    uint64 firstSector;   // relative to the entire drive
+    uint64 length;        // length in MFT records
+  };
+
+  NTFS_Block* m_pBlocks;
+  uint32 m_count;
+};
+
 #endif // !defined(AFX_NTFSX__9363C7D2_D3CC_4D49_BEE0_27AD025670F2__INCLUDED_)
diff --git a/src/scrounge.c b/src/scrounge.c
index 8841aa9..2cb0d71 100644
--- a/src/scrounge.c
+++ b/src/scrounge.c
@@ -31,12 +31,6 @@
 #include "locks.h"
 #include "scrounge.h"
 
-#define RET_ERROR(l) { ::SetLastError(l); bRet = TRUE; goto clean_up; }
-#define PASS_ERROR() {bRet = TRUE; goto clean_up; }
-#define RET_FATAL(l) { ::SetLastError(l); bRet = FALSE; goto clean_up; }
-#define PASS_FATAL() {bRet = FALSE; goto clean_up; }
-
-
 // ----------------------------------------------------------------------
 // Process a potential MFT Record. For directories create the directory
 // and for files write out the file.
@@ -45,7 +39,7 @@
 // hIn is an open drive handle 
 // pInfo is partition info about hIn (needs to be a ref counted pointer)
 
-BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 mftRecord, HANDLE hIn)
+BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 sector, NTFS_MFTMap* map, HANDLE hIn)
 {
 	// Declare data that needs cleaning up first
 	BOOL bRet = TRUE;					// Return value
@@ -54,15 +48,16 @@ BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 mftRecord, HANDLE hIn)
 	NTFS_Attribute* pAttribData = NULL; // Data Attribute
 	NTFS_DataRun* pDataRun = NULL;		// Data runs for nonresident data
 
+  ASSERT(sector != kInvalidSector);
+
 
 	// Tracks whether or not we output a file
 	bool bFile = false;
 
 	{
-
 		// Read the MFT record
 		NTFS_Record record(pInfo);
-		if(!record.Read(mftRecord, hIn))
+		if(!record.Read(sector, hIn))
 			PASS_ERROR();
 
 		NTFS_RecordHeader* pRecord = record.GetHeader();
@@ -79,7 +74,7 @@ BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 mftRecord, HANDLE hIn)
 		FILETIME ftModified;
 		FILETIME ftAccessed;
 		DWORD fileAttrib = 0;
-		uint64 mftParent = 0;
+		uint64 mftParent = kInvalidSector;
 		byte* pResidentData = NULL;
 
 
@@ -124,7 +119,8 @@ BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 mftRecord, HANDLE hIn)
 					fileAttrib |= FILE_ATTRIBUTE_SYSTEM;
 
 				// Parent Directory
-				mftParent = NTFS_RefToSector(*pInfo, pFileName->refParent);
+        if(map)
+          mftParent = map->SectorForIndex(pFileName->refParent & 0xFFFFFFFFFFFF);
 
 				// Namespace
 				nameSpace = pFileName->nameSpace;
@@ -138,20 +134,35 @@ BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 mftRecord, HANDLE hIn)
 			RET_ERROR(ERROR_NTFS_INVALID);
 
 
-
 		// Check if it's the root
 		// If so then bumm out cuz we don't want to have anything to do with it
-		if(mftRecord == mftParent ||	// Root is it's own parent
-		   !wcscmp(fileName, L".") ||   // Or is called '.'
-		   !wcscmp(fileName, kNTFS_MFTName)) // Or it's the MFT
+		if(sector == mftParent ||	// Root is it's own parent
+		   !wcscmp(fileName, L"."))   // Or is called '.'
 		    RET_ERROR(ERROR_SUCCESS);
 
 
+    // If it's the MFT then we read that in
+    if(!wcscmp(fileName, kNTFS_MFTName))
+    {
+      if(map)
+      {
+        printf("[Processing MFT] ");
 
-		// Create Parent folders
-		if(!ProcessMFTRecord(pInfo, mftParent, hIn))
-			PASS_FATAL();
+        // We found the MFT, let's load it up
+        if(!map->Load(&record, hIn))
+          PASS_FATAL();
+      }
 
+      RET_ERROR(ERROR_SUCCESS);
+    }
+
+
+    if(mftParent != kInvalidSector)
+    {
+		  // Create Parent folders
+		  if(!ProcessMFTRecord(pInfo, mftParent, map, hIn))
+			  PASS_FATAL();
+    }
 
 
 		// If it's a folder then create it
@@ -381,6 +392,58 @@ void PrintLastError()
 
 BOOL ScroungeMFTRecords(PartitionInfo* pInfo, HANDLE hIn)
 {
+	uint64 numRecords = 0;
+
+	// Save current directory away
+	TCHAR curDir[MAX_PATH + 1];
+	GetCurrentDirectory(MAX_PATH, curDir);
+
+  NTFS_MFTMap map(pInfo);
+
+  // Try and find the MFT at the given location
+  uint64 sector = pInfo->offMFT + pInfo->firstSector;
+
+  if(sector < pInfo->lastSector)
+  {
+    BOOL bRet = ProcessMFTRecord(pInfo, sector, &map, hIn);
+
+		fputc('\n', stdout);
+
+		if(!bRet)
+			return FALSE;
+  }
+
+  uint64 length = map.GetLength();
+  if(length == 0)
+  {
+    printf("[MFT not found at specified location. No folder structure]\n");
+    return ScroungeRawRecords(pInfo, hIn);
+  }
+
+  for(uint64 i = 1; i < length; i += (kNTFS_RecordLen / kSectorSize))
+  {
+  	// Move to right output directory
+	  SetCurrentDirectory(curDir);
+
+    // TODO: Warn when invalid
+    uint64 sector = map.SectorForIndex(i);
+    if(sector == kInvalidSector)
+      continue;
+
+		// Then process it
+		BOOL bRet = ProcessMFTRecord(pInfo, sector, &map, hIn);
+
+		fputc('\n', stdout);
+
+		if(!bRet)
+			return FALSE;
+	}
+
+	return TRUE;
+}
+
+BOOL ScroungeRawRecords(PartitionInfo* pInfo, HANDLE hIn)
+{
 	byte buffSec[kSectorSize];
 	DWORD dwDummy = 0;
 
@@ -418,7 +481,7 @@ BOOL ScroungeMFTRecords(PartitionInfo* pInfo, HANDLE hIn)
 			SetCurrentDirectory(curDir);
 
 			// Then process it
-			BOOL bRet = ProcessMFTRecord(pInfo, sec, hIn);
+			BOOL bRet = ProcessMFTRecord(pInfo, sec, NULL, hIn);
 
 			fputc('\n', stdout);
 
diff --git a/src/scrounge.h b/src/scrounge.h
index dd908ea..b425db9 100644
--- a/src/scrounge.h
+++ b/src/scrounge.h
@@ -23,9 +23,14 @@
 #ifndef __SCROUNGE_H__
 #define __SCROUNGE_H__
 
+#define RET_ERROR(l) { ::SetLastError(l); bRet = TRUE; goto clean_up; }
+#define PASS_ERROR() {bRet = TRUE; goto clean_up; }
+#define RET_FATAL(l) { ::SetLastError(l); bRet = FALSE; goto clean_up; }
+#define PASS_FATAL() {bRet = FALSE; goto clean_up; }
 
-BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 mftRecord, HANDLE hIn);
+BOOL ProcessMFTRecord(PartitionInfo* pInfo, uint64 sector, NTFS_MFTMap* map, HANDLE hIn);
 BOOL ScroungeMFTRecords(PartitionInfo* pInfo, HANDLE hIn);
+BOOL ScroungeRawRecords(PartitionInfo* pInfo, HANDLE hIn);
 void PrintLastError();
 
 
diff --git a/src/usuals.h b/src/usuals.h
index 4193675..8961592 100644
--- a/src/usuals.h
+++ b/src/usuals.h
@@ -23,7 +23,7 @@
 #ifndef __USUALS_H__20010822
 #define __USUALS_H__20010822
 
-#include <debug.h>
+#include <win32/debug.h>
 
 typedef unsigned __int64 uint64;
 typedef unsigned long uint32;