Fixes. Updates

7 files changed, 80 insertions, 33 deletions

diff --git a/src/compat.h b/src/compat.h
index a73637f..c2b4128 100644
--- a/src/compat.h
+++ b/src/compat.h
@@ -213,7 +213,6 @@ void* reallocf(void* p, size_t sz);
   #endif
 #endif
 
-
 /*
  * Depending on the OS we use different width characters
  * for file names and file access. Before enabling wide
@@ -223,6 +222,11 @@ void* reallocf(void* p, size_t sz);
  */
 
 #ifdef _WIN32
+
+  #ifdef SIZEOF_WHCHAR_T != 2
+    #error Incompatible size of wchar_t
+  #endif
+
   /* On windows we use UCS2 */
   typedef wchar_t fchar_t;
   #define FC_WIDE 1
diff --git a/src/list.c b/src/list.c
index fc30ac0..faeeecd 100644
--- a/src/list.c
+++ b/src/list.c
@@ -38,6 +38,7 @@ int printNTFSInfo(int dd, uint64 tblSector)
   size_t sz;
   ntfs_bootsector* boot;
 
+  /* TODO: Check the range and don't print if out of range */
 	pos = SECTOR_TO_BYTES(tblSector);
 
   if(lseek64(dd, pos, SEEK_SET) == -1)
@@ -130,6 +131,7 @@ void scroungeListDrive(char* drive)
   if(dd == -1)
     err(1, "couldn't open drive: %s", drive);
 
+  printf(kPrintData);
   printf(kPrintDrivePath, drive);
   printPartitionInfo(dd, 0);
   close(dd);
diff --git a/src/main.c b/src/main.c
index 9f6e4a2..ef5be8e 100644
--- a/src/main.c
+++ b/src/main.c
@@ -44,7 +44,7 @@ usage: scrounge [-d drive] [-m mftoffset] [-c clustersize] [-o outdir] start end
 #else /* Not WIN32 */
 
 const char kPrintHelp[]       = "\
-usage: scrounge -l                                                   \n\
+usage: scrounge -l disk                                              \n\
   List all drive partition information.                              \n\
                                                                      \n\
 usage: scrounge -s disk                                              \n\
@@ -68,7 +68,7 @@ usage: scrounge [-m mftoffset] [-c clustersize] [-o outdir] disk start end  \n\
 #define MODE_SEARCH   3 
 
 /* Forward decls */
-int usage();
+void usage();
 
 int main(int argc, char* argv[])
 {
@@ -126,6 +126,7 @@ int main(int argc, char* argv[])
 #endif
 
     /* help mode */
+    case '?':
     case 'h':
       usage();
       break;
@@ -171,9 +172,9 @@ int main(int argc, char* argv[])
       break;
 
     default:
-      ASSERT(false);
+      usage();
       break;
-    }
+    };
 
     if(mode != MODE_SCROUNGE && mode != 0)
       break;
@@ -261,7 +262,7 @@ int main(int argc, char* argv[])
   return 0;
 }
 
-int usage()
+void usage()
 {
   fprintf(stderr, "%s", kPrintHelp);
   exit(2);
diff --git a/src/ntfs.h b/src/ntfs.h
index f35fac0..ae7ea6e 100644
--- a/src/ntfs.h
+++ b/src/ntfs.h
@@ -23,6 +23,9 @@
 #include "usuals.h"
 #include "stddef.h"
 #include "drive.h"
+#include "compat.h"
+
+typedef int16 ntfs_char;
 
 #pragma pack(1)
 
@@ -200,7 +203,7 @@ bool ntfs_isbetternamespace(byte n1, byte n2);
 bool ntfs_dofixups(byte* cluster, uint32 size);
 
 /* TODO: Move these declarations elsewhere */
-char* unicode_transcode16to8(const wchar_t* src, size_t len);
-wchar_t* unicode_transcode8to16(const char* src, wchar_t* out, size_t len);
+char* unicode_transcode16to8(const ntfs_char* src, size_t len);
+ntfs_char* unicode_transcode8to16(const char* src, ntfs_char* out, size_t len);
 
 #endif /* __NTFS_H__ */
diff --git a/src/ntfsx.c b/src/ntfsx.c
index 6c4dd45..fd60f50 100644
--- a/src/ntfsx.c
+++ b/src/ntfsx.c
@@ -365,6 +365,18 @@ void ntfsx_mftmap_destroy(ntfsx_mftmap* map)
   }
 }
 
