UnrealAgent.cs

// Copyright 2014 Xamarin Inc.
// For details, see LICENSE.txt.

using System;
using System.IO;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;
using System.Net.Sockets;
using System.Runtime.InteropServices;

#if AGENT_CLIENT
namespace XamarinStudio.Unreal.Projects
#else
namespace UnrealEngine.MainDomain
#endif
{
	abstract class UnrealAgent : IDisposable
	{
		string gameRoot, engineRoot, agentFile;
		bool disposed;

		protected readonly object ConnectionCreationLock = new object();
		UnrealAgentConnection connection;

		protected UnrealAgent (string engineRoot, string gameRoot)
		{
			this.engineRoot = engineRoot;
			this.gameRoot = gameRoot;
			this.agentFile = Path.Combine (gameRoot, ".xamarin-ide");
		}

		public string GameRoot {
			get { return gameRoot; }
		}

		public string EngineRoot {
			get { return engineRoot; }
		}

		public string AgentFile {
			get { return agentFile; }
		}

		protected bool IsDisposed {
			get { return disposed; }
		}

		public bool IsConnected {
			get { return connection != null; }
		}

		public Task Connect (CancellationToken token)
		{
			if (IsConnected)
				return TaskFromResult ((object)null);

			return new ConnectTask (this, StartTarget (), token).Task
				//HACK: small delay before sending commands if launching the editor, or it can crash
				.ContinueWith (t => {
					t.Wait ();
					if (t.IsCompleted)
						Thread.Sleep (5000);
				});
		}

		protected abstract Process StartTarget ();

		class ConnectTask
		{
			TaskCompletionSource<object> tcs = new TaskCompletionSource<object> ();
			readonly UnrealAgent agent;
			readonly Process target;

			public ConnectTask (UnrealAgent agent, Process target, CancellationToken token)
			{
				this.agent = agent;
				this.target = target;

				if (token.CanBeCanceled)
					token.Register (OnCancelled);

				agent.Connected += OnConnected;
				agent.Disconnected += OnDisconnected;
				target.Exited += OnProcessExited;
				target.EnableRaisingEvents = true;

				if (agent.IsConnected)
					OnConnected ();
				else if (target.HasExited)
					OnDisconnected ();
			}

			public Task Task {
				get { return tcs.Task; }
			}

			void OnProcessExited (object sender, EventArgs e)
			{
				if (target.ExitCode != 0 && tcs.TrySetException (new Exception ("Target exited unexpectedly")))
					Dispose ();
			}

			void OnCancelled ()
			{
				if (tcs.TrySetCanceled ())
					Dispose ();
			}

			void OnConnected ()
			{
				if (tcs.TrySetResult (null))
					Dispose ();
			}

			void OnDisconnected ()
			{
				if (tcs.TrySetException (new Exception ("Unable to connect to target")))
					Dispose ();
			}

			void Dispose ()
			{
				agent.Connected -= OnConnected;
				agent.Disconnected -= OnDisconnected;
				target.Exited -= OnProcessExited;
				target.EnableRaisingEvents = false;
			}
		}


		/// <summary>
		/// Subclass should call this when a connection is starting.
		/// </summary>
		protected UnrealAgentConnection OnConnecting (TcpClient client)
		{
			lock (ConnectionCreationLock) {
				if (connection != null)
					connection.Dispose ();
				connection = new UnrealAgentConnection (client, HandleCommandBase, HandleDisposed);
			}
			return connection;
		}

		/// <summary>
		/// Subclass should call this to close the connection.
		/// </summary>
		protected void CloseConnection ()
		{
			lock (ConnectionCreationLock) {
				if (connection != null)
					connection.Close ();
				connection = null;
			}
		}

		void HandleDisposed (UnrealAgentConnection c)
		{
			lock (ConnectionCreationLock) {
				if (c == connection)
					connection = null;
			}

			var evt = Disconnected;
			if (evt == null)
				return;
			try {
				evt ();
			} catch (Exception ex) {
				UnrealAgentHelper.Error (ex, "Unhandled exception");
			}
		}

		bool HandleCommandBase (string name, string[] args)
		{
			if (name == "Connected")
			{
				var e = Connected;
				if (e != null) {
					e ();
				}
				return true;
			}

			return HandleCommand (name, args);
		}

		protected abstract bool HandleCommand (string name, string[] args);

		public virtual void Dispose ()
		{
			if (disposed)
				return;

			lock (ConnectionCreationLock) {
				if (disposed)
					return;
				disposed = true;
			}

			Dispose (true);
			GC.SuppressFinalize (this);
		}

		protected virtual void Dispose (bool disposing)
		{
			if (disposing && connection != null) {
				connection.Close ();
				connection = null;
			}
		}

		~UnrealAgent()
		{
			Dispose (false);
		}

		public event Action Connected;
		public event Action Disconnected;

		bool SendSync (string name, params object[] args)
		{
			var c = connection;
			if (c != null)
				return c.Send(name, args);
			return false;
		}

		protected Task<bool> Send (string name, params object[] args)
		{
			if (connection == null)
				return TaskFromResult (false);

			return LogExceptions (Task.Factory.StartNew (() => SendSync (name, args)));
		}

		protected Task<bool> ConnectAndSend (CancellationToken token, string name, params object[] args)
		{
			return ConnectAndSend (token, true, name, args);
		}

		/// <summary>
		/// Sends a command to the remote process, launching it if necessary, and giving it focus.
		/// </summary>
		/// <returns>Task that completes when the connection is established or fails, or the remote process exits.</returns>
		/// <param name="token">Cancellation token.</param>
		/// <param name="focus">Whether to focus the remote process.</param>
		/// <param name="name">Command. May be null.</param>
		/// <param name="args">Format arguments for the command.</param>
		protected Task<bool> ConnectAndSend (CancellationToken token, bool focus, string name, params object[] args)
		{
			if (IsConnected) {
				bool success = !focus || GiveFocusToRemoteProcess ();
				if (name == null)
					return TaskFromResult (success);

				//don't try to reconnect on send failures since the reconnect would likely break too
				//ignore the tiny race that could happen if the remote process quits during the send
				return Send (name, args);
			}

			return LogExceptions (Connect (token).ContinueWith (t => {
				t.Wait ();
				if (name == null)
					return true;
				return SendSync (name, args);
			}));
		}

		static Task<T> TaskFromResult<T>(T result)
		{
			var tcs = new TaskCompletionSource<T> ();
			tcs.SetResult (result);
			return tcs.Task;
		}

		//ensures exceptions are observed
		static Task<T> LogExceptions<T> (Task<T> task)
		{
			task.ContinueWith (t => {
				UnrealAgentHelper.Error (t.Exception.Flatten ().InnerException, "Error in task");
			}, TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously);

			return task;
		}

		bool GiveFocusToRemoteProcess ()
		{
			var c = connection;
			if (c == null)
				return false;

			return GiveFocusToProcess (c.RemotePid);
		}

		public static bool GiveFocusToProcess (int pid)
		{
			if (UnrealAgentHelper.IsWindows)
				return GiveFocusToRemoteWindowsProcess (pid);
			if (UnrealAgentHelper.IsMac)
				return GiveFocusToRemoteMacProcess (pid);
			throw new NotImplementedException ();
		}

		static bool GiveFocusToRemoteWindowsProcess (int pid)
		{
			IntPtr handle;
			if (GetMainWindowHandle (pid, out handle) && SetForegroundWindow (handle)) {
				SetFocus (handle);
				return true;
			}
			return false;
		}

		[DllImport ("user32.dll")]
		[return: MarshalAs (UnmanagedType.Bool)]
        static extern bool SetForegroundWindow (IntPtr hWnd);

		[DllImport ("user32.dll")]
        static extern IntPtr SetFocus (IntPtr hWnd);

		//Process.MainWindowHandle doesn't work on Mono
		static bool GetMainWindowHandle (int pid, out IntPtr handle)
		{
			var result = IntPtr.Zero;

			EnumWindows ((hWnd, lParam) => {
				uint winPid;
				if (GetWindowThreadProcessId (hWnd, out winPid) != 0 && pid == winPid && IsWindowVisible (hWnd)) {
					result = hWnd;
					return false;
				}
				//return true means continue enumerating
				return true;
			}, IntPtr.Zero);

			handle = result;
			return result != IntPtr.Zero;
		}

		[DllImport ("user32.dll")]
		[return: MarshalAs (UnmanagedType.Bool)]
		static extern bool IsWindowVisible (IntPtr hWnd);

		[DllImport("user32.dll")]
		[return: MarshalAs(UnmanagedType.Bool)]
		static extern bool EnumWindows (EnumWindowsProc lpEnumFunc, IntPtr lParam);

		delegate bool EnumWindowsProc (IntPtr hWnd, IntPtr lParam);

		[DllImport("user32.dll", SetLastError=true)]
		static extern uint GetWindowThreadProcessId(IntPtr hWnd, out uint lpdwProcessId);

		static bool GiveFocusToRemoteMacProcess (int pid)
		{
			var args = string.Format (
				"-e \"tell application \\\"System Events\\\" to set frontmost of the first process whose unix id is {0} to true\"",
				pid
			);
			Process.Start (new ProcessStartInfo ("/usr/bin/osascript", args) { UseShellExecute = false });
			return true;
		}
	}
}

