Fixed VB style async builder handling

pamidur · Apr 7, 2022 · 7e58016 · 7e58016
1 parent 445b5c3
commit 7e58016
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 11 deletions.
diff --git a/AspectInjector.sln b/AspectInjector.sln
@@ -64,6 +64,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "AspectInjector.Tests.Generi
 		{FC862191-9081-4A58-96FF-B57F17AC40A9} = {FC862191-9081-4A58-96FF-B57F17AC40A9}
 	EndProjectSection
 EndProject
+Project("{F184B08F-C81C-45F6-A57F-5ABD9991F28F}") = "AspectInjector.Tests.VBRuntime", "tests\AspectInjector.Tests.VBRuntime\AspectInjector.Tests.VBRuntime.vbproj", "{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -265,6 +267,18 @@ Global
 		{BF019B75-CECE-46E7-9E3D-620B2CAB6A96}.Release|Any CPU.Build.0 = Release|Any CPU
 		{BF019B75-CECE-46E7-9E3D-620B2CAB6A96}.Release|x86.ActiveCfg = Release|Any CPU
 		{BF019B75-CECE-46E7-9E3D-620B2CAB6A96}.Release|x86.Build.0 = Release|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Debug|x86.ActiveCfg = Debug|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Debug|x86.Build.0 = Debug|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.DebugTests|Any CPU.ActiveCfg = DebugTests|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.DebugTests|Any CPU.Build.0 = DebugTests|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.DebugTests|x86.ActiveCfg = Debug|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.DebugTests|x86.Build.0 = Debug|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Release|Any CPU.Build.0 = Release|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Release|x86.ActiveCfg = Release|Any CPU
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D}.Release|x86.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -285,6 +299,7 @@ Global
 		{7BD22418-F4A9-422B-9FDE-A20708BA0648} = {8F9E5ECE-A6EB-4DC7-B3E2-20319BF21607}
 		{230C6FAF-11A1-4318-A482-69C7252652B1} = {2DAB661A-26B2-440A-BC54-59F2C8404F12}
 		{BF019B75-CECE-46E7-9E3D-620B2CAB6A96} = {2DAB661A-26B2-440A-BC54-59F2C8404F12}
+		{1C5CC61E-3080-42D8-BEE0-6113C8C57D5D} = {2DAB661A-26B2-440A-BC54-59F2C8404F12}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {B153C707-A96B-404C-B0E6-2225AE14B304}

