Harden MCP request limits, key lifecycle, and connection guides

README.md

@@ -57,6 +57,12 @@ includes successful requests. Tool calls are additionally limited to 120 per
 minute for each access-key and tool pair. Limits are process-local and therefore
 apply independently to each Halo replica.
 
+Authenticated request handling is cancelled after a 30-second deadline, and
+concurrent authenticated requests are capped at 100 globally and 16 per access
+key; requests beyond the cap receive `429 Too Many Requests`. Attachment uploads
+reserve their decoded size against in-flight byte budgets of 64 MiB globally and
+32 MiB per key before decoding, and excess uploads fail with `RATE_LIMITED`.
+
 Use a dedicated, least-privilege key. Do not put keys in URLs, configuration
 files committed to source control, shell history, or logs.
 

src/main/java/run/halo/mcpserver/HaloMcpServer.java

@@ -6,7 +6,6 @@
 import io.modelcontextprotocol.server.McpStatelessAsyncServer;
 import io.modelcontextprotocol.server.transport.DefaultServerTransportSecurityValidator;
 import io.modelcontextprotocol.spec.McpSchema;
-import java.time.Duration;
 import org.springframework.ai.mcp.server.webflux.transport.WebFluxStatelessServerTransport;
 import org.springframework.stereotype.Component;
 import org.springframework.web.reactive.function.server.RouterFunction;
@@ -53,7 +52,6 @@ class HaloMcpServer {
                 .capabilities(McpSchema.ServerCapabilities.builder()
                         .tools(false)
                         .build())
-                .requestTimeout(Duration.ofSeconds(30))
                 .tools(builtInTools.specifications())
                 .build();
     }

src/main/java/run/halo/mcpserver/KeyedInFlightLimiter.java

@@ -0,0 +1,72 @@
+package run.halo.mcpserver;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+/** Process-local global and per-key in-flight budgets. */
+public final class KeyedInFlightLimiter {
+
+    private final int perKeyLimit;
+    private final int maxTrackedKeys;
+    private final Semaphore globalPermits;
+    private final ConcurrentHashMap<String, Semaphore> keyPermits = new ConcurrentHashMap<>();
+
+    public KeyedInFlightLimiter(int globalLimit, int perKeyLimit, int maxTrackedKeys) {
+        this.perKeyLimit = perKeyLimit;
+        this.maxTrackedKeys = maxTrackedKeys;
+        this.globalPermits = new Semaphore(globalLimit);
+    }
+
+    public Lease tryAcquire(String keyId, int permits) {
+        if (permits < 0 || keyPermits.size() >= maxTrackedKeys && !keyPermits.containsKey(keyId)) {
+            return null;
+        }
+        if (!globalPermits.tryAcquire(permits)) {
+            return null;
+        }
+        var acquired = new AtomicReference<Semaphore>();
+        keyPermits.compute(keyId, (key, existing) -> {
+            var semaphore = existing == null ? new Semaphore(perKeyLimit) : existing;
+            if (!semaphore.tryAcquire(permits)) {
+                return existing;
+            }
+            acquired.set(semaphore);
+            return semaphore;
+        });
+        var semaphore = acquired.get();
+        if (semaphore == null) {
+            globalPermits.release(permits);
+            return null;
+        }
+        return new Lease(keyId, semaphore, permits);
+    }
+
+    public final class Lease implements AutoCloseable {
+
+        private final String keyId;
+        private final Semaphore semaphore;
+        private final int permits;
+        private final AtomicBoolean released = new AtomicBoolean();
+
+        private Lease(String keyId, Semaphore semaphore, int permits) {
+            this.keyId = keyId;
+            this.semaphore = semaphore;
+            this.permits = permits;
+        }
+
+        @Override
+        public void close() {
+            if (released.compareAndSet(false, true)) {
+                globalPermits.release(permits);
+                keyPermits.compute(keyId, (key, mapped) -> {
+                    semaphore.release(permits);
+                    return mapped == semaphore && semaphore.availablePermits() == perKeyLimit
+                            ? null
+                            : mapped;
+                });
+            }
+        }
+    }
+}

