Poison SBRP

src/SourceBuild/content/eng/tools/tasks/Microsoft.DotNet.SourceBuild.Tasks.LeakDetection/CheckForPoison.cs

@@ -12,8 +12,11 @@
 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.Text.RegularExpressions;
 using System.Xml;
 using System.Xml.Linq;
 
@@ -147,6 +150,10 @@ public class CheckForPoison : Task
 
         private const string PoisonMarker = "POISONED";
 
+        private const string SbrpAttributeType = "System.Reflection.AssemblyMetadataAttribute";
+
+        private const string SbrpAttributeValuePattern = "source\\s?source\\-build\\-reference\\-packages";
+
         private record CandidateFileEntry(string ExtractedPath, string DisplayPath);
 
         public override bool Execute()
@@ -298,7 +305,11 @@ private static PoisonedFileEntry CheckSingleFile(IEnumerable<CatalogPackageEntry
             try
             {
                 AssemblyName asm = AssemblyName.GetAssemblyName(fileToCheck);
-                if (IsAssemblyPoisoned(fileToCheck))
+                if (!candidate.DisplayPath.Contains("SourceBuildReferencePackages") && IsAssemblyFromSbrp(fileToCheck))
+                {
+                    poisonEntry.Type |= PoisonType.SourceBuildReferenceAssembly;
+                }
+                else if (IsAssemblyPoisoned(fileToCheck))
                 {
                     poisonEntry.Type |= PoisonType.AssemblyAttribute;
                 }
@@ -332,6 +343,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)
+            {
+                BlobReader blobReader = reader.GetBlobReader(attr.Value);
+                string attributeValue = Encoding.UTF8.GetString(blobReader.ReadBytes(blobReader.Length));
+                attributeValue = Regex.Replace(attributeValue, @"\p{C}+", string.Empty);
+                return Regex.IsMatch(attributeValue, SbrpAttributeValuePattern);
+            }
+            return false;
+        }
+
         private static PoisonedFileEntry ExtractAndCheckZipFileOnly(IEnumerable<CatalogPackageEntry> catalogedPackages, CandidateFileEntry candidate, string markerFileName, string tempDir, Queue<CandidateFileEntry> futureFilesToCheck)
         {
             var poisonEntry = new PoisonedFileEntry();

src/SourceBuild/content/eng/tools/tasks/Microsoft.DotNet.SourceBuild.Tasks.LeakDetection/PoisonType.cs

@@ -11,5 +11,6 @@ internal enum PoisonType
         Hash = 1,
         AssemblyAttribute = 2,
         NupkgFile = 4,
+        SourceBuildReferenceAssembly = 8,
     }
 }