+static void mftmap_expand(ntfsx_mftmap* map, uint32* allocated)
+{
+  if(map->_count >= *allocated)
+  {
+    (*allocated) += 16;
+    map->_blocks = (struct _ntfsx_mftmap_block*)reallocf(map->_blocks, 
+                  (*allocated) * sizeof(struct _ntfsx_mftmap_block));
+    if(!(map->_blocks))
+ 			errx(1, "out of memory");
+  }
+}
+
 bool ntfsx_mftmap_load(ntfsx_mftmap* map, ntfsx_record* record, int dd)
 {
   bool ret = true;
@@ -377,6 +389,7 @@ bool ntfsx_mftmap_load(ntfsx_mftmap* map, ntfsx_record* record, int dd)
     uint64 length;
     uint64 firstSector;
     uint32 allocated;
+    uint64 total;
 
     /* TODO: Check here whether MFT has already been loaded */
   
@@ -403,6 +416,9 @@ bool ntfsx_mftmap_load(ntfsx_mftmap* map, ntfsx_record* record, int dd)
 
     map->_count = 0;
     allocated = 0;
+    mftmap_expand(map, &allocated);
+
+    total = nonres->cbAttribData / kSectorSize;
 
 		/* Now loop through the data run */
 		if(ntfsx_datarun_first(datarun))
@@ -412,14 +428,7 @@ bool ntfsx_mftmap_load(ntfsx_mftmap* map, ntfsx_record* record, int dd)
         if(datarun->sparse)
           RETWARNBX("invalid mft. sparse data runs");
 
-        if(map->_count >= allocated)
-        {
-          allocated += 16;
-          map->_blocks = (struct _ntfsx_mftmap_block*)reallocf(map->_blocks, 
-                  allocated * sizeof(struct _ntfsx_mftmap_block));
-          if(!(map->_blocks))
-      			errx(1, "out of memory");
-        }
+        mftmap_expand(map, &allocated);
 
         ASSERT(map->info->cluster != 0);
 
@@ -434,8 +443,36 @@ bool ntfsx_mftmap_load(ntfsx_mftmap* map, ntfsx_record* record, int dd)
         map->_blocks[map->_count].length = length;
         map->_blocks[map->_count].firstSector = firstSector;
         map->_count++;
+
+        /*
+         * In some cases the data runs for the MFT don't specify the entire
+         * MFT file, and so we track the remainder and tack it onto
+         * the last block.
+         */
+        total -= length;
 			} 
 			while(ntfsx_datarun_next(datarun));
+
+    }
+
+    if(total > 0)
+    {
+      if(map->_count == 0)
+      {
+        /*
+         * When no data runs were found we start off right
+         * at the MFT and go for the specified length.
+         */
+        ASSERT(allocated > 0);
+        map->_blocks[0].length = total;
+        map->_blocks[0].firstSector = map->info->mft + map->info->first;
+      }
+
+      else
+      {
+        /* Add the remainder of the missing blocks here */
+        map->_blocks[map->_count - 1].length += total;
+      }
     }
 
     ret = true;
