Check for SBRP attribute when looking for poisoning

dotnet · Oct 17, 2023 · 5e596d9 · 5e596d9
1 parent a87d1d5
commit 5e596d9
Show file tree

Hide file tree

Showing 2 changed files with 45 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
@@ -141,6 +141,10 @@ public class CheckForPoison : Task
 
         private const string PoisonMarker = "POISONED";
 
+        private const string SbrpAttributeType = "System.Reflection.AssemblyMetadataAttribute";
+
+        private const string SbrpAttributeValue = "source-build-reference-packages";
+
         public override bool Execute()
         {
             IEnumerable<PoisonedFileEntry> poisons = GetPoisonedFiles(FilesToCheck.Select(f => f.ItemSpec), HashCatalogFilePath, MarkerFileName);
@@ -286,7 +290,11 @@ private static PoisonedFileEntry CheckSingleFile(IEnumerable<CatalogPackageEntry
             try
             {
                 AssemblyName asm = AssemblyName.GetAssemblyName(fileToCheck);
-                if (IsAssemblyPoisoned(fileToCheck))
+                if (IsAssemblyFromSbrp(fileToCheck))
+                {
+                    poisonEntry.Type |= PoisonType.ReferenceAssemblyAttribute;
+                }
+                else if (IsAssemblyPoisoned(fileToCheck))
                 {
                     poisonEntry.Type |= PoisonType.AssemblyAttribute;
                 }
@@ -320,6 +328,41 @@ private static bool IsAssemblyPoisoned(string path)
             return false;
         }
 
+        private static bool IsAssemblyFromSbrp(string assemblyPath)
+        {
+            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 => IsAttributeSbrp(reader, attr));
+        }
+
+        private static bool IsAttributeSbrp(MetadataReader reader, CustomAttribute attr)
+        {
+            string attributeType = string.Empty;
+
+            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);
+                    attributeType = $"{reader.GetString(tref.Namespace)}.{reader.GetString(tref.Name)}";
+                }
+            }
+
+            if (attributeType == SbrpAttributeType)
+            {
+                byte[] data = reader.GetBlobBytes(attr.Value);
+                string attributeValue = Encoding.UTF8.GetString(data);
+
+                return attributeValue.Contains(SbrpAttributeValue);
+            }
+            return false;
+        }
+
         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
@@ -11,5 +11,6 @@ internal enum PoisonType
         Hash = 1,
         AssemblyAttribute = 2,
         NupkgFile = 4,
+        ReferenceAssemblyAttribute = 8,
     }
 }