UnrealAgentConnection.cs

// Copyright 2014 Xamarin Inc.
// For details, see LICENSE.txt.

using System;
using System.IO;
using System.Net.Sockets;
using System.Text;
using System.Diagnostics;
using System.Collections.Generic;

#if AGENT_CLIENT
//for ProcessArgumentBuilder only
using MonoDevelop.Core.Execution;
#endif

#if AGENT_CLIENT
namespace XamarinStudio.Unreal.Projects
#else
namespace UnrealEngine.MainDomain
#endif
{
	class UnrealAgentConnection : IDisposable
	{
		const string protocolVersion = "1.0";

		#if AGENT_CLIENT
		const string sendHandshake = "XAMARIN UNREAL AGENT " + protocolVersion + " CLIENT";
		const string recvHandshake = "XAMARIN UNREAL AGENT " + protocolVersion + " SERVER";
		#else
		const string sendHandshake = "XAMARIN UNREAL AGENT " + protocolVersion + " SERVER";
		const string recvHandshake = "XAMARIN UNREAL AGENT " + protocolVersion + " CLIENT";
        #endif

		bool disposed;

		TcpClient client;
		TextWriter writer;
		TextReader reader;
		int remotePid;

		readonly Func<string,string[],bool> commandHandler;
		readonly Action<UnrealAgentConnection> disposedHandler;

		public UnrealAgentConnection (
			TcpClient client,
			Func<string,string[],bool> commandHandler,
			Action<UnrealAgentConnection> disposedHandler)
		{
			this.client = client;
			this.commandHandler = commandHandler;
			this.disposedHandler = disposedHandler;
		}

		public int RemotePid { get { return remotePid; } }

		public void Run ()
		{
			try {
				if (!Connect ())
					return;

				try {
					commandHandler ("Connected", null);
				} catch (Exception ex) {
					if (!disposed)
						UnrealAgentHelper.Error (ex, "Error handling command Connected");
				}

				string line;
				while ((line = Read ()) != null) {
					string command;
					string[] args;
					ParseCommand (line, out command, out args);

					//TODO: this should eventually be less verbose
					UnrealAgentHelper.Log ("Command {0}", line);

					if (command == "CLOSE"){
						UnrealAgentHelper.Log ("Disconnecting.");
						return;
					}

					try {
						if (commandHandler (command, args))
							continue;
					} catch (Exception ex) {
						if (!disposed)
							UnrealAgentHelper.Error (ex, "Error handling command " + command);
					}

					UnrealAgentHelper.Error ("Unknown command " + command);
				}
			} catch (Exception ex) {
				if (!disposed)
					UnrealAgentHelper.Error (ex, "Unhandled error in agent thread");
			} finally {
				Dispose ();
			}
		}

		bool Connect ()
		{
			var stream = client.GetStream ();

			//don't wait forever on writes if other end is unresponsive
			stream.WriteTimeout = 10000;

			writer = new StreamWriter (stream, Encoding.UTF8);
			reader = new StreamReader (stream, Encoding.UTF8);

			var pid = Process.GetCurrentProcess ().Id;

			if (!Write (sendHandshake + " " + pid))
				return false;

			var line = Read ();
			if (line == null)
				return false;

			if (!line.StartsWith (recvHandshake, StringComparison.Ordinal)
				|| !int.TryParse (line.Substring (recvHandshake.Length).Trim (), out remotePid))
			{
				UnrealAgentHelper.Error ("Bad handshake.");
				return false;
			}

			UnrealAgentHelper.Log ("Handshake succeeded, connected to {0}", remotePid);
			return true;
		}

		bool Write (string line)
		{
			var w = writer;
			if (w == null)
				return false;
			try {
				lock (w) {
					w.WriteLine (line);
					w.Flush ();
				}
				return true;
			} catch (Exception ex) {
				if (!disposed) {
					UnrealAgentHelper.Error (ex, "Unhandled error in agent");
					Dispose ();
				}
			}
			return false;
		}

		string Read ()
		{
			var r = reader;
			if (r == null)
				return null;
			try {
				return r.ReadLine ();
			} catch {
				if (disposed)
					return null;
				throw;
			}
		}

		static void ParseCommand (string line, out string command, out string[] args)
		{
			int idx = line.IndexOf (' ');
			if (idx < 0) {
				command = line;
				args = new string[0];
				return;
			}
			command = line.Substring (0, idx);

			args = ProcessArgumentBuilder.Parse(line.Substring(idx + 1));
		}

		public bool Send (string name, params object[] args)
		{
			if (args.Length == 0)
				return Write (name);

			var pb = new ProcessArgumentBuilder ();
			pb.Add(name);
			foreach (object o in args) {
				pb.AddQuoted (o.ToString ());
			}

			return Write (pb.ToString ());
		}

		public void Close ()
		{
			UnrealAgentHelper.Log ("Disconnecting");
			try {
				Write ("CLOSE");
				// Analysis disable once EmptyGeneralCatchClause
			} catch {
			}

			Dispose ();
		}

		public void Dispose ()
		{
			if (disposed)
				return;
			disposed = true;

			var w = writer;
			if (w != null) {
				writer = null;
				w.Dispose ();
			}

			var r = reader;
			if (r != null) {
				reader = null;
				r.Dispose ();
			}

			var c = client;
			if (c != null) {
				client = null;
				try {
					c.Close ();
				} catch (Exception ex) {
					UnrealAgentHelper.Error (ex, "Unhandled error closing connection");
				}
			}

			disposedHandler (this);
		}
	}
}