Add ability to inject into module initializer (#165)

* Add option to inject into module initializer * Change AutoInit into InitMode enum * Add unit tests for InitMode.ModuleLoad
Keboo · Oct 5, 2019 · 02e7fbd · 02e7fbd
1 parent 5a9a49c
commit 02e7fbd
Show file tree

Hide file tree

Showing 10 changed files with 178 additions and 12 deletions.
diff --git a/AutoDI.AspNetCore/WebHostBuilderMixins.cs b/AutoDI.AspNetCore/WebHostBuilderMixins.cs
@@ -1,6 +1,6 @@
 using System.Reflection;
 using Microsoft.AspNetCore.Hosting;
-[assembly:AutoDI.Settings(AutoInit = false)]
+[assembly:AutoDI.Settings(InitMode = AutoDI.InitMode.None)]
 namespace AutoDI.AspNetCore
 {
     public static class WebHostBuilderMixins

diff --git a/AutoDI.Build.Tests/ManualInjectionOfContainer.cs b/AutoDI.Build.Tests/ManualInjectionOfContainer.cs
@@ -48,7 +48,7 @@ public void CanManuallyInjectTheGeneratedContainer()
 //<type:ConsoleApplication/>
 //<ref: AutoDI />
 //<weaver: AutoDI.Build.ProcessAssemblyTask />
-//<raw:[assembly:AutoDI.Settings(AutoInit = false)] />
+//<raw:[assembly:AutoDI.Settings(InitMode = AutoDI.InitMode.None)] />
 namespace ManualInjectionNamespace
 {
     using AutoDI;

diff --git a/AutoDI.Build.Tests/ModuleLoadInjectionTests.cs b/AutoDI.Build.Tests/ModuleLoadInjectionTests.cs
@@ -0,0 +1,80 @@
+using AutoDI.AssemblyGenerator;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using ModuleLoadInjectionNamespace;
+using System.Reflection;
+using System.Threading.Tasks;
+
+namespace AutoDI.Build.Tests
+{
+    [TestClass]
+    public class ModuleLoadInjectionTests
+    {
+        private static Assembly _testAssembly;
+
+        [ClassInitialize]
+        public static async Task Initialize(TestContext context)
+        {
+            var gen = new Generator();
+
+            _testAssembly = (await gen.Execute()).SingleAssembly();
+        }
+
+        [ClassCleanup]
+        public static void Cleanup()
+        {
+            DI.Dispose(_testAssembly);
+        }
+
+        [TestMethod]
+        public void InitThrowsIfModuleLoadAssembly()
+        {
+            try
+            {
+                DI.Init(_testAssembly);
+            }
+            catch (TargetInvocationException e)
+                when (e.InnerException is AlreadyInitializedException)
+            {
+                return;
+            }
+            Assert.Fail($"Excepted {nameof(AlreadyInitializedException)}");
+        }
+
+        [TestMethod]
+        public void TryInitReturnsFalseOnFirstInvocation()
+        {
+            Assert.IsFalse(DI.TryInit(_testAssembly));
+        }
+
+        [TestMethod]
+        public void LibraryDependenciesAreInjected()
+        {
+            dynamic sut = _testAssembly.CreateInstance<ModuleLoadingLibrary>();
+            Assert.IsTrue(((object)sut.Service).Is<Service>());
+        }
+    }
+}
+
+//<assembly />
+//<ref: AutoDI />
+//<weaver: AutoDI.Build.ProcessAssemblyTask />
+//<raw:[assembly:AutoDI.Settings(InitMode = AutoDI.InitMode.ModuleLoad)] />
+namespace ModuleLoadInjectionNamespace
+{
+    using AutoDI;
+
+    public class ModuleLoadingLibrary
+    {
+        public IService Service { get; }
+
+        public ModuleLoadingLibrary([Dependency] IService service = null)
+        {
+            Service = service;
+        }
+    }
+
+    public interface IService { }
+
+    public class Service : IService { }
+}
+//</assembly>
diff --git a/AutoDI.Build.Tests/SettingsTests.cs b/AutoDI.Build.Tests/SettingsTests.cs
@@ -20,5 +20,13 @@ public void DebugCodeGenerationDefaultsOff()
 
             Assert.AreEqual((int)CodeLanguage.None, (int)settings.DebugCodeGeneration);
         }
+
+        [TestMethod]
+        public void InitModeDefaultsEntryPoint()
+        {
+            var settings = new Settings();
+
+            Assert.AreEqual((int)InitMode.EntryPoint, (int)settings.InitMode);
+        }
     }
 }
