SocialPub/PrivaPub/Infrastructure/RateLimiting.cs

using Microsoft.AspNetCore.RateLimiting;

using PrivaPub.Federation.Objects;
using PrivaPub.Federation.Signing;
using PrivaPub.Infrastructure.Statistics;
using PrivaPub.Models.Statistics;

using System.Threading.RateLimiting;

namespace PrivaPub.Infrastructure
{
	public static class RateLimiting
	{
		public const string Accounts = "accounts";
		public const string Inbox = "inbox";

		public static IServiceCollection PrivaPubRateLimiting(this IServiceCollection service) =>
			service.AddRateLimiter(options =>
			{
				options.RejectionStatusCode = StatusCodes.Status429TooManyRequests;
				options.OnRejected = (context, _) =>
				{
					var http = context.HttpContext;
					var ledger = http.RequestServices.GetService<IInteractionLedger>();
					var policy = http.GetEndpoint()?.Metadata.GetMetadata<EnableRateLimitingAttribute>()?.PolicyName;
					if (policy == Inbox)
						ledger?.Record(new InteractionEvent
						{
							Channel = Interactions.Receive,
							Host = Interactions.HostOf(SenderOrigin(http.Request)),
							Status = StatusCodes.Status429TooManyRequests,
							Outcome = Interactions.Refused,
							Reason = "rate-limited",
							Inbox = http.Request.Path.Value?.EndsWith("/mouth", StringComparison.Ordinal) == true ? "personal" : "shared"
						}, hostClaimed: true);
					else
						ledger?.CountServer(ServerSections.Client, $"{policy ?? "unknown"}:429");
					return ValueTask.CompletedTask;
				};
				options.AddPolicy(Accounts, context => RateLimitPartition.GetFixedWindowLimiter(
					context.Connection.RemoteIpAddress?.ToString() ?? "unknown",
					_ => new FixedWindowRateLimiterOptions { PermitLimit = 10, Window = TimeSpan.FromMinutes(1), QueueLimit = 0 }));
				options.AddPolicy(Inbox, context => RateLimitPartition.GetTokenBucketLimiter(
					SenderOrigin(context.Request) ?? "unsigned:" + context.Connection.RemoteIpAddress,
					_ => new TokenBucketRateLimiterOptions
					{
						TokenLimit = 300,
						TokensPerPeriod = 50,
						ReplenishmentPeriod = TimeSpan.FromSeconds(10),
						QueueLimit = 0
					}));
			});

		static string SenderOrigin(HttpRequest request)
		{
			var signature = request.Headers["Signature"].ToString();
			return string.IsNullOrEmpty(signature) ? default : Origin.Of(HttpSignatures.Parse(signature)?.KeyId);
		}
	}
}