Merge pull request #407 from Washi1337/feature/ca-validation

Add CA parameter count validation
Washi1337 · Jan 26, 2023 · 5c92f02 · 5c92f02
2 parents be42dd4 + dc4ffd9
commit 5c92f02
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 0 deletions.
diff --git a/src/AsmResolver.DotNet/Builder/DotNetDirectoryBuffer.CodedIndices.cs b/src/AsmResolver.DotNet/Builder/DotNetDirectoryBuffer.CodedIndices.cs
@@ -17,6 +17,11 @@ private void AddCustomAttribute(MetadataToken ownerToken, CustomAttribute attrib
         {
             var table = Metadata.TablesStream.GetSortedTable<CustomAttribute, CustomAttributeRow>(TableIndex.CustomAttribute);
 
+            // Ensure the signature defines the right parameters w.r.t. the constructor's parameters.
+            if (attribute.Constructor is not null && attribute.Signature is not null)
+                attribute.Signature.IsCompatibleWith(attribute.Constructor, ErrorListener);
+
+            // Add it.
             var encoder = Metadata.TablesStream.GetIndexEncoder(CodedIndex.HasCustomAttribute);
             var row = new CustomAttributeRow(
                 encoder.EncodeToken(ownerToken),

diff --git a/src/AsmResolver.DotNet/Signatures/CustomAttributeSignature.cs b/src/AsmResolver.DotNet/Signatures/CustomAttributeSignature.cs
@@ -156,6 +156,43 @@ protected override void WriteContents(in BlobSerializationContext context)
             for (int i = 0; i < NamedArguments.Count; i++)
                 NamedArguments[i].Write(context);
         }
+
+        /// <summary>
+        /// Validates whether the signature is compatible with the provided attribute constructor.
+        /// </summary>
+        /// <param name="constructor">The constructor to validate against.</param>
+        /// <returns><c>true</c> if the constructor is compatible, <c>false</c> otherwise.</returns>
+        public bool IsCompatibleWith(ICustomAttributeType constructor)
+        {
+            return IsCompatibleWith(constructor, EmptyErrorListener.Instance);
+        }
+
+        /// <summary>
+        /// Validates whether the signature is compatible with the provided attribute constructor.
+        /// </summary>
+        /// <param name="constructor">The constructor to validate against.</param>
+        /// <param name="listener">The object responsible for reporting any errors during the validation of the signature.</param>
+        /// <returns><c>true</c> if the constructor is compatible, <c>false</c> otherwise.</returns>
+        public virtual bool IsCompatibleWith(ICustomAttributeType constructor, IErrorListener listener)
+        {
+            var signature = constructor.Signature;
+
+            int expectedCount = signature?.ParameterTypes.Count ?? 0;
+            if (expectedCount != FixedArguments.Count)
+            {
+                listener.MetadataBuilder(
+                    $"Custom attribute constructor {constructor.SafeToString()} expects {expectedCount} arguments but the signature provided {FixedArguments.Count} arguments.");
+                return false;
+            }
+
+            if (signature?.SentinelParameterTypes.Count > 0)
+            {
+                listener.MetadataBuilder($"Custom attribute constructor {constructor.SafeToString()} defines sentinel parameters.");
+                return false;
+            }
+
+            return true;
+        }
     }
 
 }
diff --git a/src/AsmResolver.DotNet/Signatures/SerializedCustomAttributeSignature.cs b/src/AsmResolver.DotNet/Signatures/SerializedCustomAttributeSignature.cs
@@ -75,5 +75,11 @@ protected override void WriteContents(in BlobSerializationContext context)
             else
                 _reader.Fork().WriteToOutput(context.Writer);
         }
+
+        /// <inheritdoc />
+        public override bool IsCompatibleWith(ICustomAttributeType constructor, IErrorListener listener)
+        {
+            return !IsInitialized || base.IsCompatibleWith(constructor, listener);
+        }
     }
 }
diff --git a/test/AsmResolver.DotNet.Tests/CustomAttributeTest.cs b/test/AsmResolver.DotNet.Tests/CustomAttributeTest.cs
@@ -646,5 +646,25 @@ public void NamedGenericTypeNullArgument(bool rebuild, bool access)
 
             Assert.Null(argument.Argument.Element);
         }
+
+        [Fact]
+        public void TestSignatureCompatibility()
+        {
+            var module = new ModuleDefinition("Dummy");
+            var factory = module.CorLibTypeFactory;
+            var ctor = factory.CorLibScope
+                .CreateTypeReference("System", "CLSCompliantAttribute")
+                .CreateMemberReference(".ctor", MethodSignature.CreateInstance(factory.Void, factory.Boolean))
+                .ImportWith(module.DefaultImporter);
+
+            var attribute = new CustomAttribute(ctor);
+
+            // Empty signature is not compatible with a ctor that takes a boolean.
+            Assert.False(attribute.Signature!.IsCompatibleWith(attribute.Constructor!));
+
+            // If we add it, it should be compatible again.
+            attribute.Signature.FixedArguments.Add(new CustomAttributeArgument(factory.Boolean, true));
+            Assert.True(attribute.Signature!.IsCompatibleWith(attribute.Constructor!));
+        }
     }
 }