From 8f26e86cd97f24d9d3c44360e7d29a2c36dc1453 Mon Sep 17 00:00:00 2001
From: Ladislav Zezula <ladislav.zezula@avast.com>
Date: Tue, 20 Oct 2020 16:14:26 +0200
Subject: [PATCH] Fixed issue #872

---
 include/retdec/pelib/ImageLoader.h |  6 ++++++
 src/pelib/ImageLoader.cpp          |  3 +++
 src/pelib/RelocationsDirectory.cpp | 11 ++++++++++-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/include/retdec/pelib/ImageLoader.h b/include/retdec/pelib/ImageLoader.h
index f0e390e3f..d3c3c3845 100644
--- a/include/retdec/pelib/ImageLoader.h
+++ b/include/retdec/pelib/ImageLoader.h
@@ -202,6 +202,11 @@ class ImageLoader
 		return (sectionIndex < sections.size()) ? &sections[sectionIndex] : nullptr;
 	}
 
+	std::uint64_t getSizeOfFile() const
+	{
+		return fileSize;
+	}
+
 	std::uint64_t getOrdinalMask() const
 	{
 		return (uint64_t)1 << (getImageBitability() - 1);
@@ -446,6 +451,7 @@ class ImageLoader
 	PELIB_IMAGE_OPTIONAL_HEADER optionalHeader;         // 32/64-bit optional header
 	ByteBuffer rawFileData;                             // Loaded content of the image in case it couldn't have been mapped
 	LoaderError ldrError;
+	std::uint64_t fileSize;                             // Size of the raw file
 	std::uint32_t windowsBuildNumber;
 	std::uint32_t ntSignature;
 	std::uint32_t maxSectionCount;
diff --git a/src/pelib/ImageLoader.cpp b/src/pelib/ImageLoader.cpp
index 13c5278a4..48230f3cf 100644
--- a/src/pelib/ImageLoader.cpp
+++ b/src/pelib/ImageLoader.cpp
@@ -871,6 +871,9 @@ int PeLib::ImageLoader::Load(
 {
 	int fileError;
 
+	// Remember the size of the file for later use
+	fileSize = fileData.size();
+
 	// Check and capture DOS header
 	fileError = captureDosHeader(fileData);
 	if(fileError != ERROR_NONE)
diff --git a/src/pelib/RelocationsDirectory.cpp b/src/pelib/RelocationsDirectory.cpp
index 0c3736f1f..2aed0fb5c 100644
--- a/src/pelib/RelocationsDirectory.cpp
+++ b/src/pelib/RelocationsDirectory.cpp
@@ -26,6 +26,7 @@ namespace PeLib
 		std::uint32_t rva = imageLoader.getDataDirRva(PELIB_IMAGE_DIRECTORY_ENTRY_BASERELOC);
 		std::uint32_t size = imageLoader.getDataDirSize(PELIB_IMAGE_DIRECTORY_ENTRY_BASERELOC);
 		std::uint32_t sizeOfImage = imageLoader.getSizeOfImage();
+		std::uint64_t sizeOfFile = imageLoader.getSizeOfFile();
 
 		// Check for relocations out of image
 		if(rva >= sizeOfImage || (rva + size) < rva || (rva + size) > sizeOfImage)
@@ -34,9 +35,17 @@ namespace PeLib
 			return ERROR_INVALID_FILE;
 		}
 
+		// Check for relocations out of file
+		if(size > sizeOfFile)
+		{
+			RelocationsDirectory::setLoaderError(LDR_ERROR_RELOCATIONS_OUT_OF_IMAGE);
+			return ERROR_INVALID_FILE;
+		}
+
 		// Read the entire relocation directory from the image
 		std::vector<std::uint8_t> vRelocDirectory(size);
-		imageLoader.readImage(vRelocDirectory.data(), rva, size);
+		if(imageLoader.readImage(vRelocDirectory.data(), rva, size) != size)
+			return ERROR_INVALID_FILE;
 
 		// Parse the relocations directory
 		read(vRelocDirectory.data(), size, sizeOfImage);