diff --git a/src/scrounge.c b/src/scrounge.c
index 4350e5e..5627d36 100644
--- a/src/scrounge.c
+++ b/src/scrounge.c
@@ -43,7 +43,7 @@ void processRecordFileBasics(partitioninfo* pi, ntfsx_record* record, filebasics
     byte* resident = NULL;
     ntfs_attribfilename* filename;
     byte nameSpace;
-    wchar_t* name;
+    ntfs_char* name;
     size_t len;
 #ifndef FC_WIDE
     char* temp;
@@ -82,7 +82,7 @@ void processRecordFileBasics(partitioninfo* pi, ntfsx_record* record, filebasics
         basics->accessed = filename->timeRead;
 
 				/* File Name */
-        name = (wchar_t*)(((byte*)filename) + sizeof(ntfs_attribfilename));
+        name = (ntfs_char*)(((byte*)filename) + sizeof(ntfs_attribfilename));
         len = filename->cFileName;
         if(len > MAX_PATH)
           len = MAX_PATH;
diff --git a/src/unicode.c b/src/unicode.c
index f6197b6..b715a47 100644
--- a/src/unicode.c
+++ b/src/unicode.c
@@ -1,6 +1,6 @@
 
 #include "usuals.h"
-#include <wchar.h>
+#include "ntfs.h"
 
 /*
  * Transcode UCS2 to UTF8. 
@@ -9,13 +9,13 @@
  * resulting length is unpredictable, this function
  * allocates it's own memory.
  */
-char* unicode_transcode16to8(const wchar_t* src, size_t len)
+char* unicode_transcode16to8(const ntfs_char* src, size_t len)
 {
   char* ret = NULL;
   size_t alloc = 0;
   size_t pos = 0;
-  const wchar_t* c;
-  const wchar_t* e;
+  const ntfs_char* c;
+  const ntfs_char* e;
 
   /* Allocate 1.25 times the length initially */
   alloc = len + (len / 4) + 1;
@@ -67,7 +67,7 @@ char* unicode_transcode16to8(const wchar_t* src, size_t len)
  * Since a semi predictable length of the resulting data is 
  * known, the caller should allocate the memory for this conversion.
  */
-wchar_t* unicode_transcode8to16(const char* src, wchar_t* out, size_t len)
+ntfs_char* unicode_transcode8to16(const char* src, ntfs_char* out, size_t len)
 {
   /* Note: out should always be at least as long as src in chars */
 
@@ -91,10 +91,10 @@ wchar_t* unicode_transcode8to16(const char* src, wchar_t* out, size_t len)
        (c[2] & 0xC0) == 0x80 &&
        (c[3] & 0xC0) == 0x80)
     {
-      out[pos++] = (wchar_t)(((wchar_t)c[0] & 7) << 18 |
-                             ((wchar_t)c[1] & 63) << 12 |
-                             ((wchar_t)c[2] & 63) << 6 |
-                             ((wchar_t)c[3] & 63));
+      out[pos++] = (ntfs_char)(((ntfs_char)c[0] & 7) << 18 |
+                             ((ntfs_char)c[1] & 63) << 12 |
+                             ((ntfs_char)c[2] & 63) << 6 |
+                             ((ntfs_char)c[3] & 63));
       c += 3;
     }
 
@@ -104,9 +104,9 @@ wchar_t* unicode_transcode8to16(const char* src, wchar_t* out, size_t len)
             (c[1] & 0xC0) == 0x80 &&
             (c[2] & 0xC0) == 0x80)
     {
-      out[pos++] = (wchar_t)(((wchar_t)c[0] & 15) << 12 |
-                             ((wchar_t)c[1] & 63) << 6 |
-                             ((wchar_t)c[2] & 63));
+      out[pos++] = (ntfs_char)(((ntfs_char)c[0] & 15) << 12 |
+                             ((ntfs_char)c[1] & 63) << 6 |
+                             ((ntfs_char)c[2] & 63));
       c += 2;
     }
 
@@ -115,15 +115,15 @@ wchar_t* unicode_transcode8to16(const char* src, wchar_t* out, size_t len)
             (c[0] & 0xE0) == 0xC0 &&
             (c[1] & 0xC0) == 0x80)
     {
-      out[pos++] = (wchar_t)(((wchar_t)c[0] & 31) << 6 |
-                             ((wchar_t)c[1] & 63));
+      out[pos++] = (ntfs_char)(((ntfs_char)c[0] & 31) << 6 |
+                             ((ntfs_char)c[1] & 63));
       c += 1;
     }
 
     /* First bit set */
     else if(!(c[0] & 0x80))  
     {
-      out[pos++] = (wchar_t)c[0];
+      out[pos++] = (ntfs_char)c[0];
     }
 
     /* Invalid encoding */