Detect and mark as reference assembly

dotnet · Sep 26, 2023 · 22b41a2 · 22b41a2
1 parent f027e14
commit 22b41a2
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/...ontent/eng/tools/tasks/Microsoft.DotNet.SourceBuild.Tasks.LeakDetection/CheckForPoison.cs b/...ontent/eng/tools/tasks/Microsoft.DotNet.SourceBuild.Tasks.LeakDetection/CheckForPoison.cs
@@ -12,6 +12,8 @@
 using System.IO.Compression;
 using System.Linq;
 using System.Reflection;
+using System.Reflection.Metadata;
+using System.Reflection.PortableExecutable;
 using System.Security.Cryptography;
 using System.Text;
 using System.Xml;
@@ -288,7 +290,14 @@ private static PoisonedFileEntry CheckSingleFile(IEnumerable<CatalogPackageEntry
                 AssemblyName asm = AssemblyName.GetAssemblyName(fileToCheck);
                 if (IsAssemblyPoisoned(fileToCheck))
                 {
-                    poisonEntry.Type |= PoisonType.AssemblyAttribute;
+                    if (fileToCheck.EndsWith(".dll") && HasAttributeOfType(fileToCheck, "System.Runtime.CompilerServices.ReferenceAssemblyAttribute"))
+                    {
+                        poisonEntry.Type |= PoisonType.ReferenceAssemblyAttribute;
+                    }
+                    else
+                    {
+                        poisonEntry.Type |= PoisonType.AssemblyAttribute;
+                    }
                 }
             }
             catch
@@ -320,6 +329,43 @@ private static bool IsAssemblyPoisoned(string path)
             return false;
         }
 
+        private static bool HasAttributeOfType(string assemblyPath, string desiredAttributeType)
+        {
+            using var stream = new FileStream(assemblyPath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite);
+            using var peReader = new PEReader(stream);
+
+            MetadataReader reader = peReader.GetMetadataReader();
+            return reader.CustomAttributes.Select(attrHandle => reader.GetCustomAttribute(attrHandle))
+                    .Any(attr => IsCustomAttributeOfType(reader, attr, desiredAttributeType));
+        }
+
+        private static bool IsCustomAttributeOfType(MetadataReader reader, CustomAttribute attr, string desiredAttributeType)
+        {
+            string customType = "";
+            if (attr.Constructor.Kind == HandleKind.MethodDefinition)
+            {
+                MethodDefinition mdef = reader.GetMethodDefinition((MethodDefinitionHandle)attr.Constructor);
+                TypeDefinition tdef = reader.GetTypeDefinition(mdef.GetDeclaringType());
+                customType = $"{reader.GetString(tdef.Namespace)}.{reader.GetString(tdef.Name)}";
+            }
+            else if (attr.Constructor.Kind == HandleKind.MemberReference)
+            {
+                MemberReference mref = reader.GetMemberReference((MemberReferenceHandle)attr.Constructor);
+
+                if (mref.Parent.Kind == HandleKind.TypeReference)
+                {
+                    TypeReference tref = reader.GetTypeReference((TypeReferenceHandle)mref.Parent);
+                    customType = $"{reader.GetString(tref.Namespace)}.{reader.GetString(tref.Name)}";
+                }
+                else if (mref.Parent.Kind == HandleKind.TypeDefinition)
+                {
+                    TypeDefinition tdef = reader.GetTypeDefinition((TypeDefinitionHandle)mref.Parent);
+                    customType = $"{reader.GetString(tdef.Namespace)}.{reader.GetString(tdef.Name)}";
+                }
+            }
+            return customType == desiredAttributeType;
+        }
+
         private static PoisonedFileEntry ExtractAndCheckZipFileOnly(IEnumerable<CatalogPackageEntry> catalogedPackages, string zipToCheck, string markerFileName, string tempDir, Queue<string> futureFilesToCheck)
         {
             var poisonEntry = new PoisonedFileEntry();

diff --git a/...ld/content/eng/tools/tasks/Microsoft.DotNet.SourceBuild.Tasks.LeakDetection/PoisonType.cs b/...ld/content/eng/tools/tasks/Microsoft.DotNet.SourceBuild.Tasks.LeakDetection/PoisonType.cs
@@ -10,6 +10,7 @@ internal enum PoisonType
         None = 0,
         Hash = 1,
         AssemblyAttribute = 2,
+        ReferenceAssemblyAttribute = 3,
         NupkgFile = 4,
     }
 }