diff --git a/Libraries/SPTarkov.DI/Annotations/Injectable.cs b/Libraries/SPTarkov.DI/Annotations/Injectable.cs index 44c62578..111991bc 100644 --- a/Libraries/SPTarkov.DI/Annotations/Injectable.cs +++ b/Libraries/SPTarkov.DI/Annotations/Injectable.cs @@ -1,7 +1,7 @@ namespace SPTarkov.DI.Annotations; [AttributeUsage(AttributeTargets.Class, Inherited = false)] -public class Injectable(InjectionType injectionType = InjectionType.Scoped, Type? type = null, int typePriority = int.MaxValue) : Attribute +public class Injectable(InjectionType injectionType = InjectionType.Scoped, Type? typeOverride = null, int typePriority = int.MaxValue) : Attribute { public InjectionType InjectionType { @@ -14,6 +14,12 @@ public class Injectable(InjectionType injectionType = InjectionType.Scoped, Type get; set; } = typePriority; + + public Type? TypeOverride + { + get; + set; + } = typeOverride; } public enum InjectionType diff --git a/Libraries/SPTarkov.DI/DependencyInjectionHandler.cs b/Libraries/SPTarkov.DI/DependencyInjectionHandler.cs index 900ef644..a7dddc46 100644 --- a/Libraries/SPTarkov.DI/DependencyInjectionHandler.cs +++ b/Libraries/SPTarkov.DI/DependencyInjectionHandler.cs @@ -61,10 +61,24 @@ public class DependencyInjectionHandler throw new Exception("Invalid usage of DependencyInjectionHandler, this is a one time use service!"); } _oneTimeUseFlag = true; - var sortedInjectableTypes = _injectedTypeNames.Values - .Select(t => - new TypeRefContainer(((Injectable[]) Attribute.GetCustomAttributes(t, typeof(Injectable)))[0], t, t)) - .OrderBy(tRef => tRef.InjectableAttribute.TypePriority); + var typeRefValues = _injectedTypeNames.Values + .Select(t => new TypeRefContainer(((Injectable[]) Attribute.GetCustomAttributes(t, typeof(Injectable)))[0], t, t)); + // All the components that have a type override, we need to find them and remove them before injecting everything + var componentsToRemove = typeRefValues.Where(tr => tr.InjectableAttribute.TypeOverride != null).Select(tr => + string.IsNullOrEmpty(tr.InjectableAttribute.TypeOverride!.FullName) + ? $"{tr.InjectableAttribute.TypeOverride.Namespace}.{tr.InjectableAttribute.TypeOverride.Name}" + : tr.InjectableAttribute.TypeOverride.FullName!) + .ToHashSet(); + // All the components without the removed overrides + var cleanedComponents = typeRefValues.Where(tr => + { + var name = string.IsNullOrEmpty(tr.Type.FullName) + ? $"{tr.Type.Namespace}.{tr.Type.Name}" + : tr.Type.FullName!; + return !componentsToRemove.Contains(name); + }); + // All the components sorted and ready to be inserted into the DI container + var sortedInjectableTypes = cleanedComponents.OrderBy(tRef => tRef.InjectableAttribute.TypePriority); foreach (var typeRefToInject in sortedInjectableTypes) { diff --git a/Libraries/SPTarkov.Server.Core/Utils/Watermark.cs b/Libraries/SPTarkov.Server.Core/Utils/Watermark.cs index d6cec585..fcf6dc1a 100644 --- a/Libraries/SPTarkov.Server.Core/Utils/Watermark.cs +++ b/Libraries/SPTarkov.Server.Core/Utils/Watermark.cs @@ -88,7 +88,7 @@ public class Watermark : IOnLoad sptConfig = _configServer.GetConfig(); } - public Task OnLoad() + public virtual Task OnLoad() { var description = _watermarkLocale.GetDescription(); var warning = _watermarkLocale.GetWarning();