diff --git a/AutoDI.Build/AutoDI.Build.csproj b/AutoDI.Build/AutoDI.Build.csproj
@@ -28,6 +28,7 @@
     <Compile Include="..\AutoDI\Constants.cs" Link="Constants.cs" />
     <Compile Include="..\AutoDI\DebugLogLevel.cs" Link="DebugLogLevel.cs" />
     <Compile Include="..\AutoDI\IncludeAssemblyAttribute.cs" Link="IncludeAssemblyAttribute.cs" />
+    <Compile Include="..\AutoDI\InitMode.cs" Link="InitMode.cs" />
     <Compile Include="..\AutoDI\Lifetime.cs" Link="Lifetime.cs" />
     <Compile Include="..\AutoDI\MapAttribute.cs" Link="MapAttribute.cs" />
     <Compile Include="..\AutoDI\SettingsAttribute.cs" Link="SettingsAttribute.cs" />

diff --git a/AutoDI.Build/ProcessAssemblyTask.InjectContainer.cs b/AutoDI.Build/ProcessAssemblyTask.InjectContainer.cs
@@ -1,4 +1,5 @@
-using Mono.Cecil;
+using System.Linq;
+using Mono.Cecil;
 using Mono.Cecil.Cil;
 
 namespace AutoDI.Build
@@ -20,5 +21,47 @@ private void InjectInitCall(MethodReference initMethod)
                 Logger.Debug($"No entry point in {ModuleDefinition.FileName}. Skipping container injection.", DebugLogLevel.Default);
             }
         }
+
+        private void InjectModuleCctorInitCall(MethodReference initMethod)
+        {
+            var moduleClass = ModuleDefinition.Types.FirstOrDefault(t => t.Name == "<Module>");
+            if (moduleClass == null)
+            {
+                Logger.Debug($"No module class in {ModuleDefinition.FileName}. Skipping container injection.", DebugLogLevel.Default);
+                return;
+            }
+
+            var cctor = FindOrCreateCctor(moduleClass);    
+            if (cctor != null)
+            {
+                Logger.Debug("Injecting AutoDI .cctor init call", DebugLogLevel.Verbose);
+
+                var injector = new Injector(cctor);
+                injector.Insert(OpCodes.Ldnull);
+                injector.Insert(OpCodes.Call, initMethod);
+            }
+            else
+            {
+                Logger.Debug($"Couldn't find or create .cctor in {ModuleDefinition.FileName}. Skipping container injection.", DebugLogLevel.Default);
+            }
+
+        }
+
+        private MethodDefinition FindOrCreateCctor(TypeDefinition moduleClass)
+        {
+            var cctor = moduleClass.Methods.FirstOrDefault(m => m.Name == ".cctor");
+            if (cctor == null)
+            {
+                var attributes = MethodAttributes.Private
+                    | MethodAttributes.HideBySig
+                    | MethodAttributes.Static
+                    | MethodAttributes.SpecialName
+                    | MethodAttributes.RTSpecialName;
+                cctor = new MethodDefinition(".cctor", attributes, ModuleDefinition.ImportReference(typeof(void)));
+                moduleClass.Methods.Add(cctor);
+                cctor.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
+            }
+            return cctor;
+        }
     }
 }
