Merge pull request #312 from avast-tl/feature-unpacker-limit-overall-…

…memory

Add support for limiting the overall memory of our unpacker (#290)

scripts/retdec-decompiler.sh

@@ -86,8 +86,8 @@ print_help()
 	echo "               --static-code-sigfile path             Adds additional signature file for static code detection."
 	echo "               --static-code-archive path             Adds additional signature file for static code detection from given archive."
 	echo "               --no-default-static-signatures         No default signatures for statically linked code analysis are loaded (options static-code-sigfile/archive are still available)."
-	echo "               --max-memory bytes                     Limits the maximal memory of fileinfo, bin2llvmir, and llvmir2hll into the given number of bytes."
-	echo "               --no-memory-limit                      Disables the default memory limit (half of system RAM) of fileinfo, bin2llvmir, and llvmir2hll."
+	echo "               --max-memory bytes                     Limits the maximal memory of fileinfo, unpacker, bin2llvmir, and llvmir2hll into the given number of bytes."
+	echo "               --no-memory-limit                      Disables the default memory limit (half of system RAM) of fileinfo, unpacker, bin2llvmir, and llvmir2hll."
 }
 SCRIPT_NAME=$0
 GETOPT_SHORTOPT="a:e:hkl:m:o:p:"
@@ -986,6 +986,13 @@ if [ "$MODE" = "bin" ] || [ "$MODE" = "raw" ]; then
 	## Unpacking.
 	##
 	UNPACK_PARAMS=(--extended-exit-codes --output "$OUT_UNPACKED" "$IN")
+	if [ ! -z "$MAX_MEMORY" ]; then
+		UNPACK_PARAMS+=(--max-memory "$MAX_MEMORY")
+	elif [ -z "$NO_MEMORY_LIMIT" ]; then
+		# By default, we want to limit the memory of retdec-unpacker into half
+		# of system RAM to prevent potential black screens on Windows (#270).
+		UNPACK_PARAMS+=(--max-memory-half-ram)
+	fi
 
 	if [ "$GENERATE_LOG" ]; then
 		LOG_UNPACKER_OUTPUT="$($UNPACK_SH "${UNPACK_PARAMS[@]}" 2>&1)"

scripts/retdec-unpacker.sh

@@ -59,6 +59,8 @@ print_help()
 	echo "    -h,        --help                 Print this help message."
 	echo "    -e,        --extended-exit-codes  Use more granular exit codes than just 0/1."
 	echo "    -o FILE,   --output FILE          Output file (default: file-unpacked)."
+	echo "               --max-memory N         Limit the maximal memory of retdec-unpacker to N bytes."
+	echo "               --max-memory-half-ram  Limit the maximal memory of retdec-unpacker to half of system RAM."
 }
 
 #
@@ -106,6 +108,11 @@ try_to_unpack()
 	local UNPACKER_EXIT_CODE_PREPROCESSING_ERROR=3
 
 	UNPACKER_PARAMS=("$IN" -o "$OUT")
+	if [ ! -z "$MAX_MEMORY" ]; then
+		UNPACKER_PARAMS+=(--max-memory "$MAX_MEMORY")
+	elif [ ! -z "$MAX_MEMORY_HALF_RAM" ]; then
+		UNPACKER_PARAMS+=(--max-memory-half-ram)
+	fi
 	echo ""
 	echo "##### Trying to unpack $IN into $OUT by using generic unpacker..."
 	echo "RUN: $UNPACKER ${UNPACKER_PARAMS[@]}"
@@ -157,7 +164,7 @@ try_to_unpack()
 
 SCRIPT_NAME=$0
 GETOPT_SHORTOPT="eho:"
-GETOPT_LONGOPT="extended-exit-codes,help,output:"
+GETOPT_LONGOPT="extended-exit-codes,help,output:,max-memory:,max-memory-half-ram"
 
 # Check script arguments.
 PARSED_OPTIONS=$(getopt -o "$GETOPT_SHORTOPT" -l "$GETOPT_LONGOPT" -n "$SCRIPT_NAME" -- "$@")
@@ -180,6 +187,19 @@ while true; do
 		[ "$OUT" ] && print_error_and_die "Duplicate option: -o|--output"
 		OUT="$2"
 		shift 2;;
