Nextus-Web/src/main/java/dev/coph/nextusweb/server/ratelimit/RateLimitGate.java

package dev.coph.nextusweb.server.ratelimit;

import dev.coph.nextusweb.server.router.Request;
import dev.coph.nextusweb.server.router.Response;

import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Request-pipeline entry point that applies a {@link RateLimitConfig} to incoming requests and
 * surfaces the outcome as standard {@code X-RateLimit-*} response headers.
 *
 * <p>For each request the gate evaluates every {@link RateLimitConfig.Rule rule} that applies
 * to the request path. If any rule denies the request, evaluation stops and that denial is
 * returned; otherwise the strictest (lowest remaining) allowance is returned so the headers
 * reflect the tightest applicable budget.</p>
 *
 * <p>A daemon background thread periodically triggers cleanup of stale limiter state. The gate
 * should be {@link #shutdown() shut down} when the server stops.</p>
 */
public final class RateLimitGate {

    /**
     * Default idle age after which per-key limiter state is eligible for eviction.
     */
    private static final long DEFAULT_STALE_AFTER_NANOS = 10L * 60 * 1_000_000_000L;

    /**
     * The rule set this gate enforces.
     */
    private final RateLimitConfig config;
    /**
     * Idle age (nanoseconds) after which a limiter's per-key state may be evicted.
     */
    private final long staleAfterNanos;
    /**
     * Single-threaded scheduler driving periodic cleanup of stale buckets.
     */
    private final ScheduledExecutorService cleanup;

    /**
     * Creates a gate for the given configuration and starts a background cleanup task that runs
     * every five minutes on a daemon thread, evicting per-key state idle for more than ten
     * minutes.
     *
     * @param config the rate-limit rules to enforce
     */
    public RateLimitGate(RateLimitConfig config) {
        this(config, DEFAULT_STALE_AFTER_NANOS);
    }

    /**
     * Creates a gate with an explicit idle age before per-key limiter state is evicted.
     *
     * @param config          the rate-limit rules to enforce
     * @param staleAfterNanos idle age in nanoseconds after which per-key state is evicted; must
     *                        be positive
     */
    public RateLimitGate(RateLimitConfig config, long staleAfterNanos) {
        if (staleAfterNanos <= 0) throw new IllegalArgumentException("staleAfterNanos must be > 0");
        this.config = config;
        this.staleAfterNanos = staleAfterNanos;
        this.cleanup = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread t = new Thread(r, "ratelimit-cleanup");
            t.setDaemon(true);
            return t;
        });
        cleanup.scheduleAtFixedRate(this::doCleanup, 5, 5, TimeUnit.MINUTES);
    }

    /**
     * Periodic cleanup hook invoked by the background scheduler. Asks every configured limiter to
     * evict per-key state idle for longer than {@link #staleAfterNanos}. A failure cleaning one
     * limiter must not abort the others or kill the scheduler, so each call is guarded.
     */
    private void doCleanup() {
        for (RateLimiter limiter : config.allLimiters()) {
            try {
                limiter.cleanup(staleAfterNanos);
            } catch (RuntimeException ignored) {
            }
        }
    }

    /**
     * Writes the standard rate-limit headers ({@code X-RateLimit-Limit},
     * {@code X-RateLimit-Remaining}, and {@code Retry-After} when denied) onto a response.
     *
     * <p>Does nothing when {@code result} is {@code null} (no rule applied). The retry hint is
     * rounded up to whole seconds as required by the {@code Retry-After} header.</p>
     *
     * @param result the limiting result, may be {@code null}
     * @param res    the response to decorate
     */
    public static void applyHeaders(RateLimiter.Result result, Response res) {
        if (result == null) return;
        res.header("X-RateLimit-Limit", String.valueOf(result.limit()));
        res.header("X-RateLimit-Remaining", String.valueOf(Math.max(0, result.remaining())));
        if (!result.allowed()) {
            res.header("Retry-After", String.valueOf((result.retryAfterMillis() + 999) / 1000));
        }
    }

    /**
     * Evaluates all rules applicable to the given path and decides whether the request may
     * proceed.
     *
     * <p>Each rule's key is namespaced with the rule name to keep buckets from different rules
     * independent. The first denial short-circuits and is returned immediately; if every rule
     * allows the request, the result with the least remaining quota is returned.</p>
     *
     * @param req      the incoming request, used by key resolvers
     * @param path     the request path used to select rules
     * @param clientIp the resolved client IP (honouring trusted proxies), used as a key-resolver
     *                 fallback
     * @return the limiting result, or {@code null} if no rule applies to the path
     */
    public RateLimiter.Result check(Request req, String path, String clientIp) {
        List<RateLimitConfig.Rule> rules = config.rulesFor(path);
        if (rules.isEmpty()) return null;

        long now = System.nanoTime();
        RateLimiter.Result strictest = null;

        for (var rule : rules) {
            String key = rule.name() + ":" + rule.keyResolver().resolve(req, clientIp);
            RateLimiter.Result result = rule.limiter().tryAcquire(key, now);

            if (!result.allowed()) return result;

            if (strictest == null || result.remaining() < strictest.remaining()) {
                strictest = result;
            }
        }
        return strictest;
    }

    /**
     * Stops the background cleanup scheduler. Should be called when the server shuts down.
     */
    public void shutdown() {
        cleanup.shutdown();
    }
}