* Handle malformed number of data directories.

Author: vdisasmdev

PE.DataDirectories.pas

@@ -25,9 +25,15 @@   TDataDirectories = class
 
     procedure Clear;
 
+    procedure NullInvalid(const Msg: TMsgMgr);
+
     // Load array of TImageDataDirectory (va,size) from stream.
-    procedure LoadDirectoriesFromStream(Stream: TStream; const Msg: TMsgMgr;
-      DeclaredCount: integer);
+    procedure LoadDirectoriesFromStream(
+      Stream: TStream;
+      const Msg: TMsgMgr;
+      DeclaredCount: uint32; // # of rva and sizes
+      MaxBytes: uint16
+      );
 
     // Save array of TImageDataDirectory (va,size) to stream.
     // Return saved size.
@@ -190,14 +196,46 @@ function TDataDirectories.GetName(Index: integer): string;
     Result := '';
 end;
 
+procedure TDataDirectories.NullInvalid(const Msg: TMsgMgr);
+var
+  i: integer;
+  needToNullDir: boolean;
+begin
+  // Check RVAs.
+  for i := 0 to self.Count - 1 do
+  begin
+    // Empty dir is ok.
+    if FItems[i].IsEmpty then
+      continue;
+
+    needToNullDir := False;
+
+    if (FItems[i].Size = 0) and (FItems[i].RVA <> 0) then
+    begin
+      Msg.Write(SCategoryDataDirecory, 'Directory # %d has size = 0.', [i]);
+      needToNullDir := true;
+    end
+    else if not TPEImage(FPE).RVAExists(FItems[i].RVA) then
+    begin
+      Msg.Write(SCategoryDataDirecory, 'Directory # %d RVA is not in image.', [i]);
+      needToNullDir := true;
+    end;
+
+    if needToNullDir and (PO_NULL_INVALID_DIRECTORY in TPEImage(FPE).Options) then
+    begin
+      FItems[i] := NULL_IMAGE_DATA_DIRECTORY;
+      Msg.Write(SCategoryDataDirecory, 'Directory # %d nulled.', [i]);
+    end;
+  end;
+end;
+
 procedure TDataDirectories.LoadDirectoriesFromStream;
 var
-  CountToEOF: integer; // Count from StartOfs to EOF.
+  MaxCountPossible: integer;
   CountToRead: integer;
-  SizeToEOF: uint64;
+  SizeToFileEnd: uint64;
   Size: uint32;
-  i: integer;
-  needToNullDir: boolean;
+  MaxSizePossible: uint16;
 begin
   Clear;
 
@@ -207,26 +245,35 @@ procedure TDataDirectories.LoadDirectoriesFromStream;
     exit;
   end;
 
-  SizeToEOF := (Stream.Size - Stream.Position);
+  SizeToFileEnd := (Stream.Size - Stream.Position);
 
-  CountToEOF := SizeToEOF div SizeOf(TImageDataDirectory);
+  // Max size available for dirs.
+  if SizeToFileEnd > MaxBytes then
+    MaxSizePossible := MaxBytes
+  else
+    MaxSizePossible := SizeToFileEnd;
 
-  // File can have part of dword stored. It must be extended with zeros.
-  if (SizeToEOF mod SizeOf(TImageDataDirectory)) <> 0 then
-    inc(CountToEOF);
+  MaxCountPossible := MaxSizePossible div SizeOf(TImageDataDirectory);
 
-  CountToRead := DeclaredCount;
+  // File can have part of dword stored. It must be extended with zeros.
+  if (MaxSizePossible mod SizeOf(TImageDataDirectory)) <> 0 then
+    inc(MaxCountPossible);
 
   if DeclaredCount <> TYPICAL_NUMBER_OF_DIRECTORIES then
-    Msg.Write(SCategoryDataDirecory, 'Non-usual count of directories (%d).', [DeclaredCount]);
+    Msg.Write(SCategoryDataDirecory, 'Non-usual count of directories declared (0x%x).', [DeclaredCount]);
 
-  if DeclaredCount > CountToEOF then
+  if DeclaredCount > MaxCountPossible then
   begin
-    CountToRead := CountToEOF;
+    CountToRead := MaxCountPossible;
 
     Msg.Write(SCategoryDataDirecory,
-      'Declared count of directories is greater than file can contain (%d > %d).',
-      [DeclaredCount, CountToEOF]);
+      'Declared count of directories is greater than file can contain (0x%x > 0x%x).',
+      [DeclaredCount, MaxCountPossible]);
+    Msg.Write(SCategoryDataDirecory, 'Fall back to 0x%x.', [MaxCountPossible]);
+  end
+  else
+  begin
+    CountToRead := DeclaredCount;
   end;
 
   // Read data directories.
@@ -243,32 +290,7 @@ procedure TDataDirectories.LoadDirectoriesFromStream;
   // Set final count.
   self.Count := CountToRead;
 
-  // Check RVAs.
-  for i := 0 to self.Count - 1 do
-  begin
-    // Empty dir is ok.
-    if FItems[i].IsEmpty then
-      continue;
-
-    needToNullDir := False;
-
-    if (FItems[i].Size = 0) and (FItems[i].RVA <> 0) then
-    begin
-      Msg.Write(SCategoryDataDirecory, 'Directory # %d has size = 0.', [i]);
-      needToNullDir := true;
-    end
-    else if not TPEImage(FPE).RVAExists(FItems[i].RVA) then
-    begin
-      Msg.Write(SCategoryDataDirecory, 'Directory # %d RVA is not in image.', [i]);
-      needToNullDir := true;
-    end;
-
-    if needToNullDir and (PO_NULL_INVALID_DIRECTORY in TPEImage(FPE).Options) then
-    begin
-      FItems[i] := NULL_IMAGE_DATA_DIRECTORY;
-      Msg.Write(SCategoryDataDirecory, 'Directory # %d nulled.', [i]);
-    end;
-  end;
+  NullInvalid(Msg);
 end;
 
 function TDataDirectories.SaveToStream(Index: integer; Stream: TStream): boolean;

PE.Image.pas

@@ -1594,6 +1594,11 @@ function TPEImage.LoadFromStream(AStream: TStream; AParseStages: TParserFlags;
 
   // Safe read optional header.
   OptHdrSizeRead := FOptionalHeader.ReadFromStream(AStream, FImageBitSize, -1);
+
+  // Can't read more bytes then available.
+  if OptHdrSizeRead > FFileHeader.SizeOfOptionalHeader then
+    raise Exception.Create('Read size of opt. header > FileHeader.SizeOfOptionalHeader');
+
   DataDirOfs := AStream.Position;
 
   // Load Section Headers.
@@ -1608,8 +1613,10 @@ function TPEImage.LoadFromStream(AStream: TStream; AParseStages: TParserFlags;
   if OptHdrSizeRead <> 0 then
   begin
     AStream.Position := DataDirOfs;
+
     FDataDirectories.LoadDirectoriesFromStream(AStream, Msg,
-      FOptionalHeader.NumberOfRvaAndSizes // declared count
+      FOptionalHeader.NumberOfRvaAndSizes,              // declared count
+      FFileHeader.SizeOfOptionalHeader - OptHdrSizeRead // bytes left in optional header
       );
   end;