Check for SBRP attribute when looking for poison

dotnet · Oct 16, 2023 · ff0ea37 · ff0ea37
1 parent 6a1ba20
commit ff0ea37
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 52 deletions.
diff --git a/src/SourceBuild/content/Directory.Build.props b/src/SourceBuild/content/Directory.Build.props
@@ -134,8 +134,6 @@
     <PoisonMarkerFile>.prebuilt.xml</PoisonMarkerFile>
     <SourceBuiltPoisonReportDataFile>$(PackageReportDir)poison-source-built-catalog.xml</SourceBuiltPoisonReportDataFile>
     <SourceBuiltPoisonMarkerFile>.source-built.xml</SourceBuiltPoisonMarkerFile>
-    <CurrentSbrpPoisonReportDataFile>$(PackageReportDir)poison-current-sbrp-catalog.xml</CurrentSbrpPoisonReportDataFile>
-    <CurrentSbrpPoisonMarkerFile>.current-sbrp.xml</CurrentSbrpPoisonMarkerFile>
     <ProjectAssetsJsonArchiveFile>$(PackageReportDir)all-project-assets-json-files.zip</ProjectAssetsJsonArchiveFile>
     <ProdConManifestFile>$(PackageReportDir)prodcon-build.xml</ProdConManifestFile>
     <PoisonedReportFile>$(PackageReportDir)poisoned.txt</PoisonedReportFile>

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
@@ -31,6 +31,12 @@ public class CheckForPoison : Task
         [Required]
         public ITaskItem[] FilesToCheck { get; set; }
 
+        /// <summary>
+        /// The path of the project directory to the FilesToCheck.
+        /// </summary>
+        [Required]
+        public string ProjectDirPath { get; set; }
+
         /// <summary>
         /// The output path for an XML poison report, if desired.
         /// </summary>
@@ -143,6 +149,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);
@@ -285,19 +295,16 @@ private static PoisonedFileEntry CheckSingleFile(IEnumerable<CatalogPackageEntry
                 }
             }
 
-            try
+           try
             {
                 AssemblyName asm = AssemblyName.GetAssemblyName(fileToCheck);
-                if (IsAssemblyPoisoned(fileToCheck))
+                if (IsAssemblyFromSbrp(fileToCheck))
                 {
-                    if (fileToCheck.EndsWith(".dll") && HasAttributeOfType(fileToCheck, "System.Runtime.CompilerServices.ReferenceAssemblyAttribute"))
-                    {
-                        poisonEntry.Type |= PoisonType.ReferenceAssemblyAttribute;
-                    }
-                    else
-                    {
-                        poisonEntry.Type |= PoisonType.AssemblyAttribute;
-                    }
+                    poisonEntry.Type |= PoisonType.ReferenceAssemblyAttribute;
+                }
+                else if (IsAssemblyPoisoned(fileToCheck))
+                {
+                    poisonEntry.Type |= PoisonType.AssemblyAttribute;
                 }
             }
             catch
@@ -329,41 +336,39 @@ private static bool IsAssemblyPoisoned(string path)
             return false;
         }
 
-        private static bool HasAttributeOfType(string assemblyPath, string desiredAttributeType)
+        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 => IsCustomAttributeOfType(reader, attr, desiredAttributeType));
+                    .Any(attr => IsAttributeSbrp(reader, attr));
         }
 
-        private static bool IsCustomAttributeOfType(MetadataReader reader, CustomAttribute attr, string desiredAttributeType)
+        private static bool IsAttributeSbrp(MetadataReader reader, CustomAttribute attr)
         {
-            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)
+            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);
-                    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)}";
+                    attributeType = $"{reader.GetString(tref.Namespace)}.{reader.GetString(tref.Name)}";
                 }
             }
-            return customType == desiredAttributeType;
+
+            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)

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,7 +10,7 @@ internal enum PoisonType
         None = 0,
         Hash = 1,
         AssemblyAttribute = 2,
-        ReferenceAssemblyAttribute = 3,
         NupkgFile = 4,
+        ReferenceAssemblyAttribute = 8,
     }
 }
diff --git a/src/SourceBuild/content/repo-projects/source-build-reference-packages.proj b/src/SourceBuild/content/repo-projects/source-build-reference-packages.proj
@@ -24,25 +24,4 @@
       SourcePath="$(LocalNuGetPackageCacheDirectory)" />
   </Target>
 
-  <Import Project="Sdk.targets" Sdk="Microsoft.NET.Sdk" />
-
-  <UsingTask AssemblyFile="$(LeakDetectionTasksAssembly)" TaskName="MarkAndCatalogPackages" />
-
-  <Target Name="PoisonSbrp"
-          Condition="'$(EnablePoison)' == 'true'"
-          Inputs="$(MSBuildProjectFullPath)"
-          Outputs="$(RepoCompletedSemaphorePath)PoisonCurrentSbrp.complete"
-          AfterTargets="Build">
-    <ItemGroup>
-      <CurrentSbrp Include="$(ReferencePackagesDir)**/*.nupkg" />
-    </ItemGroup>
-
-    <Message Importance="High" Text="[$([System.DateTime]::Now.ToString('HH:mm:ss.ff'))] Poisoning current SBRP for leak detection." />
-
-    <MarkAndCatalogPackages PackagesToMark="@(CurrentSbrp)" CatalogOutputFilePath="$(CurrentSbrpPoisonReportDataFile)" MarkerFileName="$(CurrentSbrpPoisonMarkerFile)" />
-
-    <WriteLinesToFile File="$(RepoCompletedSemaphorePath)PoisonCurrentSbrp.complete" Overwrite="true" />
-    <Message Importance="High" Text="[$([System.DateTime]::Now.ToString('HH:mm:ss.ff'))] Done poisoning current SBRP." />
-  </Target>
-
 </Project>