RAGECOOP-V/Client/Loader/LoaderContext.cs

using System;
using System.Collections.Concurrent;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Windows.Forms;
using GTA.UI;
using SHVDN;
using Console = GTA.Console;
using Script = GTA.Script;

namespace RageCoop.Client.Loader
{
    public class LoaderContext : MarshalByRefObject, IDisposable
    {
        #region PRIMARY-LOADING-LOGIC

        public static ConcurrentDictionary<string, LoaderContext> LoadedDomains =>
            new ConcurrentDictionary<string, LoaderContext>(_loadedDomains);

        private static readonly ConcurrentDictionary<string, LoaderContext> _loadedDomains =
            new ConcurrentDictionary<string, LoaderContext>();

        public bool UnloadRequested;
        public string BaseDirectory => AppDomain.CurrentDomain.BaseDirectory;
        private ScriptDomain CurrentDomain => ScriptDomain.CurrentDomain;

        public static void CheckForUnloadRequest()
        {
            lock (_loadedDomains)
            {
                foreach (var p in _loadedDomains.Values)
                    if (p.UnloadRequested)
                        Unload(p);
            }
        }

        public static bool IsLoaded(string dir)
        {
            return _loadedDomains.ContainsKey(Path.GetFullPath(dir).ToLower());
        }

        public static LoaderContext Load(string dir)
        {
            lock (_loadedDomains)
            {
                dir = Path.GetFullPath(dir).ToLower();
                if (IsLoaded(dir)) throw new Exception("Already loaded");
                ScriptDomain newDomain = null;
                try
                {
                    dir = Path.GetFullPath(dir).ToLower();
                    Directory.CreateDirectory(dir);
                    Exception e = null;
                    // Load domain in main thread
                    Main.QueueToMainThread(() =>
                    {
                        try
                        {
                            /*
                            var assemblies = new List<string>();
                            assemblies.Add(typeof(DomainLoader).Assembly.Location);
                            assemblies.AddRange(typeof(DomainLoader).Assembly.GetReferencedAssemblies()
                                .Select(x => Assembly.Load(x.FullName).Location)
                                .Where(x => !string.IsNullOrEmpty(x)));
                            */

                            // Delete API assemblies
                            Directory.GetFiles(dir, "ScriptHookVDotNet*", SearchOption.AllDirectories).ToList()
                                .ForEach(x => File.Delete(x));
                            var ctxAsm = Path.Combine(dir, "RageCoop.Client.Loader.dll");
                            if (File.Exists(ctxAsm)) File.Delete(ctxAsm);

                            newDomain = ScriptDomain.Load(
                                Directory.GetParent(typeof(ScriptDomain).Assembly.Location).FullName, dir);
                            newDomain.AppDomain.SetData("Primary", ScriptDomain.CurrentDomain);
                            newDomain.AppDomain.SetData("Console",
                                ScriptDomain.CurrentDomain.AppDomain.GetData("Console"));
                            var context = (LoaderContext)newDomain.AppDomain.CreateInstanceFromAndUnwrap(
                                typeof(LoaderContext).Assembly.Location,
                                typeof(LoaderContext).FullName, false,
                                BindingFlags.Instance | BindingFlags.NonPublic,
                                null,
                                new object[] { }
                                , null, null);
                            newDomain.AppDomain.SetData("RageCoop.Client.LoaderContext", context);
                            newDomain.Start();
                            _loadedDomains.TryAdd(dir, context);
                        }
                        catch (Exception ex)
                        {
                            e = ex;
                        }
                    });
                    // Wait till next tick
                    Script.Yield();
                    if (e != null) throw e;
                    return _loadedDomains[dir];
                }
                catch (Exception ex)
                {
                    Notification.Show(ex.ToString());
                    Console.Error(ex);
                    if (newDomain != null) ScriptDomain.Unload(newDomain);
                    throw;
                }
            }
        }

        public static void Unload(LoaderContext domain)
        {
            lock (_loadedDomains)
            {
                Exception ex = null;
                var name = domain.CurrentDomain.Name;
                Console.Info("Unloading domain: " + name);
                Main.QueueToMainThread(() =>
                {
                    try
                    {
                        if (!_loadedDomains.TryRemove(domain.BaseDirectory.ToLower(), out _))
                            throw new Exception("Failed to remove domain from list");
                        domain.Dispose();
                        ScriptDomain.Unload(domain.CurrentDomain);
                    }
                    catch (Exception e)
                    {
                        ex = e;
                        Notification.Show(ex.ToString());
                    }
                });
                Script.Yield();
                if (ex != null) throw ex;
                Console.Info("Unloaded domain: " + name);
            }
        }

        public static void Unload(string dir)
        {
            Unload(_loadedDomains[Path.GetFullPath(dir).ToLower()]);
        }

        public static void UnloadAll()
        {
            lock (_loadedDomains)
            {
                foreach (var d in _loadedDomains.Values.ToArray()) Unload(d);
            }
        }

        #endregion

        #region LOAD-CONTEXT

        private LoaderContext()
        {
            AppDomain.CurrentDomain.DomainUnload += (s, e) => Dispose();
            PrimaryDomain.Tick += Tick;
            PrimaryDomain.KeyEvent += KeyEvent;
            Console.Info(
                $"Loaded domain: {AppDomain.CurrentDomain.FriendlyName}, {AppDomain.CurrentDomain.BaseDirectory}");
        }

        public static ScriptDomain PrimaryDomain => AppDomain.CurrentDomain.GetData("Primary") as ScriptDomain;

        public static LoaderContext CurrentContext =>
            AppDomain.CurrentDomain.GetData("RageCoop.Client.LoaderContext") as LoaderContext;

        /// <summary>
        ///     Request the current domain to be unloaded
        /// </summary>
        public static void RequestUnload()
        {
            if (PrimaryDomain == null)
                throw new NotSupportedException(
                    "Current domain not loaded by the loader therfore cannot be unloaded automatically");
            CurrentContext.UnloadRequested = true;
        }

        private void Tick()
        {
            CurrentDomain.DoTick();
        }

        private void KeyEvent(Keys keys, bool status)
        {
            CurrentDomain.DoKeyEvent(keys, status);
        }

        public override object InitializeLifetimeService()
        {
            // Return null to avoid lifetime restriction on the marshaled object.
            return null;
        }

        public void Dispose()
        {
            lock (this)
            {
                if (PrimaryDomain == null) return;
                PrimaryDomain.Tick -= Tick;
                PrimaryDomain.KeyEvent -= KeyEvent;
                AppDomain.CurrentDomain.SetData("Primary", null);
            }
        }

        #endregion
    }
}