Framework for testing Swarm protocol

src/Tgstation.Server.Host/Controllers/IRequestSwarmRegistrationParser.cs

@@ -0,0 +1,19 @@
+using System;
+
+using Microsoft.AspNetCore.Http;
+
+namespace Tgstation.Server.Host.Controllers
+{
+	/// <summary>
+	/// Parses the swarm registration header from a <see cref="HttpRequest"/>.
+	/// </summary>
+	public interface IRequestSwarmRegistrationParser
+	{
+		/// <summary>
+		/// Gets the swarm registration <see cref="Guid"/> from the headers of a given <paramref name="request"/>.
+		/// </summary>
+		/// <param name="request">The <see cref="HttpRequest"/>, must contain a valid <see cref="Swarm.SwarmConstants.RegistrationIdHeader"/>.</param>
+		/// <returns>The parsed registration ID.</returns>
+		Guid GetRequestRegistrationId(HttpRequest request);
+	}
+}

src/Tgstation.Server.Host/Controllers/RequestSwarmRegistrationParser.cs

@@ -0,0 +1,22 @@
+using System;
+using System.Linq;
+
+using Microsoft.AspNetCore.Http;
+
+using Tgstation.Server.Host.Swarm;
+
+namespace Tgstation.Server.Host.Controllers
+{
+	/// <inheritdoc />
+	sealed class RequestSwarmRegistrationParser : IRequestSwarmRegistrationParser
+	{
+		/// <inheritdoc />
+		public Guid GetRequestRegistrationId(HttpRequest request)
+		{
+			if (request == null)
+				throw new ArgumentNullException(nameof(request));
+
+			return Guid.Parse(request.Headers[SwarmConstants.RegistrationIdHeader].First());
+		}
+	}
+}

src/Tgstation.Server.Host/Controllers/SwarmController.cs

@@ -29,13 +29,18 @@ namespace Tgstation.Server.Host.Controllers
 		/// <summary>
 		/// Get the current registration <see cref="Guid"/> from the <see cref="ControllerBase.Request"/>.
 		/// </summary>
-		Guid RequestRegistrationId => Guid.Parse(Request.Headers[SwarmConstants.RegistrationIdHeader].First());
+		internal Guid RequestRegistrationId => requestRegistrationParser.GetRequestRegistrationId(Request);
 
 		/// <summary>
 		/// The <see cref="ISwarmOperations"/> for the <see cref="SwarmController"/>.
 		/// </summary>
 		readonly ISwarmOperations swarmOperations;
 
+		/// <summary>
+		/// The <see cref="IRequestSwarmRegistrationParser"/> for the <see cref="SwarmController"/>.
+		/// </summary>
+		readonly IRequestSwarmRegistrationParser requestRegistrationParser;
+
 		/// <summary>
 		/// The <see cref="IAssemblyInformationProvider"/> for the <see cref="SwarmController"/>.
 		/// </summary>
@@ -55,16 +60,19 @@ namespace Tgstation.Server.Host.Controllers
 		/// Initializes a new instance of the <see cref="SwarmController"/> class.
 		/// </summary>
 		/// <param name="swarmOperations">The value of <see cref="swarmOperations"/>.</param>
