From a9830f38e3fc87959885ec910f4d9a1788ea0e6c Mon Sep 17 00:00:00 2001
From: Aristarkh Zagorodnikov <xm@drive.net>
Date: Sun, 10 Feb 2019 23:50:03 +0300
Subject: [PATCH] Improve StackTraceHiddenAttribute handling (#83)

* Add .NET Core 2.1 as target for tests to allow testing for StackTraceHiddenAttribute handling

* Customize EnhancedStackTrace "hidden" frames selection strategy depending on the runtime, and optimize attribute access for non-reflection-only assemblies

* Add .NET Core 2.1 as target for benchmarks to allow benchmarking for it
---
 .../EnhancedStackTrace.Frames.cs              | 127 +++++++++++-------
 .../Ben.Demystifier.Benchmarks.csproj         |   2 +-
 .../Ben.Demystifier.Test.csproj               |   2 +-
 3 files changed, 78 insertions(+), 53 deletions(-)

diff --git a/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs b/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs
index 5897f70..5c8ec4f 100644
--- a/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs
+++ b/src/Ben.Demystifier/EnhancedStackTrace.Frames.cs
@@ -19,6 +19,8 @@ namespace System.Diagnostics
 {
     public partial class EnhancedStackTrace
     {
+        private static readonly Type StackTraceHiddenAttibuteType = Type.GetType("System.Diagnostics.StackTraceHiddenAttribute", false);
+
         private static List<EnhancedStackFrame> GetFrames(Exception exception)
         {
             if (exception == null)
@@ -645,60 +647,59 @@ namespace System.Diagnostics
         private static bool ShowInStackTrace(MethodBase method)
         {
             Debug.Assert(method != null);
-            try
+            var type = method.DeclaringType;
+            if (type == typeof(Task<>) && method.Name == "InnerInvoke")
             {
-                var type = method.DeclaringType;
-                if (type == typeof(Task<>) && method.Name == "InnerInvoke")
+                return false;
+            }
+            if (type == typeof(Task))
+            {
+                switch (method.Name)
+                {
+                    case "ExecuteWithThreadLocal":
+                    case "Execute":
+                    case "ExecutionContextCallback":
+                    case "ExecuteEntry":
+                    case "InnerInvoke":
+                        return false;
+                }
+            }
+            if (type == typeof(ExecutionContext))
+            {
+                switch (method.Name)
+                {
+                    case "RunInternal":
+                    case "Run":
+                        return false;
+                }
+            }
+
+            if (StackTraceHiddenAttibuteType != null)
+            {
+                // Don't show any methods marked with the StackTraceHiddenAttribute
+                // https://github.com/dotnet/coreclr/pull/14652
+                if (IsStackTraceHidden(method))
                 {
                     return false;
                 }
-                if (type == typeof(Task))
-                {
-                    switch (method.Name)
-                    {
-                        case "ExecuteWithThreadLocal":
-                        case "Execute":
-                        case "ExecutionContextCallback":
-                        case "ExecuteEntry":
-                        case "InnerInvoke":
-                            return false;
-                    }
-                }
-                if (type == typeof(ExecutionContext))
-                {
-                    switch (method.Name)
-                    {
-                        case "RunInternal":
-                        case "Run":
-                            return false;
-                    }
-                }
+            }
 
-                // Don't show any methods marked with the StackTraceHiddenAttribute
+            if (type == null)
+            {
+                return true;
+            }
+
+            if (StackTraceHiddenAttibuteType != null)
+            {
+                // Don't show any types marked with the StackTraceHiddenAttribute
                 // https://github.com/dotnet/coreclr/pull/14652
-                foreach (var attibute in EnumerableIList.Create(method.GetCustomAttributesData()))
+                if (IsStackTraceHidden(type))
                 {
-                    // internal Attribute, match on name
-                    if (attibute.AttributeType.Name == "StackTraceHiddenAttribute")
-                    {
-                        return false;
-                    }
+                    return false;
                 }
-
-                if (type == null)
-                {
-                    return true;
-                }
-
-                foreach (var attibute in EnumerableIList.Create(type.GetCustomAttributesData()))
-                {
-                    // internal Attribute, match on name
-                    if (attibute.AttributeType.Name == "StackTraceHiddenAttribute")
-                    {
-                        return false;
-                    }
-                }
-
+            }
+            else
+            {
                 // Fallbacks for runtime pre-StackTraceHiddenAttribute
                 if (type == typeof(ExceptionDispatchInfo) && method.Name == "Throw")
                 {
@@ -723,15 +724,39 @@ namespace System.Diagnostics
                     return false;
                 }
             }
-            catch
-            {
-                // GetCustomAttributesData can throw
-                return true;
-            }
 
             return true;
         }
 
+        private static bool IsStackTraceHidden(MemberInfo memberInfo)
+        {
+            if (!memberInfo.Module.Assembly.ReflectionOnly)
+            {
+                return memberInfo.GetCustomAttributes(StackTraceHiddenAttibuteType, false).Length != 0;
+            }
+
+            EnumerableIList<CustomAttributeData> attributes;
+            try
+            {
+                attributes = EnumerableIList.Create(memberInfo.GetCustomAttributesData());
+            }
+            catch (NotImplementedException)
+            {
+                return false;
+            }
+
+            foreach (var attribute in attributes)
+            {
+                // reflection-only attribute, match on name
+                if (attribute.AttributeType.FullName == StackTraceHiddenAttibuteType.FullName)
+                {
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
         private static bool TryResolveStateMachineMethod(ref MethodBase method, out Type declaringType)
         {
             Debug.Assert(method != null);
diff --git a/test/Ben.Demystifier.Benchmarks/Ben.Demystifier.Benchmarks.csproj b/test/Ben.Demystifier.Benchmarks/Ben.Demystifier.Benchmarks.csproj
index 75d99f0..b208351 100644
--- a/test/Ben.Demystifier.Benchmarks/Ben.Demystifier.Benchmarks.csproj
+++ b/test/Ben.Demystifier.Benchmarks/Ben.Demystifier.Benchmarks.csproj
@@ -1,6 +1,6 @@
 <Project Sdk="Microsoft.NET.Sdk" ToolsVersion="15.0">
   <PropertyGroup>
-    <TargetFrameworks>netcoreapp2.0;net462</TargetFrameworks>
+    <TargetFrameworks>netcoreapp2.1;netcoreapp2.0;net462</TargetFrameworks>
     <Configuration>Release</Configuration>
     <OutputType>Exe</OutputType>
   </PropertyGroup>
diff --git a/test/Ben.Demystifier.Test/Ben.Demystifier.Test.csproj b/test/Ben.Demystifier.Test/Ben.Demystifier.Test.csproj
index ee86c07..bf3aa90 100644
--- a/test/Ben.Demystifier.Test/Ben.Demystifier.Test.csproj
+++ b/test/Ben.Demystifier.Test/Ben.Demystifier.Test.csproj
@@ -1,7 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk">
   
   <PropertyGroup>
-    <TargetFrameworks>netcoreapp2.0;net46</TargetFrameworks>
+    <TargetFrameworks>netcoreapp2.1;netcoreapp2.0;net46</TargetFrameworks>
     <SignAssembly>true</SignAssembly>
     <AssemblyOriginatorKeyFile>..\..\src\Ben.Demystifier\key.snk</AssemblyOriginatorKeyFile>
     <IsPackable>false</IsPackable>