src/main/java/run/halo/mcpserver/McpAccessKeyService.java

@@ -127,14 +127,28 @@ Mono<McpKeyAuthenticationToken> authenticate(
                         .filter(Boolean::booleanValue)
                         .filter(ignored -> McpIpAllowlist.allows(
                                 accessKey.getSpec().getAllowedIpRanges(), remoteAddress))
-                        .flatMap(ignored -> touch(accessKey).thenReturn(new McpKeyAuthenticationToken(
+                        .flatMap(ignored -> revalidate(accessKey))
+                        .flatMap(fresh -> touch(fresh).thenReturn(new McpKeyAuthenticationToken(
                                 parsed.id(),
-                                accessKey.getSpec().getDisplayName(),
-                                accessKey.getSpec().getKeyPrefix(),
-                                accessKey.getSpec().getOwnerName(),
-                                accessKey.getSpec().getAllowedTools() == null
+                                fresh.getSpec().getDisplayName(),
+                                fresh.getSpec().getKeyPrefix(),
+                                fresh.getSpec().getOwnerName(),
+                                fresh.getSpec().getAllowedTools() == null
                                         ? Set.of()
-                                        : accessKey.getSpec().getAllowedTools()))));
+                                        : fresh.getSpec().getAllowedTools()))));
+    }
+
+    /**
+     * Re-fetches the key after asynchronous password verification so a rotation, disablement,
+     * scope change, or deletion committed meanwhile invalidates this authentication. Status-only
+     * writes such as last-used updates do not affect the comparison. This relies on the extension
+     * client returning a freshly deserialized instance per fetch; a shared mutable instance would
+     * make the spec comparison trivially equal and silently defeat this check.
+     */
+    private Mono<McpAccessKey> revalidate(McpAccessKey snapshot) {
+        return client.fetch(McpAccessKey.class, snapshot.getMetadata().getName())
+                .filter(this::active)
+                .filter(fresh -> fresh.getSpec().equals(snapshot.getSpec()));
     }
 
     private Mono<McpAccessKey> get(String id) {

src/main/java/run/halo/mcpserver/McpAuthorization.java

@@ -32,6 +32,10 @@ public <T> Mono<T> authorize(String toolName, Supplier<Mono<T>> action) {
         }));
     }
 