diff --git a/external/FluentIL b/external/FluentIL
diff --git a/src/AspectInjector.Core.Advice/Weavers/Processes/AfterAsyncWeaveProcess.cs b/src/AspectInjector.Core.Advice/Weavers/Processes/AfterAsyncWeaveProcess.cs
@@ -11,16 +11,17 @@
 namespace AspectInjector.Core.Advice.Weavers.Processes
 {
     internal class AfterAsyncWeaveProcess : AfterStateMachineWeaveProcessBase
-    {       
+    {
         private readonly TypeReference _builder;
         private readonly TypeReference _asyncResult;
 
-        public AfterAsyncWeaveProcess(ILogger log, MethodDefinition target, InjectionDefinition injection) 
+        public AfterAsyncWeaveProcess(ILogger log, MethodDefinition target, InjectionDefinition injection)
             : base(log, target, injection)
         {
             SetStateMachine(GetStateMachine());
 
-            var builderField = _stateMachine.Fields.First(f => f.Name == "<>t__builder");
+            var builderField = _stateMachine.Fields.First(f => f.Name == "<>t__builder" || f.Name == "$Builder");
+
             _builder = builderField.FieldType;
             _asyncResult = (_builder as IGenericInstance)?.GenericArguments.FirstOrDefault();
         }
@@ -30,8 +31,8 @@ private TypeReference GetStateMachine()
             var smRef = _method.CustomAttributes.First(ca => ca.AttributeType.Match(StandardType.AsyncStateMachineAttribute))
                 .GetConstructorValue<TypeReference>(0);
 
-            if (smRef.HasGenericParameters)            
-                smRef = _method.Body.Variables.First(v => v.VariableType.Resolve() == smRef).VariableType;            
+            if (smRef.HasGenericParameters)
+                smRef = _method.Body.Variables.First(v => v.VariableType.Resolve() == smRef).VariableType;
 
             return smRef;
         }
@@ -99,13 +100,24 @@ protected override Cut LoadReturnTypeArgument(in Cut pc, AdviceArgument paramete
         }
 
         protected override void InsertStateMachineCall(PointCut code)
-        {
-            var method = _builder.Resolve().Methods.First(m => m.Name == "Start").MakeReference(_builder);
+        {
+            var smctor = _method.Body.Instructions.Where(i => i.OpCode == OpCodes.Newobj && i.Operand is MethodReference mr && mr.DeclaringType.Resolve() == _stateMachine).ToArray();
 
-            var v = _method.Body.Variables.First(vr => vr.VariableType.Resolve().Match(_stateMachine));
-            var loadVar = v.VariableType.IsValueType ? (PointCut)((in Cut c) => c.LoadRef(v)) : (in Cut c) => c.Load(v);
+            if (smctor.Length > 0) // DEBUG build
+            {
+                _method.Body.AfterInstruction(smctor[0], (in Cut cut) => cut.Dup().Here(code));
+            }
+            else // RELEASE build
+            {
+                var v = _method.Body.Variables.First(vr => vr.VariableType.Resolve().Match(_stateMachine));
+                var sminit = _method.Body.Instructions.Where(i => i.OpCode == OpCodes.Initobj && i.Operand is TypeReference tr && tr.Resolve() == _stateMachine).ToArray();
+
+                if (sminit.Length > 0) // VB                                   
+                    _method.Body.AfterInstruction(sminit[0], (in Cut cut) => cut.LoadRef(v).Here(code));
+                else // C#         
+                    _method.Body.AfterEntry((in Cut cut) => cut.LoadRef(v).Here(code));
 
-            _method.Body.OnCall(method, (in Cut cut) => cut.Prev().Prev().Prev().Here(loadVar).Here(code));
+            }
         }
     }
 }
diff --git a/tests/AspectInjector.Tests.VBRuntime/AspectInjector.Tests.VBRuntime.vbproj b/tests/AspectInjector.Tests.VBRuntime/AspectInjector.Tests.VBRuntime.vbproj
@@ -0,0 +1,23 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFrameworks>netcoreapp3.1;net5.0;net6.0;net471;net461</TargetFrameworks>
+    <Configurations>Debug;Release;DebugTests</Configurations>
+  </PropertyGroup>
+  <ItemGroup>
+    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.0.0" />
+    <PackageReference Include="Microsoft.NETFramework.ReferenceAssemblies" Version="1.0.2">
+      <PrivateAssets>all</PrivateAssets>
+      <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
+    </PackageReference>
+    <PackageReference Include="xunit" Version="2.4.1" />
+    <PackageReference Include="xunit.runner.visualstudio" Version="2.4.3">
+      <PrivateAssets>all</PrivateAssets>
+      <IncludeAssets>runtime; build; native; contentfiles; analyzers</IncludeAssets>
+    </PackageReference>
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\src\AspectInjector.Broker\AspectInjector.Broker.csproj" />
+    <ProjectReference Include="..\AspectInjector.Tests.RuntimeAssets\AspectInjector.Tests.RuntimeAssets.csproj" />
+  </ItemGroup>
+</Project>
diff --git a/tests/AspectInjector.Tests.VBRuntime/Issue_0182.vb b/tests/AspectInjector.Tests.VBRuntime/Issue_0182.vb
@@ -0,0 +1,42 @@
+Imports AspectInjector.Broker
+Imports AspectInjector.Tests.Assets
+Imports Xunit
+
+Namespace AspectInjector.Tests.VBRuntime
+
+    Public Class Issue_0182
+
+        <Fact>
+        Public Async Sub AdviceBefore_Methods_Passes()
+            Dim vb = New TestClassVbAspect()
+            Dim result = Await vb.TestAsyncMethod("Test")
+            Assert.Equal("Test", result)
+        End Sub
+
+        Partial Friend NotInheritable Class TestClassVbAspect
+            <InjectInstanceAspect>
+            Public Async Function TestAsyncMethod(Of T)(ByVal obj As T) As Task(Of T)
+                Return Await Task.FromResult(obj)
+            End Function
+        End Class
+
+        <AttributeUsage(AttributeTargets.All, AllowMultiple:=True)>
+        <Injection(GetType(InstanceAspect))>
+        Public Class InjectInstanceAspect
+            Inherits Attribute
+        End Class
+
+        <Aspect(Scope.PerInstance)>
+        Public Class InstanceAspect
+            Inherits TestAspectBase
+
+            <Advice(Kind.After)>
+            Public Sub AspectMethod()
+            End Sub
+
+        End Class
+
+    End Class
+
+End Namespace
+