diff --git a/JetHerald/Configs.cs b/JetHerald/Configs.cs index 9f53084..33690ca 100644 --- a/JetHerald/Configs.cs +++ b/JetHerald/Configs.cs @@ -19,4 +19,11 @@ { public string Token { get; set; } } + + public class Timeout + { + public int DebtLimitSeconds { get; set; } + public int HeartbeatCost { get; set; } + public int ReportCost { get; set; } + } } diff --git a/JetHerald/Controllers/HeartbeatController.cs b/JetHerald/Controllers/HeartbeatController.cs index b187899..c6ae3c2 100644 --- a/JetHerald/Controllers/HeartbeatController.cs +++ b/JetHerald/Controllers/HeartbeatController.cs @@ -1,12 +1,10 @@ using System; -using System.IO; -using System.Runtime.Serialization; -using System.Text; using System.Text.Json; using System.Text.Json.Serialization; using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; using Microsoft.AspNetCore.Mvc; -using Microsoft.AspNetCore.Mvc.ModelBinding; +using Microsoft.Extensions.Options; namespace JetHerald.Controllers { @@ -15,11 +13,15 @@ namespace JetHerald.Controllers { Db Db { get; } JetHeraldBot Herald { get; } + LeakyBucket Timeouts { get; } + Options.Timeout Config { get; } - public HeartbeatController(Db db, JetHeraldBot herald) + public HeartbeatController(Db db, JetHeraldBot herald, LeakyBucket timeouts, IOptions cfgOptions) { Herald = herald; + Timeouts = timeouts; Db = db; + Config = cfgOptions.Value; } // tfw when you need to manually parse body and query @@ -71,12 +73,16 @@ namespace JetHerald.Controllers else if (!t.WriteToken.Equals(args.WriteToken, StringComparison.Ordinal)) return StatusCode(403); + if (Timeouts.IsTimedOut(t.TopicId)) + return StatusCode(StatusCodes.Status429TooManyRequests); var affected = await Db.ReportHeartbeat(t.TopicId, heart, args.ExpiryTimeout); if (affected == 1) await Herald.HeartbeatSent(t, heart); + Timeouts.ApplyCost(t.TopicId, Config.HeartbeatCost); + return new OkResult(); } diff --git a/JetHerald/Controllers/ReportController.cs b/JetHerald/Controllers/ReportController.cs index 9a675dd..dca583c 100644 --- a/JetHerald/Controllers/ReportController.cs +++ b/JetHerald/Controllers/ReportController.cs @@ -1,7 +1,9 @@ using System; using System.Text.Json; using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; using Microsoft.AspNetCore.Mvc; +using Microsoft.Extensions.Options; namespace JetHerald.Controllers { @@ -10,11 +12,15 @@ namespace JetHerald.Controllers { Db Db { get; } JetHeraldBot Herald { get; } + LeakyBucket Timeouts { get; } + Options.Timeout Config { get; } - public ReportController(Db db, JetHeraldBot herald) + public ReportController(Db db, JetHeraldBot herald, LeakyBucket timeouts, IOptions cfgOptions) { Herald = herald; + Timeouts = timeouts; Db = db; + Config = cfgOptions.Value; } [Route("api/report")] @@ -55,7 +61,13 @@ namespace JetHerald.Controllers else if (!t.WriteToken.Equals(args.WriteToken, StringComparison.OrdinalIgnoreCase)) return StatusCode(403); + if (Timeouts.IsTimedOut(t.TopicId)) + return StatusCode(StatusCodes.Status429TooManyRequests); + await Herald.PublishMessage(t, args.Message); + + Timeouts.ApplyCost(t.TopicId, Config.ReportCost); + return new OkResult(); } diff --git a/JetHerald/LeakyBucket.cs b/JetHerald/LeakyBucket.cs new file mode 100644 index 0000000..f213c66 --- /dev/null +++ b/JetHerald/LeakyBucket.cs @@ -0,0 +1,39 @@ +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using Microsoft.Extensions.Options; + +namespace JetHerald +{ + public class LeakyBucket + { + private readonly ConcurrentDictionary expiryDates = new(); + private readonly Options.Timeout config; + + public LeakyBucket(IOptions cfgOptions) + { + config = cfgOptions.Value; + } + + public bool IsTimedOut(uint key) + { + var debtLimit = DateTime.Now.AddSeconds(config.DebtLimitSeconds); + var time = expiryDates.GetValueOrDefault(key, DateTime.Now); + Console.WriteLine(time); + return time > debtLimit; + } + + public void ApplyCost(uint key, int cost) + { + expiryDates.AddOrUpdate(key, + key => DateTime.Now.AddSeconds(cost), + (key, oldDebt) => + { + if (oldDebt < DateTime.Now) + return DateTime.Now.AddSeconds(cost); + + return oldDebt.AddSeconds(cost); + }); + } + } +} \ No newline at end of file diff --git a/JetHerald/Startup.cs b/JetHerald/Startup.cs index 793a7d9..68de981 100644 --- a/JetHerald/Startup.cs +++ b/JetHerald/Startup.cs @@ -22,8 +22,10 @@ namespace JetHerald services.Configure(Configuration.GetSection("ConnectionStrings")); services.Configure(Configuration.GetSection("Telegram")); services.Configure(Configuration.GetSection("Discord")); + services.Configure(Configuration.GetSection("Timeout")); services.AddSingleton(); services.AddSingleton(); + services.AddSingleton(); services.AddMvc().SetCompatibilityVersion(CompatibilityVersion.Version_3_0); } diff --git a/JetHerald/appsettings.json b/JetHerald/appsettings.json index a605bc3..8ead099 100644 --- a/JetHerald/appsettings.json +++ b/JetHerald/appsettings.json @@ -7,5 +7,10 @@ "Telegram": { "UseProxy": "true" }, + "Timeout": { + "DebtLimitSeconds": 600, + "HeartbeatCost": 60, + "ReportCost": 60 + }, "PathBase": "/" } \ No newline at end of file