From 664042ec447089fbeff4a36c5660c60713859d64 Mon Sep 17 00:00:00 2001
From: Neo <neo@rayshift.io>
Date: Sun, 1 Nov 2020 02:24:02 +0000
Subject: [PATCH] Add endpoint routing

---
 KamihamaWeb/Controllers/APIController.cs | 31 ++++++++++++++++++++++--
 KamihamaWeb/Interfaces/IMasterService.cs |  1 +
 KamihamaWeb/Services/MasterService.cs    | 22 +++++++++++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/KamihamaWeb/Controllers/APIController.cs b/KamihamaWeb/Controllers/APIController.cs
index ac41d7d..bc5c9b7 100644
--- a/KamihamaWeb/Controllers/APIController.cs
+++ b/KamihamaWeb/Controllers/APIController.cs
@@ -2,10 +2,13 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
+using KamihamaWeb.Interfaces;
 using KamihamaWeb.Models;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.Extensions.Configuration;
+using Serilog;
 
 // ReSharper disable InconsistentNaming
 
@@ -15,17 +18,41 @@ namespace KamihamaWeb.Controllers
     [ApiController]
     public class APIController : ControllerBase
     {
-        public APIController(IConfiguration config)
+        public APIController(IConfiguration config, IMasterSingleton master)
         {
             _config = config;
+            _master = master;
         }
 
         private IConfiguration _config { get; set; }
+        private IMasterSingleton _master { get; set; }
         [Route("endpoint")]
         public IActionResult GetEndpoint()
         {
             var response = new APIResult(200, "ok");
-            response.Add("endpoint", _config["MagiRecoServer:Endpoint"]);
+            if (_master.Endpoints.Count == 0) // Endpoint server
+            {
+                response.Add("endpoint", _config["MagiRecoServer:Endpoint"]);
+            }
+            else // Master server
+            {
+                if (Request.Headers.ContainsKey("CF-IPCountry"))
+                {
+                    if (_master.Endpoints.ContainsKey(Request.Headers["CF-IPCountry"]))
+                    {
+                        response.Add("endpoint", _master.Endpoints[Request.Headers["CF-IPCountry"]]);
+                    }
+                    else
+                    {
+                        response.Add("endpoint", _master.Endpoints["*"]);
+                    }
+                }
+                else  {
+                    Log.Warning("No Cloudflare IP header found.");
+                    response.Add("endpoint", _master.Endpoints["*"]);
+                }
+            }
+
             response.Add("version", int.Parse(_config["MagiRecoServer:Version"]));
             response.Add("max_threads", int.Parse(_config["MagiRecoServer:MaxThreads"]));
             return response;
diff --git a/KamihamaWeb/Interfaces/IMasterService.cs b/KamihamaWeb/Interfaces/IMasterService.cs
index fa9e352..743ecc8 100644
--- a/KamihamaWeb/Interfaces/IMasterService.cs
+++ b/KamihamaWeb/Interfaces/IMasterService.cs
@@ -17,6 +17,7 @@ namespace KamihamaWeb.Interfaces
         public long AssetsCurrentVersion { get; set; }
         public Dictionary<string, GamedataAsset> EnglishMasterAssets { get; set; }
         public Dictionary<string, Dictionary<string, GamedataAsset>> GamedataAssets { get; set; }
+        public Dictionary<string, string> Endpoints { get; set; }
         public Task<string> ProvideJson(string which);
         public Task<bool> RunUpdate();
     }
diff --git a/KamihamaWeb/Services/MasterService.cs b/KamihamaWeb/Services/MasterService.cs
index d29bae2..2ae816c 100644
--- a/KamihamaWeb/Services/MasterService.cs
+++ b/KamihamaWeb/Services/MasterService.cs
@@ -45,10 +45,32 @@ namespace KamihamaWeb.Services
             _cache = ((RedisCache) cache).GetConnection().GetDatabase();
             _rest = rest;
             _builder = builder;
+
+            if (_config["MagiRecoServer:Type"] == "master")
+            {
+                Log.Information("This is a master server, populating endpoints.");
+
+                foreach (var item in _config.GetSection("MagiRecoNodes").Get<Dictionary<string, string>>())
+                {
+                    Log.Information($"{item.Key} -> {item.Value}");
+                    foreach (var endpoint in item.Key.Split(","))
+                    {
+                        Endpoints.Add(endpoint, item.Value);
+                    }
+                }
+
+                if (!Endpoints.ContainsKey("*"))
+                {
+                    Log.Fatal("Missing * endpoint in config! Please add one!");
+                    throw new Exception("Missing * endpoint in config! Please add one!");
+                }
+            }
             Task.Run(Initialize);
         }
         public Guid Guid { get; set; }
 
+        public Dictionary<string, string> Endpoints { get; set; } = new Dictionary<string, string>();
+
         public List<string> ModdedAssetLists = new List<string>()
         {
             "asset_main",