+	--max-memory-half-ram)
+		[ "$MAX_MEMORY_HALF_RAM" ] && print_error_and_die "Duplicate option: --max-memory-half-ram"
+		[ "$MAX_MEMORY" ] && print_error_and_die "Clashing options: --max-memory-half-ram and --max-memory"
+		MAX_MEMORY_HALF_RAM="1"
+		shift;;
+	--max-memory)
+		[ "$MAX_MEMORY" ] && print_error_and_die "Duplicate option: --max-memory"
+		[ "$MAX_MEMORY_HALF_RAM" ] && print_error_and_die "Clashing options: --max-memory and --max-memory-half-ram"
+		MAX_MEMORY="$2"
+		if [[ ! "$MAX_MEMORY" =~ ^[0-9]+$ ]]; then
+			print_error_and_die "Invalid value for --max-memory: $MAX_MEMORY (expected a positive integer)"
+		fi
+		shift 2;;
 	--)								# Input file.
 		if [ $# -eq 2 ]; then
 			IN="$2"

src/unpackertool/unpacker.cpp

@@ -4,10 +4,13 @@
  * @copyright (c) 2017 Avast Software, licensed under the MIT license
  */
 
+#include <cstddef>
 #include <iostream>
 #include <memory>
 
+#include "retdec/utils/conversion.h"
 #include "retdec/utils/filesystem_path.h"
+#include "retdec/utils/memory.h"
 #include "retdec/cpdetect/cpdetect.h"
 #include "retdec/fileformat/fileformat.h"
 #include "arg_handler.h"
@@ -26,7 +29,8 @@ enum ExitCode
 	EXIT_CODE_OK = 0, ///< Unpacker ended successfully.
 	EXIT_CODE_NOTHING_TO_DO, ///< There was not found matching plugin.
 	EXIT_CODE_UNPACKING_FAILED, ///< At least one plugin failed at the unpacking of the file.
-	EXIT_CODE_PREPROCESSING_ERROR ///< Error with preprocessing of input file before unpacking.
+	EXIT_CODE_PREPROCESSING_ERROR, ///< Error with preprocessing of input file before unpacking.
+	EXIT_CODE_MEMORY_LIMIT_ERROR ///< There was an error when setting the memory limit.
 };
 
 bool detectPackers(const std::string& inputFile, std::vector<retdec::cpdetect::DetectResult>& detectedPackers)
@@ -115,6 +119,39 @@ ExitCode processArgs(ArgHandler& handler, char argc, char** argv)
 
 	bool brute = handler["brute"]->used;
 
+	// --max-memory N
+	if (handler["max-memory"]->used)
+	{
+		auto maxMemoryLimitStr = handler["max-memory"]->input;
+
+		std::size_t maxMemoryLimit = 0;
+		auto conversionSucceeded = strToNum(maxMemoryLimitStr, maxMemoryLimit);
+		if (!conversionSucceeded)
+		{
+			std::cerr << "Invalid value for --max-memory: '"
+				<< maxMemoryLimitStr << "'!\n";
+			return EXIT_CODE_MEMORY_LIMIT_ERROR;
+		}
+
+		auto limitingSucceeded = limitSystemMemory(maxMemoryLimit);
+		if (!limitingSucceeded)
+		{
+			std::cerr << "Failed to limit memory to "
+				<< maxMemoryLimitStr << " bytes!\n";
+			return EXIT_CODE_MEMORY_LIMIT_ERROR;
+		}
+	}
+	// --max-memory-half-ram
+	else if (handler["max-memory-half-ram"]->used)
+	{
+		auto limitingSucceeded = limitSystemMemoryToHalfOfTotalSystemMemory();
+		if (!limitingSucceeded)
+		{
+			std::cerr << "Failed to limit memory to half of system RAM!\n";
+			return EXIT_CODE_MEMORY_LIMIT_ERROR;
+		}
+	}
+
 	// -h|--help
 	if (handler["help"]->used)
 	{
@@ -176,13 +213,17 @@ int main(int argc, char** argv)
 			"\n"
 			"Non-group optional arguments:\n"
 			"   -b|--brute             Tell unpacker to run plugins in the brute mode. Plugins may or may not\n"
-			"                          implement brute methods for unpacking. They can completely ignore this argument."
+			"                          implement brute methods for unpacking. They can completely ignore this argument.\n"
+			"   --max-memory N         Limit maximal memory to N bytes.\n"
+			"   --max-memory-half-ram  Limit maximal memory to half of system RAM."
 	);
 
 	handler.registerArg('h', "help", false);
 	handler.registerArg('o', "output", true);
 	handler.registerArg('p', "plugins", false);
 	handler.registerArg('b', "brute", false);
+	handler.registerArg('m', "max-memory", true);
+	handler.registerArg('M', "max-memory-half-ram", false);
 
 	return processArgs(handler, argc, argv);
 }