From 43d86fecea9bad45a361bda0816aae7b34d3fff1 Mon Sep 17 00:00:00 2001
From: CreateAndInject <ssyfzy@126.com>
Date: Thu, 6 Jan 2022 00:07:18 +0800
Subject: [PATCH] Validate data in case allocate too many memory, fix #442

---
 src/DotNet/CustomAttributeReader.cs |  4 +++-
 src/DotNet/MD/BlobStream.cs         |  5 +++--
 src/DotNet/SignatureReader.cs       | 14 +++++++-------
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/DotNet/CustomAttributeReader.cs b/src/DotNet/CustomAttributeReader.cs
index 46c7db465..f58ab5221 100644
--- a/src/DotNet/CustomAttributeReader.cs
+++ b/src/DotNet/CustomAttributeReader.cs
@@ -208,6 +208,8 @@ static CustomAttribute Read(ModuleDef module, ref DataReader reader, ICustomAttr
 		/// <returns>A list of <see cref="CANamedArgument"/>s or <c>null</c> if some error
 		/// occurred.</returns>
 		internal static List<CANamedArgument> ReadNamedArguments(ModuleDef module, ref DataReader reader, int numNamedArgs, GenericParamContext gpContext) {
+			if (numNamedArgs > 0x10000)
+				return null;
 			try {
 				var caReader = new CustomAttributeReader(module, ref reader, gpContext);
 				var namedArgs = caReader.ReadNamedArguments(numNamedArgs);
@@ -528,7 +530,7 @@ CAArgument ReadArrayArgument(SZArraySig arrayType) {
 			int arrayCount = reader.ReadInt32();
 			if (arrayCount == -1) {	// -1 if it's null
 			}
-			else if (arrayCount < 0)
+			else if (arrayCount < 0 || arrayCount > 0x10000)
 				throw new CABlobParserException("Array is too big");
 			else {
 				var array = new List<CAArgument>(arrayCount);
diff --git a/src/DotNet/MD/BlobStream.cs b/src/DotNet/MD/BlobStream.cs
index 51019276a..984ab0af7 100644
--- a/src/DotNet/MD/BlobStream.cs
+++ b/src/DotNet/MD/BlobStream.cs
@@ -46,8 +46,9 @@ public byte[] Read(uint offset) {
 		/// <param name="offset">Offset of blob</param>
 		/// <returns>A new stream</returns>
 		public DataReader CreateReader(uint offset) {
-			TryCreateReader(offset, out var reader);
-			return reader;
+			if (TryCreateReader(offset, out var reader))
+				return reader;
+			return default;
 		}
 
 		/// <summary>
diff --git a/src/DotNet/SignatureReader.cs b/src/DotNet/SignatureReader.cs
index 1792bfa4c..21d00d8c9 100644
--- a/src/DotNet/SignatureReader.cs
+++ b/src/DotNet/SignatureReader.cs
@@ -482,12 +482,12 @@ CallingConventionSig ReadSig() {
 
 		T ReadSig<T>(T methodSig) where T : MethodBaseSig {
 			if (methodSig.Generic) {
-				if (!reader.TryReadCompressedUInt32(out uint count))
+				if (!reader.TryReadCompressedUInt32(out uint count) || count > 0x10000)
 					return null;
 				methodSig.GenParamCount = count;
 			}
 
-			if (!reader.TryReadCompressedUInt32(out uint numParams))
+			if (!reader.TryReadCompressedUInt32(out uint numParams) || numParams > 0x10000)
 				return null;
 
 			methodSig.RetType = ReadType();
@@ -513,7 +513,7 @@ T ReadSig<T>(T methodSig) where T : MethodBaseSig {
 		/// <param name="callingConvention">First byte of signature</param>
 		/// <returns>A new <see cref="LocalSig"/> instance</returns>
 		LocalSig ReadLocalSig(CallingConvention callingConvention) {
-			if (!reader.TryReadCompressedUInt32(out uint count))
+			if (!reader.TryReadCompressedUInt32(out uint count) || count > 0x10000)
 				return null;
 			var sig = new LocalSig(callingConvention, count);
 			var locals = sig.Locals;
@@ -528,7 +528,7 @@ LocalSig ReadLocalSig(CallingConvention callingConvention) {
 		/// <param name="callingConvention">First byte of signature</param>
 		/// <returns>A new <see cref="GenericInstMethodSig"/> instance</returns>
 		GenericInstMethodSig ReadGenericInstMethod(CallingConvention callingConvention) {
-			if (!reader.TryReadCompressedUInt32(out uint count))
+			if (!reader.TryReadCompressedUInt32(out uint count) || count > 0x10000)
 				return null;
 			var sig = new GenericInstMethodSig(callingConvention, count);
 			var args = sig.GenericArguments;
@@ -605,7 +605,7 @@ TypeSig ReadType() {
 
 			case ElementType.GenericInst:
 				nextType = ReadType();
-				if (!reader.TryReadCompressedUInt32(out num))
+				if (!reader.TryReadCompressedUInt32(out num) || num > 0x10000)
 					break;
 				var genericInstSig = new GenericInstSig(nextType as ClassOrValueTypeSig, num);
 				var args = genericInstSig.GenericArguments;
@@ -627,7 +627,7 @@ TypeSig ReadType() {
 				}
 				if (!reader.TryReadCompressedUInt32(out num))
 					break;
-				if (num > MaxArrayRank)
+				if (num > rank)
 					break;
 				var sizes = new List<uint>((int)num);
 				for (i = 0; i < num; i++) {
@@ -637,7 +637,7 @@ TypeSig ReadType() {
 				}
 				if (!reader.TryReadCompressedUInt32(out num))
 					break;
-				if (num > MaxArrayRank)
+				if (num > rank)
 					break;
 				var lowerBounds = new List<int>((int)num);
 				for (i = 0; i < num; i++) {