From b4af7a2f4f1ba4407a61a82b8c04441ecc5e5541 Mon Sep 17 00:00:00 2001 From: TauAkiou Date: Sat, 17 Oct 2020 21:10:07 -0400 Subject: [PATCH] Code cleanup --- AmongUsCapture/IPC/DBus/IConnectLink.cs | 18 +-- AmongUsCapture/IPC/DBus/IPCadapterDBus.cs | 146 +++++++++++------- AmongUsCapture/IPC/IPCadapter.cs | 5 +- .../IPC/Windows/IPCadapterRpcBuffer.cs | 1 + AmongUsCapture/Memory/ProcessMemory.cs | 2 +- AmongUsCapture/Memory/ProcessMemoryLinux.cs | 2 - AmongUsCapture/Memory/ProcessMemoryWindows.cs | 6 - AmongUsCapture/Program.cs | 38 +---- 8 files changed, 107 insertions(+), 111 deletions(-) diff --git a/AmongUsCapture/IPC/DBus/IConnectLink.cs b/AmongUsCapture/IPC/DBus/IConnectLink.cs index 21965b33..2e2c9e4b 100644 --- a/AmongUsCapture/IPC/DBus/IConnectLink.cs +++ b/AmongUsCapture/IPC/DBus/IConnectLink.cs @@ -7,27 +7,19 @@ namespace AmongUsCapture.DBus [DBusInterface("org.AmongUsCapture.ConnectLink")] public interface IConnectLink : IDBusObject { - Task WatchConnectInfoAsync(Action handler); + Task SendConnectUriAsync(string uri); } public class IPCLink : IConnectLink { - string _connectlink; - public string ConnectLink - { - get { return _connectlink; } - set - { - _connectlink = value; - SentLink?.Invoke(_connectlink); - } - } public ObjectPath ObjectPath => new ObjectPath("/org/AmongUsCapture/ConnectLink"); public event Action SentLink; - public Task WatchConnectInfoAsync(Action handler) + public Task SendConnectUriAsync(string uri) { - return SignalWatcher.AddAsync(this, nameof(SentLink), handler); + // Call event and send URI. + SentLink?.Invoke(uri); + return Task.CompletedTask; } } } \ No newline at end of file diff --git a/AmongUsCapture/IPC/DBus/IPCadapterDBus.cs b/AmongUsCapture/IPC/DBus/IPCadapterDBus.cs index e674db7e..7715a816 100644 --- a/AmongUsCapture/IPC/DBus/IPCadapterDBus.cs +++ b/AmongUsCapture/IPC/DBus/IPCadapterDBus.cs @@ -22,12 +22,11 @@ namespace AmongUsCapture.DBus { class IPCadapterDBus : IPCadapter { - private Thread _dbusProcessingThread; private CancellationTokenSource _cancellation = new CancellationTokenSource(); - private Connection _dbusconnection; - private IConnectLink _ipclink; - private bool _isListening; - + private bool _serverIsStarted; + private bool _isHostInstance; + private Task _serverTask; + public override URIStartResult HandleURIStart(string[] args) { var myProcessId = Process.GetCurrentProcess().Id; @@ -43,13 +42,13 @@ public override URIStartResult HandleURIStart(string[] args) Console.WriteLine(Program.GetExecutablePath()); //mutex = new Mutex(true, appName, out var createdNew); - bool createdNew = false; + _isHostInstance = false; var wasURIStart = args.Length > 0 && args[0].StartsWith(UriScheme + "://"); var result = URIStartResult.CONTINUE; if (!File.Exists(Path.Join(Settings.StorageLocation, ".amonguscapture.pid"))) { - createdNew = true; + _isHostInstance = true; } else { @@ -66,38 +65,35 @@ public override URIStartResult HandleURIStart(string[] args) var capproc = Process.GetProcessById(pidint); var assmbname = System.Reflection.Assembly.GetExecutingAssembly().GetName().Name; var runnername = System.Diagnostics.Process.GetCurrentProcess().MainModule.FileName; - - if (!capproc.ProcessName.Contains("dotnet")) + var iscapture = false; + + foreach (ProcessModule mod in capproc.Modules) { - // We're just going to assume that a dotnet process that matches this set of parameters - // is a running capture process. - throw new ArgumentException(); - + // If we find amonguscapturedll in the modules, we can be certain + // that the located pid is, in fact, an AmongUsCapture process. + if (mod.ModuleName == "AmongUsCapture.dll") + { + iscapture = true; + break; + } } - else if (!capproc.ProcessName.Contains(Path.GetFileName(runnername))) - { - throw new ArgumentException(); - } - - if (capproc.HasExited) - { + if (!iscapture || capproc.HasExited) throw new ArgumentException(); - } } catch (ArgumentException e) { // Process doesn't exist. Clear the file. Console.WriteLine($"Found stale PID file containing {pid}."); File.Delete(Path.Join(Settings.StorageLocation, ".amonguscapture.pid")); - createdNew = true; + _isHostInstance = true; } } } } - if (createdNew) + if (_isHostInstance) { using (var pidwriter = File.CreateText(Path.Join(Settings.StorageLocation, ".amonguscapture.pid"))) { @@ -106,7 +102,7 @@ public override URIStartResult HandleURIStart(string[] args) } - if (!createdNew) // send it to already existing instance if applicable, then close + if (!_isHostInstance) // send it to already existing instance if applicable, then close { if (wasURIStart) SendToken(args[0]).Wait(); @@ -114,12 +110,10 @@ public override URIStartResult HandleURIStart(string[] args) } else if (wasURIStart) // URI start on new instance, continue as normal but also handle current argument { - // if we are running, we create a file lock with our process in it. - // Also attach the pid delete handler from Program. - result = URIStartResult.PARSE; } + // Register the xdg-mime handler for discord links. RegisterProtocol(); return result; @@ -136,10 +130,16 @@ private static void RegisterProtocol() if (!File.Exists(xdg_file)) { var executingassmb = System.Reflection.Assembly.GetExecutingAssembly().Location; + Console.WriteLine(executingassmb); if (Path.HasExtension("dll")) { executingassmb = "dotnet " + executingassmb; } + else + { + executingassmb = System.Diagnostics.Process.GetCurrentProcess().MainModule.FileName; + Console.WriteLine(executingassmb); + } var xdg_file_write = new string[] { @@ -189,16 +189,20 @@ private static void RegisterProtocol() public override async Task SendToken(string jsonText) { + // Delay the send until we know for certain the server is started. + // If this isn't a host instance, we don't care. + while (!_serverIsStarted && _isHostInstance) + await Task.Delay(500); + // Send the token via DBus. using (Connection conn = new Connection(Address.Session)) { await conn.ConnectAsync(); + + var _ipclink = conn.CreateProxy("org.AmongUsCapture.ConnectLink", + "/org/AmongUsCapture/ConnectLink"); - var obj = new IPCLink(); - await conn.RegisterObjectAsync(obj); - await conn.RegisterServiceAsync("org.AmongUsCapture.ipc", ServiceRegistrationOptions.None); - - obj.ConnectLink = jsonText; + await _ipclink.SendConnectUriAsync(jsonText); } return true; @@ -210,30 +214,38 @@ public override void SendToken(string host, string connectCode) OnTokenEvent(st); } - public async override Task RegisterMinion() + public override Task RegisterMinion() { - Task.Factory.StartNew(async () => + _serverTask = Task.Factory.StartNew(async () => { - - using (_dbusconnection = new Connection(Address.Session)) + try { - await _dbusconnection.ConnectAsync(); - - _ipclink = _dbusconnection.CreateProxy("org.AmongUsCapture.ConnectLink", - "/org/AmongUsCapture/ConnectLink"); + using (Connection conn = new Connection(Address.Session)) + { + await conn.ConnectAsync(); - await _ipclink.WatchConnectInfoAsync(RespondToDbus); + var obj = new IPCLink(); + await conn.RegisterObjectAsync(obj); + await conn.RegisterServiceAsync("org.AmongUsCapture.ConnectLink", + ServiceRegistrationOptions.None); - _isListening = true; + obj.SentLink += RespondToDbus; - while (!_cancellation.IsCancellationRequested) - { - _cancellation.Token.ThrowIfCancellationRequested(); - await Task.Delay(int.MaxValue); + _serverIsStarted = true; + + while (!_cancellation.IsCancellationRequested) + { + _cancellation.Token.ThrowIfCancellationRequested(); + await Task.Delay(int.MaxValue); + } } } + catch (OperationCanceledException e) + { + Console.WriteLine("Cancel called - terminating DBus loop."); + } }); - + return Task.CompletedTask; } public override void startWithToken(string uri) @@ -241,24 +253,54 @@ public override void startWithToken(string uri) OnTokenEvent(StartToken.FromString(uri)); } - public override bool Cancel() + public override async Task Cancel() { if (!_cancellation.IsCancellationRequested) { _cancellation.Cancel(); + await CleanPid(); + await _serverTask; return true; } return false; } - private void RespondToDbus(string signalresponse) + private Task CleanPid() { - Settings.conInterface.WriteModuleTextColored("DBus", Color.Silver, - $"Received new message on DBus: \"{signalresponse}\""); + // Make sure the pidfile is cleaned up if we have one. + var pidfile = Path.Join(Settings.StorageLocation, ".amonguscapture.pid"); - signalresponse = signalresponse.Trim('\r', '\n'); + if (File.Exists(pidfile)) + { + int pid; + bool fileread; + using (var pidread = File.OpenText(Path.Join(Settings.StorageLocation, ".amonguscapture.pid"))) + { + fileread = Int32.TryParse(pidread.ReadLine(), out pid); + } + + if (!fileread) + { + // Bad read, file must be corrupt. Clear pidfile. + File.Delete(pidfile); + } + if (pid == Process.GetCurrentProcess().Id) + { + // This is our process. Delete file. + File.Delete(pidfile); + } + } + + return Task.CompletedTask; + } + + private void RespondToDbus(string signalresponse) + { + Settings.conInterface.WriteModuleTextColored("DBus", Color.Silver, + $"Received DBus Method Call: \"{signalresponse}\""); + var token = StartToken.FromString(signalresponse); OnTokenEvent(token); diff --git a/AmongUsCapture/IPC/IPCadapter.cs b/AmongUsCapture/IPC/IPCadapter.cs index 1c44e520..c624daa6 100644 --- a/AmongUsCapture/IPC/IPCadapter.cs +++ b/AmongUsCapture/IPC/IPCadapter.cs @@ -26,9 +26,6 @@ abstract class IPCadapter protected const string FriendlyName = "AmongUs Capture"; protected Mutex mutex; - - - private static IPCadapter instance; public static IPCadapter getInstance() { @@ -68,7 +65,7 @@ protected virtual void OnTokenEvent(StartToken e) public abstract void startWithToken(string uri); // This method is for implementations that might be cancelable - such as DBUS. - public virtual bool Cancel() + public virtual async Task Cancel() { return true; } diff --git a/AmongUsCapture/IPC/Windows/IPCadapterRpcBuffer.cs b/AmongUsCapture/IPC/Windows/IPCadapterRpcBuffer.cs index 3ab8d387..1fcf4878 100644 --- a/AmongUsCapture/IPC/Windows/IPCadapterRpcBuffer.cs +++ b/AmongUsCapture/IPC/Windows/IPCadapterRpcBuffer.cs @@ -4,6 +4,7 @@ using SharedMemory; using Mutex = System.Threading.Mutex; using Process = System.Diagnostics.Process; +#pragma warning disable 4014 namespace AmongUsCapture.Windows { diff --git a/AmongUsCapture/Memory/ProcessMemory.cs b/AmongUsCapture/Memory/ProcessMemory.cs index 5cb2fb4e..ceecc8d1 100644 --- a/AmongUsCapture/Memory/ProcessMemory.cs +++ b/AmongUsCapture/Memory/ProcessMemory.cs @@ -31,7 +31,7 @@ public static ProcessMemory getInstance() protected bool is64Bit; public Process process; public List modules; - public bool IsHooked { get; protected set; } + public bool IsHooked => process != null && !process.HasExited; public abstract bool HookProcess(string name); public abstract void LoadModules(); public abstract T Read(IntPtr address, params int[] offsets) where T : unmanaged; diff --git a/AmongUsCapture/Memory/ProcessMemoryLinux.cs b/AmongUsCapture/Memory/ProcessMemoryLinux.cs index 0c8296b5..c6bcf47d 100644 --- a/AmongUsCapture/Memory/ProcessMemoryLinux.cs +++ b/AmongUsCapture/Memory/ProcessMemoryLinux.cs @@ -50,8 +50,6 @@ public override bool HookProcess(string name) is64Bit = flag; LoadModules(); - IsHooked = true; - } } diff --git a/AmongUsCapture/Memory/ProcessMemoryWindows.cs b/AmongUsCapture/Memory/ProcessMemoryWindows.cs index ea1fd1d6..0b18ebdd 100644 --- a/AmongUsCapture/Memory/ProcessMemoryWindows.cs +++ b/AmongUsCapture/Memory/ProcessMemoryWindows.cs @@ -8,12 +8,6 @@ namespace AmongUsCapture { public class ProcessMemoryWindows : ProcessMemory { - private static bool is64Bit; - public static Process process; - public static List modules; - - public static bool IsHooked => process != null && !process.HasExited; - public override bool HookProcess(string name) { if (!IsHooked) diff --git a/AmongUsCapture/Program.cs b/AmongUsCapture/Program.cs index e31c2338..8ccaf3e1 100644 --- a/AmongUsCapture/Program.cs +++ b/AmongUsCapture/Program.cs @@ -66,7 +66,7 @@ private static void Main(string[] args) while (Settings.conInterface is null) Thread.Sleep(250); Task.Factory.StartNew(() => socket.Init()) .Wait(); // run socket in background. Important to wait for init to have actually finished before continuing - Task.Factory.StartNew(() => IPCadapter.getInstance().RegisterMinion()); + Task.Factory.StartNew(() => IPCadapter.getInstance().RegisterMinion()).Wait(); // Add a GLib Idle handler to fix the issue here. Idle.Add(delegate @@ -93,16 +93,15 @@ private static void OpenGUI() window.DeleteEvent += (object o, DeleteEventArgs e) => { - // Make sure the pid file is deleted if the application quits. - CleanPid(); + // Make sure that the IPC adapter has a chance to clean up after itself. + IPCadapter.getInstance().Cancel().Wait(); Application.Quit(); }; window.ShowAll(); Application.Run(); - - CleanPid(); + IPCadapter.getInstance().Cancel().Wait(); Environment.Exit(0); } @@ -111,34 +110,7 @@ public static string GetExecutablePath() { return Process.GetCurrentProcess().MainModule.FileName; } - - private static void CleanPid() - { - // Make sure the pidfile is cleaned up if we have one. - var pidfile = Path.Join(Settings.StorageLocation, ".amonguscapture.pid"); - - if (File.Exists(pidfile)) - { - int pid; - bool fileread; - using (var pidread = File.OpenText(Path.Join(Settings.StorageLocation, ".amonguscapture.pid"))) - { - fileread = Int32.TryParse(pidread.ReadLine(), out pid); - } - - if (!fileread) - { - // Bad read, file must be corrupt. Clear pidfile. - File.Delete(pidfile); - } - - if (pid == Process.GetCurrentProcess().Id) - { - // This is our process. Delete file. - File.Delete(pidfile); - } - } - } + [DllImport("kernel32.dll", SetLastError = true)] [return: MarshalAs(UnmanagedType.Bool)]