+    public Mono<String> keyId() {
+        return authentication().map(McpKeyAuthenticationToken::keyId);
+    }
+
     Mono<McpKeyAuthenticationToken> authentication() {
         return ReactiveSecurityContextHolder.getContext()
                 .map(SecurityContext::getAuthentication)

src/main/java/run/halo/mcpserver/McpInFlightLimiter.java

@@ -0,0 +1,23 @@
+package run.halo.mcpserver;
+
+import org.springframework.stereotype.Component;
+
+/** Process-local bounds on concurrent authenticated MCP request work. */
+@Component
+class McpInFlightLimiter {
+
+    static final int GLOBAL_LIMIT = 100;
+    static final int PER_KEY_LIMIT = 16;
+    static final int MAX_TRACKED_KEYS = 10_000;
+
+    private final KeyedInFlightLimiter limiter =
+            new KeyedInFlightLimiter(GLOBAL_LIMIT, PER_KEY_LIMIT, MAX_TRACKED_KEYS);
+
+    /**
+     * Reserves one global and per-key slot for the key, or returns null when either budget is
+     * exhausted. The caller must close the permit exactly once when the request terminates.
+     */
+    KeyedInFlightLimiter.Lease tryAcquire(String keyId) {
+        return limiter.tryAcquire(keyId, 1);
+    }
+}

src/main/java/run/halo/mcpserver/McpIpAllowlist.java

@@ -2,9 +2,9 @@
 
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
-import java.net.UnknownHostException;
 import java.util.LinkedHashSet;
 import java.util.Set;
+import com.google.common.net.InetAddresses;
 import org.springframework.security.util.matcher.InetAddressMatcher;
 import org.springframework.security.util.matcher.InetAddressMatchers;
 import org.springframework.util.StringUtils;
@@ -37,13 +37,12 @@ static boolean allows(Set<String> ranges, InetSocketAddress remoteAddress) {
         if (ranges == null || ranges.isEmpty()) {
             return true;
         }
-        if (remoteAddress == null) {
+        var resolved = resolveNumericAddress(remoteAddress);
+        if (resolved.isEmpty()) {
             return false;
         }
         try {
-            var address = remoteAddress.getAddress() == null
-                    ? parseNumericAddress(remoteAddress.getHostString())
-                    : remoteAddress.getAddress();
+            var address = resolved.get();
             var matchers = ranges.stream()
                     .map(McpIpAllowlist::compile)
                     .toList();
@@ -58,6 +57,23 @@ static boolean allows(Set<String> ranges, InetSocketAddress remoteAddress) {
         return false;
     }
 
+    /**
+     * Resolves the numeric client address for both IP authorization and rate limiting, whether
+     * the socket address carried a resolved InetAddress or only a numeric host string.
+     */
+    static java.util.Optional<InetAddress> resolveNumericAddress(InetSocketAddress remoteAddress) {
+        if (remoteAddress == null) {
+            return java.util.Optional.empty();
+        }
+        try {
+            return java.util.Optional.of(remoteAddress.getAddress() == null
+                    ? parseNumericAddress(remoteAddress.getHostString())
+                    : remoteAddress.getAddress());
+        } catch (IllegalArgumentException error) {
+            return java.util.Optional.empty();
+        }
+    }
+
     private static CompiledRange compile(String range) {
         var matcher = InetAddressMatchers.fromIpAddress(range);
         var slashIndex = range.indexOf('/');
@@ -69,16 +85,15 @@ private static CompiledRange compile(String range) {
         return new CompiledRange(addressLength, matcher);
     }
 
+    /**
+     * Parses a strict numeric IP literal without ever consulting DNS; anything else is rejected
+     * with an IllegalArgumentException.
+     */
     private static InetAddress parseNumericAddress(String address) {
         var value = address.startsWith("[") && address.endsWith("]")
                 ? address.substring(1, address.length() - 1)
                 : address;
-        InetAddressMatchers.fromIpAddress(value);
-        try {
-            return InetAddress.getByName(value);
-        } catch (UnknownHostException error) {
-            throw new IllegalArgumentException(error);
-        }
+        return InetAddresses.forString(value);
     }
 
     private static void validateMask(String mask, int maxBits) {

src/main/java/run/halo/mcpserver/McpKeyAuthenticationFilter.java

@@ -4,6 +4,7 @@
 import static org.springframework.http.HttpHeaders.RETRY_AFTER;
 import static org.springframework.http.HttpHeaders.WWW_AUTHENTICATE;
 
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.core.Ordered;
 import org.springframework.core.annotation.Order;
 import org.springframework.http.HttpStatus;
@@ -20,21 +21,37 @@
 class McpKeyAuthenticationFilter implements BeforeSecurityWebFilter {
 
     static final String MCP_PATH = "/mcp";
+    static final java.time.Duration DEFAULT_REQUEST_TIMEOUT = java.time.Duration.ofSeconds(30);
     private static final String BEARER_SCHEME = "Bearer ";
 
     private final McpAccessKeyService accessKeyService;
     private final McpRequestRateLimiter rateLimiter;
+    private final McpInFlightLimiter inFlightLimiter;
     private final WebHandler mcpHandler;
     private final java.util.Set<String> protocolVersions;
+    private final java.time.Duration requestTimeout;
 
+    @Autowired
     McpKeyAuthenticationFilter(
             McpAccessKeyService accessKeyService,
             McpRequestRateLimiter rateLimiter,
+            McpInFlightLimiter inFlightLimiter,
             HaloMcpServer mcpServer) {
+        this(accessKeyService, rateLimiter, inFlightLimiter, mcpServer, DEFAULT_REQUEST_TIMEOUT);
+    }
+
+    McpKeyAuthenticationFilter(
+            McpAccessKeyService accessKeyService,
+            McpRequestRateLimiter rateLimiter,
+            McpInFlightLimiter inFlightLimiter,
+            HaloMcpServer mcpServer,
+            java.time.Duration requestTimeout) {
         this.accessKeyService = accessKeyService;
         this.rateLimiter = rateLimiter;
+        this.inFlightLimiter = inFlightLimiter;
         this.mcpHandler = RouterFunctions.toWebHandler(mcpServer.routerFunction());
         this.protocolVersions = java.util.Set.copyOf(mcpServer.protocolVersions());
+        this.requestTimeout = requestTimeout;
     }
 
     @Override
@@ -56,15 +73,27 @@ public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
                     if (!hasSupportedProtocolVersion(exchange)) {
                         return badRequest(exchange).thenReturn(true);
                     }
+                    var permit = inFlightLimiter.tryAcquire(authentication.keyId());
+                    if (permit == null) {
+                        return tooManyRequests(exchange).thenReturn(true);
+                    }
                     var request = exchange.getRequest().mutate()
                             .headers(headers -> headers.remove(AUTHORIZATION))
                             .build();
-                    return mcpHandler.handle(exchange.mutate().request(request).build())
+                    // Defer so a synchronously throwing handler assembly still flows through
+                    // the error path and releases the permit.
+                    return Mono.defer(() -> mcpHandler.handle(exchange.mutate().request(request).build()))
+                            .doFinally(ignored -> permit.close())
                             .contextWrite(org.springframework.security.core.context.ReactiveSecurityContextHolder
                                     .withAuthentication(authentication))
                             .thenReturn(true);
                 })
                 .defaultIfEmpty(false)
+                // The deadline covers the whole authenticate-and-handle chain so stalled
+                // credential lookups cannot park requests either.
+                .timeout(requestTimeout)
+                .onErrorResume(java.util.concurrent.TimeoutException.class,
+                        error -> serviceUnavailable(exchange).thenReturn(true))
                 .flatMap(handled -> handled ? Mono.empty() : unauthorized(exchange));
     }
 
@@ -104,4 +133,12 @@ private static Mono<Void> tooManyRequests(ServerWebExchange exchange) {
                 RETRY_AFTER, String.valueOf(McpRequestRateLimiter.RETRY_AFTER_SECONDS));
         return exchange.getResponse().setComplete();
     }
+
+    private static Mono<Void> serviceUnavailable(ServerWebExchange exchange) {
+        if (exchange.getResponse().isCommitted()) {
+            return Mono.empty();
+        }
+        exchange.getResponse().setStatusCode(HttpStatus.SERVICE_UNAVAILABLE);
+        return exchange.getResponse().setComplete();
+    }
 }

src/main/java/run/halo/mcpserver/McpRequestRateLimiter.java

@@ -31,9 +31,9 @@ class McpRequestRateLimiter {
     }
 
     boolean allowRequest(InetSocketAddress remoteAddress) {
-        var source = remoteAddress == null || remoteAddress.getAddress() == null
-                ? "unknown"
-                : remoteAddress.getAddress().getHostAddress();
+        var source = McpIpAllowlist.resolveNumericAddress(remoteAddress)
+                .map(java.net.InetAddress::getHostAddress)
+                .orElse("unknown");
         return allow(currentWindow().requestCounts, source, REQUESTS_PER_MINUTE);
     }
 

src/main/java/run/halo/mcpserver/McpToolRegistry.java

@@ -55,7 +55,7 @@ Mono<Optional<McpSchema.CallToolResult>> executeIfContributed(
                         .findFirst())
                 .flatMap(tool -> tool
                         .map(value -> gateway(() -> authorization.require(name)
-                                        .then(execute(value.definition(), arguments)))
+                                        .then(Mono.defer(() -> execute(value.definition(), arguments))))
                                 .map(Optional::of))
                         .orElseGet(() -> Mono.just(Optional.empty())));
     }