diff --git a/AutoDI.Build/ProcessAssemblyTask.cs b/AutoDI.Build/ProcessAssemblyTask.cs
@@ -60,9 +60,20 @@ protected override bool WeaveAssembly()
 
                     ModuleDefinition.Types.Add(GenerateAutoDIClass(mapping, settings, gen, out MethodDefinition initMethod));
 
-                    if (settings.AutoInit)
+                    switch (settings.InitMode)
                     {
-                        InjectInitCall(initMethod);
+                        case InitMode.None:
+                            Logger.Debug("Skipping injections of Init method", DebugLogLevel.Verbose);
+                            break;
+                        case InitMode.EntryPoint:
+                            InjectInitCall(initMethod);
+                            break;
+                        case InitMode.ModuleLoad:
+                            InjectModuleCctorInitCall(initMethod);
+                            break;
+                        default:
+                            Logger.Warning($"Unsupported InitMode: {settings.InitMode}");
+                            break;
                     }
                 }
                 else

diff --git a/AutoDI.Build/Settings.cs b/AutoDI.Build/Settings.cs
@@ -21,8 +21,8 @@ public static Settings Load(ModuleDefinition module)
                         {
                             switch (property.Name)
                             {
-                                case nameof(SettingsAttribute.AutoInit):
-                                    settings.AutoInit = (bool)property.Argument.Value;
+                                case nameof(SettingsAttribute.InitMode):
+                                    settings.InitMode = (InitMode)property.Argument.Value;
                                     break;
                                 case nameof(SettingsAttribute.Behavior):
                                     settings.Behavior = (Behaviors)property.Argument.Value;
@@ -209,12 +209,12 @@ IList<string> GetConstructorParameterNames()
         public Behaviors Behavior { get; set; } = Behaviors.Default;
 
         /// <summary>
-        /// Automatically initialize AutoDI in assembly entry point (if available)
+        /// Specifies if and when AutoDI initializes during assembly loading.
         /// </summary>
-        public bool AutoInit { get; set; } = true;
+        public InitMode InitMode { get; set; } = InitMode.EntryPoint;
 
         /// <summary>
-        /// Generate registration calls no the container. Setting to false will negate AutoInit.
+        /// Generate registration calls no the container. Setting to false will negate InitMode.
         /// </summary>
         public bool GenerateRegistrations { get; set; } = true;
 
@@ -236,7 +236,7 @@ public override string ToString()
 
             sb.AppendLine("AutoDI Settings:");
             sb.AppendLine($"  Behavior(s): {Behavior}");
-            sb.AppendLine($"  AutoInit: {AutoInit}");
+            sb.AppendLine($"  InitMode: {InitMode}");
             sb.AppendLine($"  GenerateRegistrations: {GenerateRegistrations}");
             sb.AppendLine($"  DebugLogLevel: {DebugLogLevel}");
             sb.AppendLine($"  DebugExceptions: {DebugExceptions}");

diff --git a/AutoDI/InitMode.cs b/AutoDI/InitMode.cs
@@ -0,0 +1,23 @@
+namespace AutoDI
+{
+    /// <summary>
+    /// Controls when AutoDI is initialized
+    /// </summary>
+    public enum InitMode
+    {
+        /// <summary>
+        /// The container must be initialized manually.
+        /// </summary>
+        None,
+        /// <summary>
+        /// The container will be initialized on your program's entry point.
+        /// Recommended for executable programs.
+        /// </summary>
+        EntryPoint,
+        /// <summary>
+        /// The container will be initialized on assembly module load.
+        /// Recommended for libraries.
+        /// </summary>
+        ModuleLoad
+    }
+}
diff --git a/AutoDI/SettingsAttribute.cs b/AutoDI/SettingsAttribute.cs
@@ -6,7 +6,7 @@ namespace AutoDI
     public class SettingsAttribute : Attribute
     {
         public Behaviors Behavior { get; set; }
-        public bool AutoInit { get; set; }
+        public InitMode InitMode { get; set; }
         public bool GenerateRegistrations { get; set; }
         public bool DebugExceptions { get; set; }
         public DebugLogLevel DebugLogLevel { get; set; }