+		/// <param name="requestRegistrationParser">The value of <see cref="requestRegistrationParser"/>.</param>
 		/// <param name="assemblyInformationProvider">The value of <see cref="assemblyInformationProvider"/>.</param>
 		/// <param name="swarmConfigurationOptions">The <see cref="IOptions{TOptions}"/> containing the value of <see cref="swarmConfiguration"/>.</param>
 		/// <param name="logger">The value of <see cref="logger"/>.</param>
 		public SwarmController(
 			ISwarmOperations swarmOperations,
+			IRequestSwarmRegistrationParser requestRegistrationParser,
 			IAssemblyInformationProvider assemblyInformationProvider,
 			IOptions<SwarmConfiguration> swarmConfigurationOptions,
 			ILogger<SwarmController> logger)
 		{
 			this.swarmOperations = swarmOperations ?? throw new ArgumentNullException(nameof(swarmOperations));
+			this.requestRegistrationParser = requestRegistrationParser ?? throw new ArgumentNullException(nameof(requestRegistrationParser));
 			this.assemblyInformationProvider = assemblyInformationProvider ?? throw new ArgumentNullException(nameof(assemblyInformationProvider));
 			swarmConfiguration = swarmConfigurationOptions?.Value ?? throw new ArgumentNullException(nameof(swarmConfigurationOptions));
 			this.logger = logger;

src/Tgstation.Server.Host/Core/Application.cs

@@ -360,6 +360,7 @@ namespace Tgstation.Server.Host.Core
 			services.AddSingleton<ISwarmOperations>(x => x.GetRequiredService<SwarmService>());
 			services.AddSingleton<IServerUpdater, ServerUpdater>();
 			services.AddSingleton<IServerUpdateInitiator, ServerUpdateInitiator>();
+			services.AddSingleton<IRequestSwarmRegistrationParser, RequestSwarmRegistrationParser>();
 
 			// configure root services
 			services.AddSingleton<IJobManager, JobManager>();

src/Tgstation.Server.Host/Swarm/SwarmService.cs

@@ -51,7 +51,7 @@ namespace Tgstation.Server.Host.Swarm
 		/// <summary>
 		/// See <see cref="JsonSerializerSettings"/> for the swarm system.
 		/// </summary>
-		static readonly JsonSerializerSettings SerializerSettings = new ()
+		internal static JsonSerializerSettings SerializerSettings { get; } = new ()
 		{
 			ContractResolver = new DefaultContractResolver
 			{

tests/Tgstation.Server.Host.Tests/Swarm/SwarmRpcMapper.cs

@@ -0,0 +1,176 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net;
+using System.Net.Http;
+using System.Reflection;
+using System.Threading;
+using System.Threading.Tasks;
+
+using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Mvc;
+using Microsoft.AspNetCore.Mvc.Routing;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+using Moq;
+
+using Newtonsoft.Json;
+
+using Tgstation.Server.Common;
+using Tgstation.Server.Host.Configuration;
+using Tgstation.Server.Host.Controllers;
+using Tgstation.Server.Host.Swarm;
+
+namespace Tgstation.Server.Host.Tests.Swarm
+{
+	sealed class SwarmRpcMapper : IRequestSwarmRegistrationParser
+	{
+		List<(SwarmConfiguration, TestableSwarmNode)> configToControllers;
+		Guid? incomingRegistrationId;
+
+		public SwarmRpcMapper(Mock<IHttpClient> clientMock)
+		{
+			clientMock
+				.Setup(x => x.SendAsync(It.IsNotNull<HttpRequestMessage>(), It.IsAny<CancellationToken>()))
+				.Returns(MapRequest);
+		}
+
+		public Guid GetRequestRegistrationId(HttpRequest request)
+		{
+			Assert.IsTrue(incomingRegistrationId.HasValue);
+			var result = incomingRegistrationId.Value;
+			incomingRegistrationId = null;
+			return result;
+		}
+
+		public void Register(List<(SwarmConfiguration, TestableSwarmNode)> configToControllers)
+		{
+			this.configToControllers = configToControllers;
+		}
+
+		async Task<HttpResponseMessage> MapRequest(
+			HttpRequestMessage request,
+			CancellationToken cancellationToken)
+		{
+			var (config, node) = configToControllers.FirstOrDefault(
+				pair => pair.Item1.Address.IsBaseOf(request.RequestUri));
+
+			if (config == default)
+				Assert.Fail($"Invalid node address: {request.RequestUri}");
+
+			if (!node.Initialized)
+			{
+				throw new HttpRequestException("Can't connect to uninitialized node!");
+			}
+
+			var controller = node.Controller;
+
+			Type targetAttribute = null;
+			bool isDataRequest = false;
+			switch (request.Method.Method.ToUpperInvariant())
+			{
+				case "GET":
+					targetAttribute = typeof(HttpGetAttribute);
+					break;
+				case "POST":
+					targetAttribute = typeof(HttpPostAttribute);
+					isDataRequest = true;
+					break;
+				case "PUT":
+					targetAttribute = typeof(HttpPutAttribute);
+					isDataRequest = true;
+					break;
+				case "DELETE":
+					targetAttribute = typeof(HttpDeleteAttribute);
+					break;
+				case "PATCH":
+					targetAttribute = typeof(HttpPatchAttribute);
+					isDataRequest = true;
+					break;
+				default:
+					Assert.Fail($"Unknown request method: {request.Method.Method}");
+					break;
+			}
+
+			var stringUrl = request.RequestUri.ToString();
+			var rootIndex = stringUrl.IndexOf(SwarmConstants.ControllerRoute);
+			if (rootIndex == -1)
+				Assert.Fail($"Invalid Swarm route: {stringUrl}");
+
+			var route = stringUrl[(rootIndex + SwarmConstants.ControllerRoute.Length)..].TrimStart('/');
+
+			var controllerMethod = controller
+				.GetType()
+				.GetMethods()
+				.Select(method => (method, (HttpMethodAttribute)method.GetCustomAttribute(targetAttribute)))
+				.Where(pair => pair.Item2 != null
+					&& pair.Item2.HttpMethods.Count() == 1
+					&& pair.Item2.HttpMethods.All(supportedMethod => supportedMethod.Equals(request.Method.Method))
+					&& pair.Item2.Template == route)
+				.Select(pair => pair.method)
+				.SingleOrDefault();
+
+			if (controllerMethod == default)
+				Assert.Fail($"SwarmController has no method with attribute {targetAttribute}!");
+
+			// We're not testing OnActionExecutingAsync, that's covered by integration.
+			if (request.Headers.TryGetValues(SwarmConstants.RegistrationIdHeader, out var values) && values.Count() == 1)
+				node.RpcMapper.incomingRegistrationId = Guid.Parse(values.First());
+
+			var args = new List<object>();
+			if (isDataRequest)
+			{
+				var dataType = controllerMethod.GetParameters().First().ParameterType;
+				var json = await request.Content.ReadAsStringAsync(cancellationToken);
+				var parameter = JsonConvert.DeserializeObject(json, dataType, SwarmService.SerializerSettings);
+				args.Add(parameter);
+			}
+
+			IActionResult result;
+
+			var response = new HttpResponseMessage();
+			try
+			{
+				if (controllerMethod.ReturnType != typeof(IActionResult))
+				{
+					Assert.AreEqual(typeof(Task<IActionResult>), controllerMethod.ReturnType);
+					args.Add(cancellationToken);
+					var invocationTask = (Task<IActionResult>)controllerMethod.Invoke(controller, args.ToArray());
+					result = await invocationTask;
+				}
+				else
+				{
+					result = (IActionResult)controllerMethod.Invoke(controller, args.ToArray());
+
+					// simulate worst case, request completed but was aborted before server replied
+					cancellationToken.ThrowIfCancellationRequested();
+				}
+			}
+			catch
+			{
+				response.Dispose();
+				throw;
+			}
+
+			// manually checked all controller response types
+			// Fobid, NoContent, Conflict, StatusCode
+			if (result is ForbidResult forbidResult)
+				response.StatusCode = HttpStatusCode.Forbidden;
+			else if (result is NoContentResult noContentResult)
+				response.StatusCode = (HttpStatusCode)noContentResult.StatusCode;
+			else if (result is ConflictResult conflictResult)
+				response.StatusCode = (HttpStatusCode)conflictResult.StatusCode;
+			else if (result is ObjectResult objectResult)
+				response.StatusCode = (HttpStatusCode)objectResult.StatusCode;
+			else if (result is StatusCodeResult statusCodeResult)
+				response.StatusCode = (HttpStatusCode)statusCodeResult.StatusCode;
+			else
+			{
+				response.Dispose();
+				Assert.Fail($"Unrecognized result type: {result.GetType()}");
+			}
+
+			return response;
+		}
+	}
+}

tests/Tgstation.Server.Host.Tests/Swarm/TestSwarmProtocol.cs

@@ -0,0 +1,81 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+using Microsoft.Extensions.Logging;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+using Tgstation.Server.Host.Configuration;
+using Tgstation.Server.Host.Swarm;
+
+namespace Tgstation.Server.Host.Tests.Swarm
+{
+	[TestClass]
+	public sealed class TestSwarmProtocol : IDisposable
+	{
+		readonly ILoggerFactory loggerFactory = LoggerFactory.Create(builder => builder.AddConsole());
+		readonly HashSet<ushort> usedPorts = new ();
+
+		public void Dispose() => loggerFactory.Dispose();
+
+		[TestInitialize]
+		public void Initialize()
+		{
+			usedPorts.Clear();
+		}
+
+		[TestMethod]
+		public async Task TestInitHappensInstantlyWhenControllerIsInitialized()
+		{
+			using var controller = new TestableSwarmNode(loggerFactory, GenConfig());
+			using var node = new TestableSwarmNode(loggerFactory, GenConfig(controller.Config));
+
+			TestableSwarmNode.Link(controller, node);
+
+			var controllerInit = controller.TryInit();
+			Assert.IsTrue(controllerInit.IsCompleted);
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controllerInit);
+
+			var nodeInit = node.TryInit();
+			Assert.IsTrue(nodeInit.IsCompleted);
+			Assert.AreEqual(SwarmRegistrationResult.Success, await nodeInit);
+		}
+
+		[TestMethod]
+		public async Task TestNodeInitializeDoesNotWorkWithoutController()
+		{
+			using var controller = new TestableSwarmNode(loggerFactory, GenConfig());
+			using var node1 = new TestableSwarmNode(loggerFactory, GenConfig(controller.Config));
+			using var node2 = new TestableSwarmNode(loggerFactory, GenConfig(controller.Config));
+
+			TestableSwarmNode.Link(controller, node1, node2);
+
+			Assert.AreEqual(SwarmRegistrationResult.CommunicationFailure, await node1.TryInit());
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node2.TryInit());
+		}
+
+		SwarmConfiguration GenConfig(SwarmConfiguration controllerConfig = null)
+		{
+			ushort randPort;
+			do
+			{
+				randPort = (ushort)(Random.Shared.Next() % UInt16.MaxValue);
+			}
+			while (randPort == 0 || !usedPorts.Add(randPort));
+
+			return new SwarmConfiguration
+			{
+				Address = new Uri($"http://127.0.0.1:{randPort}"),
+				ControllerAddress = controllerConfig?.Address,
+				Identifier = $"{(controllerConfig == null ? "Controller" : "Node")}{usedPorts.Count}",
+				PrivateKey = "asdf",
+				UpdateRequiredNodeCount = (uint)usedPorts.Count - 1,
+			};
+		}
+	}
+}

tests/Tgstation.Server.Host.Tests/Swarm/TestableSwarmNode.cs

@@ -0,0 +1,174 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net.Http;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+using Elasticsearch.Net;
+
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+using Moq;
+using Moq.Language.Flow;
+
+using Tgstation.Server.Common;
+using Tgstation.Server.Host.Configuration;
+using Tgstation.Server.Host.Controllers;
+using Tgstation.Server.Host.Core;
+using Tgstation.Server.Host.Database;
+using Tgstation.Server.Host.Swarm;
+using Tgstation.Server.Host.System;
+
+namespace Tgstation.Server.Host.Tests.Swarm
+{
+	sealed class TestableSwarmNode : IDisposable, IServerUpdateExecutor
+	{
+		public SwarmController Controller { get; }
+
+		public SwarmService Service { get; }
+
+		public SwarmConfiguration Config { get; }
+
+		public SwarmRpcMapper RpcMapper { get; }
+
+		public bool Initialized { get; private set; }
+
+		readonly Mock<IHttpClient> mockHttpClient;
+		readonly Mock<IServerControl> mockServerControl;
+		readonly Mock<IDatabaseContextFactory> mockDBContextFactory;
+		readonly Mock<IDatabaseSeeder> mockDatabaseSeeder;
+		readonly ISetup<IDatabaseSeeder, Task> mockDatabaseSeederInitialize;
+
+		public static void Link(params TestableSwarmNode[] nodes)
+		{
+			var configControllerSet = nodes.Select(x => (x.Config, x)).ToList();
+
+			_ = configControllerSet.Single(x => x.Config.ControllerAddress == null);
+			Assert.IsTrue(
+				configControllerSet.All(
+					tuple1 => !String.IsNullOrWhiteSpace(tuple1.Config.PrivateKey)
+					&& configControllerSet.All(tuple2 => tuple1.Config.PrivateKey == tuple2.Config.PrivateKey)),
+				"This test doesn't support authentication issues.");
+
+			foreach (var node in nodes)
+			{
+				if (node.Config.ControllerAddress != null)
+					node.Config.UpdateRequiredNodeCount = 0;
+				else
+					node.Config.UpdateRequiredNodeCount = (uint)nodes.Length - 1;
+				node.RpcMapper.Register(configControllerSet);
+			}
+		}
+
+		public TestableSwarmNode(
+			ILoggerFactory loggerFactory,
+			SwarmConfiguration swarmConfiguration,
+			Version mockVersion = null)
+		{
+			this.Config = swarmConfiguration;
+
+			var mockOptions = new Mock<IOptions<SwarmConfiguration>>();
+			mockOptions.SetupGet(x => x.Value).Returns(swarmConfiguration);
+
+			var realVersion = new AssemblyInformationProvider().Version;
+			var mockAssemblyInformationProvider = new Mock<IAssemblyInformationProvider>();
+			mockAssemblyInformationProvider.SetupGet(x => x.Version).Returns(mockVersion ?? realVersion);
+
+			var mockDatabaseContext = Mock.Of<IDatabaseContext>();
+
+			mockDatabaseSeeder = new Mock<IDatabaseSeeder>();
+			mockDatabaseSeederInitialize = new Mock<IDatabaseSeeder>().Setup(x => x.Initialize(mockDatabaseContext, It.IsAny<CancellationToken>()));
+			mockDBContextFactory = new Mock<IDatabaseContextFactory>();
+			mockDBContextFactory
+				.Setup(x => x.UseContext(It.IsNotNull<Func<IDatabaseContext, Task>>()))
+				.Callback<Func<IDatabaseContext, Task>>((func) => func(mockDatabaseContext));
+
+			mockServerControl = new Mock<IServerControl>();
+			mockServerControl.Setup(x => x.TryStartUpdate(this, It.IsNotNull<Version>())).Returns(TryStartUpdate);
+			mockServerControl.Setup(x => x.RegisterForRestart(It.IsNotNull<IRestartHandler>())).Returns(Mock.Of<IRestartRegistration>());
+
+			var mockHttpClientFactory = new Mock<IAbstractHttpClientFactory>();
+			mockHttpClient = new Mock<IHttpClient>();
+			mockHttpClientFactory.Setup(x => x.CreateClient()).Returns(mockHttpClient.Object);
+
+			var mockAsyncDelayer = new Mock<IAsyncDelayer>();
+			mockAsyncDelayer.Setup(
+				x => x.Delay(It.IsAny<TimeSpan>(), It.IsAny<CancellationToken>()))
+				.Returns<TimeSpan, CancellationToken>(
+					(delay, ct) => Task.Delay(TimeSpan.FromSeconds(1), ct));
+
+			var mockServerUpdater = new Mock<IServerUpdater>();
+
+			RpcMapper = new SwarmRpcMapper(mockHttpClient);
+
+			Service = new SwarmService(
+				mockDBContextFactory.Object,
+				mockDatabaseSeeder.Object,
+				mockAssemblyInformationProvider.Object,
+				mockHttpClientFactory.Object,
+				mockServerControl.Object,
+				mockServerUpdater.Object,
+				mockAsyncDelayer.Object,
+				mockOptions.Object,
+				loggerFactory.CreateLogger<SwarmService>());
+
+			Controller = new SwarmController(
+				Service,
+				RpcMapper,
+				mockAssemblyInformationProvider.Object,
+				mockOptions.Object,
+				loggerFactory.CreateLogger<SwarmController>());
+		}
+
+		public void Dispose() => Service.Dispose();
+
+		public async Task<SwarmRegistrationResult?> TryInit(bool cancel = false)
+		{
+			if (Initialized)
+				Assert.Fail("Initialized twice!");
+
+			if (!cancel)
+				mockDatabaseSeederInitialize.Returns(Task.CompletedTask).Verifiable();
+			else
+				mockDatabaseSeederInitialize.ThrowsAsync(new TaskCanceledException()).Verifiable();
+
+			Task<SwarmRegistrationResult> Invoke() => Service.Initialize(default);
+
+			SwarmRegistrationResult? result;
+			if (cancel)
+			{
+				await Assert.ThrowsExceptionAsync<OperationCanceledException>(Invoke);
+				result = null;
+			}
+			else
+			{
+				result = await Invoke();
+				Initialized = true;
+			}
+
+			if (Config.ControllerAddress == null)
+				mockDatabaseSeeder.VerifyAll();
+			else
+			{
+				Assert.IsFalse(mockDatabaseSeeder.Invocations.Any());
+				Assert.IsFalse(mockDBContextFactory.Invocations.Any());
+			}
+
+			return result;
+		}
+
+		bool TryStartUpdate(IServerUpdateExecutor updateExecutor, Version newVersion)
+		{
+			throw new NotImplementedException();
+		}
+
+		public Task<bool> ExecuteUpdate(string updatePath, CancellationToken cancellationToken, CancellationToken criticalCancellationToken)
+		{
+			throw new NotImplementedException();
+		}
+	}
+}