diff --git a/README.md b/README.md index 153bfbf..47b85d9 100644 --- a/README.md +++ b/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. diff --git a/src/main/java/run/halo/mcpserver/HaloMcpServer.java b/src/main/java/run/halo/mcpserver/HaloMcpServer.java index af276c6..62cdc2d 100644 --- a/src/main/java/run/halo/mcpserver/HaloMcpServer.java +++ b/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(); } diff --git a/src/main/java/run/halo/mcpserver/KeyedInFlightLimiter.java b/src/main/java/run/halo/mcpserver/KeyedInFlightLimiter.java new file mode 100644 index 0000000..68ef6f0 --- /dev/null +++ b/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 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(); + 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; + }); + } + } + } +} diff --git a/src/main/java/run/halo/mcpserver/McpAccessKeyService.java b/src/main/java/run/halo/mcpserver/McpAccessKeyService.java index 9d90c59..cd87446 100644 --- a/src/main/java/run/halo/mcpserver/McpAccessKeyService.java +++ b/src/main/java/run/halo/mcpserver/McpAccessKeyService.java @@ -127,14 +127,28 @@ Mono 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 revalidate(McpAccessKey snapshot) { + return client.fetch(McpAccessKey.class, snapshot.getMetadata().getName()) + .filter(this::active) + .filter(fresh -> fresh.getSpec().equals(snapshot.getSpec())); } private Mono get(String id) { diff --git a/src/main/java/run/halo/mcpserver/McpAuthorization.java b/src/main/java/run/halo/mcpserver/McpAuthorization.java index bf4ca23..e6e374b 100644 --- a/src/main/java/run/halo/mcpserver/McpAuthorization.java +++ b/src/main/java/run/halo/mcpserver/McpAuthorization.java @@ -32,6 +32,10 @@ public Mono authorize(String toolName, Supplier> action) { })); } + public Mono keyId() { + return authentication().map(McpKeyAuthenticationToken::keyId); + } + Mono authentication() { return ReactiveSecurityContextHolder.getContext() .map(SecurityContext::getAuthentication) diff --git a/src/main/java/run/halo/mcpserver/McpInFlightLimiter.java b/src/main/java/run/halo/mcpserver/McpInFlightLimiter.java new file mode 100644 index 0000000..c544075 --- /dev/null +++ b/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); + } +} diff --git a/src/main/java/run/halo/mcpserver/McpIpAllowlist.java b/src/main/java/run/halo/mcpserver/McpIpAllowlist.java index 9f43198..1f9c5d4 100644 --- a/src/main/java/run/halo/mcpserver/McpIpAllowlist.java +++ b/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 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 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 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) { diff --git a/src/main/java/run/halo/mcpserver/McpKeyAuthenticationFilter.java b/src/main/java/run/halo/mcpserver/McpKeyAuthenticationFilter.java index 7803e18..e68f5dc 100644 --- a/src/main/java/run/halo/mcpserver/McpKeyAuthenticationFilter.java +++ b/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 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 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 tooManyRequests(ServerWebExchange exchange) { RETRY_AFTER, String.valueOf(McpRequestRateLimiter.RETRY_AFTER_SECONDS)); return exchange.getResponse().setComplete(); } + + private static Mono serviceUnavailable(ServerWebExchange exchange) { + if (exchange.getResponse().isCommitted()) { + return Mono.empty(); + } + exchange.getResponse().setStatusCode(HttpStatus.SERVICE_UNAVAILABLE); + return exchange.getResponse().setComplete(); + } } diff --git a/src/main/java/run/halo/mcpserver/McpRequestRateLimiter.java b/src/main/java/run/halo/mcpserver/McpRequestRateLimiter.java index 225bc95..c388ced 100644 --- a/src/main/java/run/halo/mcpserver/McpRequestRateLimiter.java +++ b/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); } diff --git a/src/main/java/run/halo/mcpserver/McpToolRegistry.java b/src/main/java/run/halo/mcpserver/McpToolRegistry.java index e6620ae..f35629d 100644 --- a/src/main/java/run/halo/mcpserver/McpToolRegistry.java +++ b/src/main/java/run/halo/mcpserver/McpToolRegistry.java @@ -55,7 +55,7 @@ Mono> 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()))); } diff --git a/src/main/java/run/halo/mcpserver/tools/AttachmentTools.java b/src/main/java/run/halo/mcpserver/tools/AttachmentTools.java index eca71cd..ea6d6aa 100644 --- a/src/main/java/run/halo/mcpserver/tools/AttachmentTools.java +++ b/src/main/java/run/halo/mcpserver/tools/AttachmentTools.java @@ -33,14 +33,17 @@ class AttachmentTools extends ToolSupport implements ToolGroup { private final ReactiveExtensionClient client; private final AttachmentService attachmentService; + private final AttachmentUploadLimiter uploadLimiter; AttachmentTools( ReactiveExtensionClient client, AttachmentService attachmentService, + AttachmentUploadLimiter uploadLimiter, McpAuthorization authorization) { super(authorization); this.client = client; this.attachmentService = attachmentService; + this.uploadLimiter = uploadLimiter; } @Override @@ -74,6 +77,29 @@ Mono upload(Map arguments) { if (encoded.length() > 12_000_000) { throw new McpToolException("INVALID_ARGUMENT", "contentBase64 exceeds the 8 MiB limit"); } + var mediaType = mediaType(arguments.get("mediaType")); + return authorization.keyId().flatMap(keyId -> { + var reservation = uploadLimiter.tryAcquire(keyId, decodedLength(encoded)); + if (reservation == null) { + return Mono.error(new McpToolException( + "RATE_LIMITED", "Too many concurrent attachment uploads; please retry shortly")); + } + return Mono.fromCallable(() -> decode(encoded)) + .flatMap(bytes -> { + var buffer = DefaultDataBufferFactory.sharedInstance.wrap(bytes); + return attachmentService + .upload(policyName, groupName, filename, Flux.just(buffer), mediaType) + .switchIfEmpty(Mono.error(new McpToolException( + "ATTACHMENT_UNAVAILABLE", "Halo did not create the attachment"))) + .map(attachment -> payload( + ContentPayloads.attachment(attachment), + "Uploaded attachment " + filename)); + }) + .doFinally(ignored -> reservation.close()); + }); + } + + private static byte[] decode(String encoded) { final byte[] bytes; try { bytes = Base64.getDecoder().decode(encoded); @@ -84,17 +110,16 @@ Mono upload(Map arguments) { throw new McpToolException( "INVALID_ARGUMENT", "Attachment content must be between 1 byte and 8 MiB"); } - var buffer = DefaultDataBufferFactory.sharedInstance.wrap(bytes); - return attachmentService - .upload( - policyName, - groupName, - filename, - Flux.just(buffer), - mediaType(arguments.get("mediaType"))) - .switchIfEmpty(Mono.error(new McpToolException( - "ATTACHMENT_UNAVAILABLE", "Halo did not create the attachment"))) - .map(attachment -> payload(ContentPayloads.attachment(attachment), "Uploaded attachment " + filename)); + return bytes; + } + + /** + * Returns the exact length valid Base64 decodes to. Invalid input may misestimate but never + * below zero; the content is still validated when it is actually decoded. + */ + private static int decodedLength(String encoded) { + var padding = encoded.endsWith("==") ? 2 : encoded.endsWith("=") ? 1 : 0; + return Math.max(encoded.length() * 3 / 4 - padding, 0); } Mono delete(Map arguments) { diff --git a/src/main/java/run/halo/mcpserver/tools/AttachmentUploadLimiter.java b/src/main/java/run/halo/mcpserver/tools/AttachmentUploadLimiter.java new file mode 100644 index 0000000..014da31 --- /dev/null +++ b/src/main/java/run/halo/mcpserver/tools/AttachmentUploadLimiter.java @@ -0,0 +1,24 @@ +package run.halo.mcpserver.tools; + +import org.springframework.stereotype.Component; +import run.halo.mcpserver.KeyedInFlightLimiter; + +/** Process-local in-flight byte budgets that bound concurrent attachment uploads. */ +@Component +class AttachmentUploadLimiter { + + static final int GLOBAL_BUDGET_BYTES = 64 * 1024 * 1024; + static final int PER_KEY_BUDGET_BYTES = 32 * 1024 * 1024; + static final int MAX_TRACKED_KEYS = 10_000; + + private final KeyedInFlightLimiter limiter = + new KeyedInFlightLimiter(GLOBAL_BUDGET_BYTES, PER_KEY_BUDGET_BYTES, MAX_TRACKED_KEYS); + + /** + * Reserves decoded bytes against the global and per-key budgets, or returns null when either + * is exhausted. The caller must close the reservation exactly once when the upload terminates. + */ + KeyedInFlightLimiter.Lease tryAcquire(String keyId, int bytes) { + return limiter.tryAcquire(keyId, bytes); + } +} diff --git a/src/test/java/run/halo/mcpserver/KeyedInFlightLimiterTest.java b/src/test/java/run/halo/mcpserver/KeyedInFlightLimiterTest.java new file mode 100644 index 0000000..4a5963d --- /dev/null +++ b/src/test/java/run/halo/mcpserver/KeyedInFlightLimiterTest.java @@ -0,0 +1,94 @@ +package run.halo.mcpserver; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.ArrayList; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +class KeyedInFlightLimiterTest { + + private static final int GLOBAL_LIMIT = 8; + private static final int PER_KEY_LIMIT = 2; + private static final int MAX_TRACKED_KEYS = 4; + + @Test + void enforcesAndReleasesBudgets() { + var limiter = new KeyedInFlightLimiter(GLOBAL_LIMIT, PER_KEY_LIMIT, MAX_TRACKED_KEYS); + var first = limiter.tryAcquire("key-one", 1); + var second = limiter.tryAcquire("key-one", 1); + + assertThat(first).isNotNull(); + assertThat(second).isNotNull(); + assertThat(limiter.tryAcquire("key-one", 1)).isNull(); + first.close(); + first.close(); + assertThat(limiter.tryAcquire("key-one", 1)).isNotNull(); + assertThat(limiter.tryAcquire("key-one", 1)).isNull(); + + var globalLimiter = new KeyedInFlightLimiter(GLOBAL_LIMIT, GLOBAL_LIMIT, GLOBAL_LIMIT + 1); + var global = new ArrayList(); + for (var i = 0; i < GLOBAL_LIMIT; i++) { + global.add(globalLimiter.tryAcquire("global-" + i, 1)); + } + assertThat(global).doesNotContainNull(); + assertThat(globalLimiter.tryAcquire("global-extra", 1)).isNull(); + } + + @Test + void forgetsReleasedKeys() { + var limiter = new KeyedInFlightLimiter(GLOBAL_LIMIT, PER_KEY_LIMIT, MAX_TRACKED_KEYS); + + for (var i = 0; i < MAX_TRACKED_KEYS + 1; i++) { + var lease = limiter.tryAcquire("key-" + i, 1); + assertThat(lease).isNotNull(); + lease.close(); + } + } + + @Test + void keepsPerKeyAccountingConsistentUnderConcurrentChurn() throws InterruptedException { + var limiter = new KeyedInFlightLimiter(GLOBAL_LIMIT, PER_KEY_LIMIT, MAX_TRACKED_KEYS); + var inFlight = new ConcurrentHashMap(); + var maximum = new ConcurrentHashMap(); + var errors = new CopyOnWriteArrayList(); + var threads = new ArrayList(); + for (var t = 0; t < 8; t++) { + var thread = new Thread(() -> { + try { + for (var i = 0; i < 2_000; i++) { + var key = "key-" + i % 4; + var lease = limiter.tryAcquire(key, 1); + if (lease == null) { + continue; + } + var count = inFlight.computeIfAbsent(key, ignored -> new AtomicInteger()); + var current = count.incrementAndGet(); + maximum.computeIfAbsent(key, ignored -> new AtomicInteger()) + .accumulateAndGet(current, Math::max); + count.decrementAndGet(); + lease.close(); + } + } catch (Throwable error) { + errors.add(error); + } + }); + threads.add(thread); + thread.start(); + } + for (var thread : threads) { + thread.join(); + } + + assertThat(errors).isEmpty(); + assertThat(maximum.values()) + .allSatisfy(max -> assertThat(max.get()).isLessThanOrEqualTo(PER_KEY_LIMIT)); + for (var i = 0; i < 4; i++) { + var lease = limiter.tryAcquire("key-" + i, PER_KEY_LIMIT); + assertThat(lease).isNotNull(); + lease.close(); + } + } +} diff --git a/src/test/java/run/halo/mcpserver/McpAccessKeyServiceTest.java b/src/test/java/run/halo/mcpserver/McpAccessKeyServiceTest.java index db2dd17..fc2a749 100644 --- a/src/test/java/run/halo/mcpserver/McpAccessKeyServiceTest.java +++ b/src/test/java/run/halo/mcpserver/McpAccessKeyServiceTest.java @@ -130,4 +130,93 @@ void authenticatesOnlyFromAnAllowedIpRange() { .isNotNull(); assertThat(created.accessKey().getStatus().getLastUsedAt()).isNotNull(); } + + @Test + void authenticatesWhenTheKeyIsUnchangedDuringVerification() { + var created = service.create( + "Automation", "admin", Set.of("halo_search_content"), Set.of(), null) + .block(); + var id = created.accessKey().getMetadata().getName(); + when(client.fetch(McpAccessKey.class, id)) + .thenReturn(Mono.just(created.accessKey())) + .thenReturn(Mono.just(copyOf(created.accessKey()))); + + assertThat(service.authenticate(created.token(), null).block()).isNotNull(); + } + + @Test + void rejectsAuthenticationWhenRotationCommitsDuringVerification() { + var created = service.create( + "Automation", "admin", Set.of("halo_search_content"), Set.of(), null) + .block(); + var id = created.accessKey().getMetadata().getName(); + var rotated = copyOf(created.accessKey()); + rotated.getSpec().setKeyHash("rotated-hash"); + when(client.fetch(McpAccessKey.class, id)) + .thenReturn(Mono.just(created.accessKey())) + .thenReturn(Mono.just(rotated)); + + assertThat(service.authenticate(created.token(), null).block()).isNull(); + } + + @Test + void rejectsAuthenticationWhenDisablementOrScopeChangeCommitsDuringVerification() { + var created = service.create( + "Automation", "admin", Set.of("halo_search_content"), Set.of(), null) + .block(); + var id = created.accessKey().getMetadata().getName(); + + var disabled = copyOf(created.accessKey()); + disabled.getSpec().setEnabled(false); + when(client.fetch(McpAccessKey.class, id)) + .thenReturn(Mono.just(created.accessKey())) + .thenReturn(Mono.just(disabled)); + assertThat(service.authenticate(created.token(), null).block()).isNull(); + + var narrowed = copyOf(created.accessKey()); + narrowed.getSpec().setAllowedTools(Set.of()); + when(client.fetch(McpAccessKey.class, id)) + .thenReturn(Mono.just(created.accessKey())) + .thenReturn(Mono.just(narrowed)); + assertThat(service.authenticate(created.token(), null).block()).isNull(); + } + + @Test + void rejectsAuthenticationWhenDeletionCommitsDuringVerification() { + var created = service.create("Automation", "admin", Set.of(), Set.of(), null).block(); + var id = created.accessKey().getMetadata().getName(); + + when(client.fetch(McpAccessKey.class, id)) + .thenReturn(Mono.just(created.accessKey())) + .thenReturn(Mono.empty()); + assertThat(service.authenticate(created.token(), null).block()).isNull(); + + var deleting = copyOf(created.accessKey()); + deleting.getMetadata().setDeletionTimestamp(Instant.now()); + when(client.fetch(McpAccessKey.class, id)) + .thenReturn(Mono.just(created.accessKey())) + .thenReturn(Mono.just(deleting)); + assertThat(service.authenticate(created.token(), null).block()).isNull(); + } + + private static McpAccessKey copyOf(McpAccessKey accessKey) { + var copy = new McpAccessKey(); + var metadata = new run.halo.app.extension.Metadata(); + metadata.setName(accessKey.getMetadata().getName()); + metadata.setVersion(accessKey.getMetadata().getVersion()); + copy.setMetadata(metadata); + var source = accessKey.getSpec(); + var spec = new McpAccessKey.Spec(); + spec.setDisplayName(source.getDisplayName()); + spec.setKeyHash(source.getKeyHash()); + spec.setKeyPrefix(source.getKeyPrefix()); + spec.setOwnerName(source.getOwnerName()); + spec.setEnabled(source.isEnabled()); + spec.setExpiresAt(source.getExpiresAt()); + spec.setAllowedTools(new java.util.LinkedHashSet<>(source.getAllowedTools())); + spec.setAllowedIpRanges(new java.util.LinkedHashSet<>(source.getAllowedIpRanges())); + copy.setSpec(spec); + copy.setStatus(accessKey.getStatus()); + return copy; + } } diff --git a/src/test/java/run/halo/mcpserver/McpInFlightLimiterTest.java b/src/test/java/run/halo/mcpserver/McpInFlightLimiterTest.java new file mode 100644 index 0000000..e225bc4 --- /dev/null +++ b/src/test/java/run/halo/mcpserver/McpInFlightLimiterTest.java @@ -0,0 +1,29 @@ +package run.halo.mcpserver; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.ArrayList; +import org.junit.jupiter.api.Test; + +class McpInFlightLimiterTest { + + @Test + void appliesRequestLimits() { + var limiter = new McpInFlightLimiter(); + var permits = new ArrayList(); + + for (var i = 0; i < McpInFlightLimiter.PER_KEY_LIMIT; i++) { + permits.add(limiter.tryAcquire("key-one")); + } + assertThat(limiter.tryAcquire("key-one")).isNull(); + var otherKey = limiter.tryAcquire("key-two"); + assertThat(otherKey).isNotNull(); + + permits.forEach(KeyedInFlightLimiter.Lease::close); + otherKey.close(); + for (var i = 0; i < McpInFlightLimiter.GLOBAL_LIMIT; i++) { + assertThat(limiter.tryAcquire("key-" + i)).isNotNull(); + } + assertThat(limiter.tryAcquire("key-extra")).isNull(); + } +} diff --git a/src/test/java/run/halo/mcpserver/McpIpAllowlistTest.java b/src/test/java/run/halo/mcpserver/McpIpAllowlistTest.java index 78d9006..1f73dea 100644 --- a/src/test/java/run/halo/mcpserver/McpIpAllowlistTest.java +++ b/src/test/java/run/halo/mcpserver/McpIpAllowlistTest.java @@ -102,6 +102,14 @@ void allowsUnknownAddressesWhenNotConfigured() { assertThat(McpIpAllowlist.allows(Set.of(), null)).isTrue(); } + @Test + void rejectsHostnamesDisguisedWithHexPrefixesAndColons() { + var disguised = InetSocketAddress.createUnresolved("1:a.attacker.example", 443); + assertThat(McpIpAllowlist.resolveNumericAddress(disguised)).isEmpty(); + assertThat(McpIpAllowlist.allows(Set.of("0.0.0.0/0"), disguised)).isFalse(); + assertThat(McpIpAllowlist.allows(Set.of(), disguised)).isTrue(); + } + @Test void failsClosedWhenStoredConfigurationIsInvalid() { var ranges = new LinkedHashSet<>(java.util.List.of("203.0.113.10", "invalid")); diff --git a/src/test/java/run/halo/mcpserver/McpKeyAuthenticationFilterTest.java b/src/test/java/run/halo/mcpserver/McpKeyAuthenticationFilterTest.java index 321c618..8005d93 100644 --- a/src/test/java/run/halo/mcpserver/McpKeyAuthenticationFilterTest.java +++ b/src/test/java/run/halo/mcpserver/McpKeyAuthenticationFilterTest.java @@ -1,6 +1,7 @@ package run.halo.mcpserver; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -13,6 +14,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; import org.springframework.http.HttpHeaders; import org.springframework.http.HttpStatus; import org.springframework.mock.http.server.reactive.MockServerHttpRequest; @@ -56,7 +58,23 @@ void setUp() { .build()); when(mcpServer.protocolVersions()).thenReturn(java.util.List.of("2025-11-25")); rateLimiter = new McpRequestRateLimiter(); - filter = new McpKeyAuthenticationFilter(accessKeyService, rateLimiter, mcpServer); + filter = new McpKeyAuthenticationFilter( + accessKeyService, rateLimiter, new McpInFlightLimiter(), mcpServer); + } + + @Test + void springCreatesTheFilterUsingTheProductionConstructor() { + try (var context = new AnnotationConfigApplicationContext()) { + context.registerBean(McpAccessKeyService.class, () -> accessKeyService); + context.registerBean(McpRequestRateLimiter.class, () -> rateLimiter); + context.registerBean(McpInFlightLimiter.class, McpInFlightLimiter::new); + context.registerBean(HaloMcpServer.class, () -> mcpServer); + context.register(McpKeyAuthenticationFilter.class); + + context.refresh(); + + assertThat(context.getBean(McpKeyAuthenticationFilter.class)).isNotNull(); + } } @Test @@ -187,4 +205,98 @@ void rateLimitsRequestsBeforeAccessKeyLookup() { .isEqualTo("60"); verify(accessKeyService, never()).authenticate(rawToken, null); } + + @Test + void timesOutAStalledMcpRequestAndReleasesThePermit() { + org.springframework.web.reactive.function.server.HandlerFunction stalled = + request -> Mono.never(); + when(mcpServer.routerFunction()).thenReturn( + RouterFunctions.route().POST("/mcp", stalled).build()); + var inFlightLimiter = new McpInFlightLimiter(); + var fastFilter = new McpKeyAuthenticationFilter( + accessKeyService, new McpRequestRateLimiter(), inFlightLimiter, mcpServer, + java.time.Duration.ofMillis(50)); + var keyId = "00000000-0000-0000-0000-000000000000"; + var rawToken = "hmcp_" + keyId + "_secret"; + when(accessKeyService.authenticate(rawToken, null)) + .thenReturn(Mono.just(new McpKeyAuthenticationToken( + keyId, "Automation", "hmcp_00000000", "admin", Set.of()))); + var exchange = MockServerWebExchange.from(MockServerHttpRequest.post(McpKeyAuthenticationFilter.MCP_PATH) + .header(HttpHeaders.AUTHORIZATION, "Bearer " + rawToken)); + + fastFilter.filter(exchange, ignored -> Mono.error(new AssertionError("Request must not continue"))) + .block(); + + assertThat(exchange.getResponse().getStatusCode()).isEqualTo(HttpStatus.SERVICE_UNAVAILABLE); + for (var i = 0; i < McpInFlightLimiter.PER_KEY_LIMIT; i++) { + assertThat(inFlightLimiter.tryAcquire(keyId)).isNotNull(); + } + } + + @Test + void timesOutAStalledAuthentication() { + var rawToken = "hmcp_00000000-0000-0000-0000-000000000000_secret"; + when(accessKeyService.authenticate(rawToken, null)).thenReturn(Mono.never()); + var fastFilter = new McpKeyAuthenticationFilter( + accessKeyService, new McpRequestRateLimiter(), new McpInFlightLimiter(), mcpServer, + java.time.Duration.ofMillis(50)); + var exchange = MockServerWebExchange.from(MockServerHttpRequest.post(McpKeyAuthenticationFilter.MCP_PATH) + .header(HttpHeaders.AUTHORIZATION, "Bearer " + rawToken)); + + fastFilter.filter(exchange, ignored -> Mono.error(new AssertionError("Request must not continue"))) + .block(); + + assertThat(exchange.getResponse().getStatusCode()).isEqualTo(HttpStatus.SERVICE_UNAVAILABLE); + } + + @Test + void rejectsRequestsWhenTheInFlightBudgetIsExhausted() { + var inFlightLimiter = new McpInFlightLimiter(); + var keyId = "00000000-0000-0000-0000-000000000000"; + for (var i = 0; i < McpInFlightLimiter.PER_KEY_LIMIT; i++) { + assertThat(inFlightLimiter.tryAcquire(keyId)).isNotNull(); + } + var rawToken = "hmcp_" + keyId + "_secret"; + when(accessKeyService.authenticate(rawToken, null)) + .thenReturn(Mono.just(new McpKeyAuthenticationToken( + keyId, "Automation", "hmcp_00000000", "admin", Set.of()))); + var exchange = MockServerWebExchange.from(MockServerHttpRequest.post(McpKeyAuthenticationFilter.MCP_PATH) + .header(HttpHeaders.AUTHORIZATION, "Bearer " + rawToken)); + + new McpKeyAuthenticationFilter(accessKeyService, rateLimiter, inFlightLimiter, mcpServer) + .filter(exchange, ignored -> Mono.error(new AssertionError("Request must not continue"))) + .block(); + + assertThat(exchange.getResponse().getStatusCode()).isEqualTo(HttpStatus.TOO_MANY_REQUESTS); + assertThat(handledPath.get()).isNull(); + } + + @Test + void releasesThePermitWhenTheHandlerAssemblyFailsSynchronously() { + org.springframework.web.reactive.function.server.RouterFunction broken = + request -> { + throw new IllegalStateException("boom"); + }; + when(mcpServer.routerFunction()).thenReturn(broken); + var inFlightLimiter = new McpInFlightLimiter(); + var brokenFilter = new McpKeyAuthenticationFilter( + accessKeyService, new McpRequestRateLimiter(), inFlightLimiter, mcpServer); + var keyId = "00000000-0000-0000-0000-000000000000"; + var rawToken = "hmcp_" + keyId + "_secret"; + when(accessKeyService.authenticate(rawToken, null)) + .thenReturn(Mono.just(new McpKeyAuthenticationToken( + keyId, "Automation", "hmcp_00000000", "admin", Set.of()))); + var exchange = MockServerWebExchange.from(MockServerHttpRequest.post(McpKeyAuthenticationFilter.MCP_PATH) + .header(HttpHeaders.AUTHORIZATION, "Bearer " + rawToken)); + + assertThatThrownBy(() -> brokenFilter + .filter(exchange, ignored -> Mono.error(new AssertionError("Request must not continue"))) + .block()) + .isInstanceOf(IllegalStateException.class) + .hasMessage("boom"); + + for (var i = 0; i < McpInFlightLimiter.PER_KEY_LIMIT; i++) { + assertThat(inFlightLimiter.tryAcquire(keyId)).isNotNull(); + } + } } diff --git a/src/test/java/run/halo/mcpserver/McpRequestRateLimiterTest.java b/src/test/java/run/halo/mcpserver/McpRequestRateLimiterTest.java index 6cc3340..601982e 100644 --- a/src/test/java/run/halo/mcpserver/McpRequestRateLimiterTest.java +++ b/src/test/java/run/halo/mcpserver/McpRequestRateLimiterTest.java @@ -37,4 +37,52 @@ void isolatesToolLimitsByKeyAndToolAndCanBeCleared() { limiter.clear(); assertThat(limiter.allowTool("key-one", "demo/one")).isTrue(); } + + @Test + void givesForwardedClientsIndependentBuckets() { + var limiter = new McpRequestRateLimiter(() -> 0L); + var first = InetSocketAddress.createUnresolved("203.0.113.10", 443); + var second = InetSocketAddress.createUnresolved("203.0.113.11", 443); + + for (var i = 0; i < McpRequestRateLimiter.REQUESTS_PER_MINUTE; i++) { + assertThat(limiter.allowRequest(first)).isTrue(); + } + assertThat(limiter.allowRequest(first)).isFalse(); + assertThat(limiter.allowRequest(second)).isTrue(); + } + + @Test + void canonicalizesResolvedAndUnresolvedFormsOfOneAddress() { + var limiter = new McpRequestRateLimiter(() -> 0L); + var resolved = new InetSocketAddress("203.0.113.10", 443); + var unresolved = InetSocketAddress.createUnresolved("203.0.113.10", 443); + + for (var i = 0; i < McpRequestRateLimiter.REQUESTS_PER_MINUTE; i++) { + assertThat(limiter.allowRequest(resolved)).isTrue(); + } + assertThat(limiter.allowRequest(unresolved)).isFalse(); + } + + @Test + void constrainsUnresolvableClientsToOneSharedBucket() { + var limiter = new McpRequestRateLimiter(() -> 0L); + + for (var i = 0; i < McpRequestRateLimiter.REQUESTS_PER_MINUTE; i++) { + assertThat(limiter.allowRequest(null)).isTrue(); + } + assertThat(limiter.allowRequest(InetSocketAddress.createUnresolved("not-an-ip", 443))) + .isFalse(); + assertThat(limiter.allowRequest(new InetSocketAddress("203.0.113.10", 443))).isTrue(); + } + + @Test + void bucketsHostnamesAsUnknownWithoutResolvingThem() { + var limiter = new McpRequestRateLimiter(() -> 0L); + var disguised = InetSocketAddress.createUnresolved("1:a.attacker.example", 443); + + for (var i = 0; i < McpRequestRateLimiter.REQUESTS_PER_MINUTE; i++) { + assertThat(limiter.allowRequest(disguised)).isTrue(); + } + assertThat(limiter.allowRequest(null)).isFalse(); + } } diff --git a/src/test/java/run/halo/mcpserver/McpToolRegistryTest.java b/src/test/java/run/halo/mcpserver/McpToolRegistryTest.java index c9c62af..ba7c543 100644 --- a/src/test/java/run/halo/mcpserver/McpToolRegistryTest.java +++ b/src/test/java/run/halo/mcpserver/McpToolRegistryTest.java @@ -6,6 +6,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -170,6 +171,34 @@ void rejectsToolWhenKeyOrProviderPermissionDenies() { .verifyComplete(); } + @Test + void defersTheProviderPermissionCallbackUntilTheKeyAllowlistPasses() { + var permissionChecks = new AtomicInteger(); + var tool = McpToolDefinition.builder() + .name("demo/secret") + .inputSchema(objectSchema(Map.of(), List.of())) + .permission(invocation -> { + permissionChecks.incrementAndGet(); + return Mono.just(true); + }) + .handler(invocation -> Mono.just(McpToolResult.success(Map.of("ok", true)))) + .build(); + providerTools(provider, "demo", tool); + + StepVerifier.create(registry.executeIfContributed("demo/secret", Map.of()) + .contextWrite(context("demo/other"))) + .assertNext(result -> assertThat(result.orElseThrow().structuredContent().toString()) + .contains("FORBIDDEN")) + .verifyComplete(); + assertThat(permissionChecks).hasValue(0); + + StepVerifier.create(registry.executeIfContributed("demo/secret", Map.of()) + .contextWrite(context("demo/secret"))) + .assertNext(result -> assertThat(result.orElseThrow().isError()).isFalse()) + .verifyComplete(); + assertThat(permissionChecks).hasValue(1); + } + @Test void resolvesProviderListForEachRequest() { var current = new java.util.concurrent.atomic.AtomicReference(tool("demo/one")); diff --git a/src/test/java/run/halo/mcpserver/tools/AttachmentToolsTest.java b/src/test/java/run/halo/mcpserver/tools/AttachmentToolsTest.java index ee72f11..18b55a6 100644 --- a/src/test/java/run/halo/mcpserver/tools/AttachmentToolsTest.java +++ b/src/test/java/run/halo/mcpserver/tools/AttachmentToolsTest.java @@ -1,11 +1,13 @@ package run.halo.mcpserver.tools; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.nio.charset.StandardCharsets; +import java.util.Base64; import java.util.Map; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -33,14 +35,100 @@ class AttachmentToolsTest { @Test void validatesAttachmentBase64BeforeUpload() { - var tools = new AttachmentTools(client, attachmentService, authorization); + var tools = new AttachmentTools( + client, attachmentService, new AttachmentUploadLimiter(), authorization); + stubKeyId("key-one"); - assertThatThrownBy(() -> tools.upload(Map.of( + StepVerifier.create(tools.upload(Map.of( "filename", "a.txt", "policyName", "local", "contentBase64", "not-base64"))) - .isInstanceOf(McpToolException.class) - .hasMessageContaining("valid Base64"); + .expectErrorSatisfies(error -> assertThat(error) + .isInstanceOf(McpToolException.class) + .hasMessageContaining("valid Base64")) + .verify(); + verify(attachmentService, never()).upload(any(), any(), any(), any(reactor.core.publisher.Flux.class), any(org.springframework.http.MediaType.class)); + } + + @Test + void rejectsUploadsWhenTheInFlightByteBudgetIsExhausted() { + var uploadLimiter = new AttachmentUploadLimiter(); + var reservation = uploadLimiter.tryAcquire( + "key-one", AttachmentUploadLimiter.PER_KEY_BUDGET_BYTES); + var tools = new AttachmentTools(client, attachmentService, uploadLimiter, authorization); + stubKeyId("key-one"); + var arguments = Map.of( + "filename", "a.txt", + "policyName", "local", + "contentBase64", Base64.getEncoder().encodeToString("hello".getBytes(StandardCharsets.UTF_8))); + + StepVerifier.create(tools.upload(arguments)) + .expectErrorSatisfies(error -> assertThat(error) + .isInstanceOf(McpToolException.class) + .hasMessageContaining("concurrent attachment uploads")) + .verify(); + verify(attachmentService, never()).upload(any(), any(), any(), any(reactor.core.publisher.Flux.class), any(org.springframework.http.MediaType.class)); + + reservation.close(); + var attachment = new Attachment(); + attachment.setMetadata(ToolSupport.metadata("a.txt")); + when(attachmentService.upload(any(), any(), any(), any(reactor.core.publisher.Flux.class), any(org.springframework.http.MediaType.class))) + .thenReturn(Mono.just(attachment)); + StepVerifier.create(tools.upload(arguments)) + .assertNext(payload -> assertThat(payload.summary()).contains("a.txt")) + .verifyComplete(); + } + + @Test + void releasesTheByteBudgetWhenTheUploadIsCancelled() { + var uploadLimiter = new AttachmentUploadLimiter(); + var tools = new AttachmentTools(client, attachmentService, uploadLimiter, authorization); + stubKeyId("key-one"); + when(attachmentService.upload(any(), any(), any(), any(reactor.core.publisher.Flux.class), any(org.springframework.http.MediaType.class))).thenReturn(Mono.never()); + + StepVerifier.create(tools.upload(Map.of( + "filename", "a.txt", + "policyName", "local", + "contentBase64", + Base64.getEncoder().encodeToString("hello".getBytes(StandardCharsets.UTF_8))))) + .thenCancel() + .verify(); + + assertThat(uploadLimiter.tryAcquire("key-one", AttachmentUploadLimiter.PER_KEY_BUDGET_BYTES)) + .isNotNull(); + } + + @Test + void admitsConcurrentUploadsUpToTheExactPerKeyByteBudget() { + var uploadLimiter = new AttachmentUploadLimiter(); + var tools = new AttachmentTools(client, attachmentService, uploadLimiter, authorization); + stubKeyId("key-one"); + when(attachmentService.upload(any(), any(), any(), any(reactor.core.publisher.Flux.class), any(org.springframework.http.MediaType.class))) + .thenReturn(Mono.never()); + var arguments = Map.of( + "filename", "a.bin", + "policyName", "local", + "contentBase64", Base64.getEncoder().encodeToString(new byte[8 * 1024 * 1024])); + + var errors = new java.util.concurrent.CopyOnWriteArrayList(); + var uploads = new java.util.ArrayList(); + for (var i = 0; i < 4; i++) { + uploads.add(tools.upload(arguments).subscribe(payload -> {}, errors::add)); + } + + // Four times 8 MiB is exactly the per-key budget: all four must be admitted. + assertThat(errors).isEmpty(); + assertThat(uploads).allSatisfy(upload -> assertThat(upload.isDisposed()).isFalse()); + + StepVerifier.create(tools.upload(arguments)) + .expectErrorSatisfies(error -> assertThat(error) + .isInstanceOf(McpToolException.class) + .hasMessageContaining("concurrent attachment uploads")) + .verify(); + + uploads.forEach(reactor.core.Disposable::dispose); + assertThat(uploadLimiter.tryAcquire("key-one", AttachmentUploadLimiter.PER_KEY_BUDGET_BYTES)) + .isNotNull(); } @Test @@ -50,7 +138,8 @@ void deletionUsesExtensionLifecycleSoReconcilerCleansStorage() { attachment.getMetadata().setVersion(1L); when(client.fetch(Attachment.class, "attachment-one")).thenReturn(Mono.just(attachment)); when(client.delete(attachment)).thenReturn(Mono.just(attachment)); - var tools = new AttachmentTools(client, attachmentService, authorization); + var tools = new AttachmentTools( + client, attachmentService, new AttachmentUploadLimiter(), authorization); StepVerifier.create(tools.delete(Map.of("name", "attachment-one"))) .assertNext(payload -> assertThat(payload.summary()).contains("attachment-one")) @@ -59,4 +148,8 @@ void deletionUsesExtensionLifecycleSoReconcilerCleansStorage() { verify(client).delete(attachment); verify(attachmentService, never()).delete(attachment); } + + private void stubKeyId(String keyId) { + when(authorization.keyId()).thenReturn(Mono.just(keyId)); + } } diff --git a/src/test/java/run/halo/mcpserver/tools/AttachmentUploadLimiterTest.java b/src/test/java/run/halo/mcpserver/tools/AttachmentUploadLimiterTest.java new file mode 100644 index 0000000..2c75ce1 --- /dev/null +++ b/src/test/java/run/halo/mcpserver/tools/AttachmentUploadLimiterTest.java @@ -0,0 +1,37 @@ +package run.halo.mcpserver.tools; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.ArrayList; +import org.junit.jupiter.api.Test; +import run.halo.mcpserver.KeyedInFlightLimiter; + +class AttachmentUploadLimiterTest { + + private static final int EIGHT_MIB = 8 * 1024 * 1024; + + @Test + void appliesUploadByteBudgets() { + var limiter = new AttachmentUploadLimiter(); + var reservations = new ArrayList(); + + for (var i = 0; i < AttachmentUploadLimiter.PER_KEY_BUDGET_BYTES / EIGHT_MIB; i++) { + reservations.add(limiter.tryAcquire("key-one", EIGHT_MIB)); + } + assertThat(limiter.tryAcquire("key-one", EIGHT_MIB)).isNull(); + assertThat(limiter.tryAcquire("key-two", AttachmentUploadLimiter.PER_KEY_BUDGET_BYTES)) + .isNotNull(); + assertThat(limiter.tryAcquire("key-three", EIGHT_MIB)).isNull(); + + reservations.getFirst().close(); + assertThat(limiter.tryAcquire("key-one", EIGHT_MIB)).isNotNull(); + } + + @Test + void rejectsReservationsLargerThanTheGlobalBudget() { + var limiter = new AttachmentUploadLimiter(); + + assertThat(limiter.tryAcquire("key-one", AttachmentUploadLimiter.GLOBAL_BUDGET_BYTES + 1)) + .isNull(); + } +} diff --git a/src/test/java/run/halo/mcpserver/tools/BuiltInToolsTest.java b/src/test/java/run/halo/mcpserver/tools/BuiltInToolsTest.java index 2e602ab..8e91257 100644 --- a/src/test/java/run/halo/mcpserver/tools/BuiltInToolsTest.java +++ b/src/test/java/run/halo/mcpserver/tools/BuiltInToolsTest.java @@ -27,7 +27,7 @@ void organizesToolsByHaloDomainAndUsesChineseConsoleDescriptions() { new CategoryTools(client, authorization), new TagTools(client, authorization), new CommentTools(client, authorization), - new AttachmentTools(client, mock(AttachmentService.class), authorization)); + new AttachmentTools(client, mock(AttachmentService.class), new AttachmentUploadLimiter(), authorization)); var names = tools.names().toList(); assertThat(tools.tools()).hasSize(29); @@ -70,7 +70,7 @@ void marksOnlyRecycleAndDeleteToolsAsDestructive() { new CategoryTools(client, authorization), new TagTools(client, authorization), new CommentTools(client, authorization), - new AttachmentTools(client, mock(AttachmentService.class), authorization)); + new AttachmentTools(client, mock(AttachmentService.class), new AttachmentUploadLimiter(), authorization)); var destructive = tools.tools().stream() .filter(tool -> Boolean.TRUE.equals(tool.protocolTool().annotations().destructiveHint())) @@ -98,7 +98,7 @@ void publishesDetailedOutputObjectSchemas() { new CategoryTools(client, authorization), new TagTools(client, authorization), new CommentTools(client, authorization), - new AttachmentTools(client, mock(AttachmentService.class), authorization)); + new AttachmentTools(client, mock(AttachmentService.class), new AttachmentUploadLimiter(), authorization)); var schemaValidator = new DefaultJsonSchemaValidator(); assertThat(tools.tools()).allSatisfy(tool -> { diff --git a/ui/src/components/AccessKeySecretModal.vue b/ui/src/components/AccessKeySecretModal.vue index 709dfab..c2f7843 100644 --- a/ui/src/components/AccessKeySecretModal.vue +++ b/ui/src/components/AccessKeySecretModal.vue @@ -26,13 +26,13 @@ const { copy, copied } = useClipboard({
接入方式
- +
diff --git a/ui/src/utils/__tests__/mcp-config.test.ts b/ui/src/utils/__tests__/mcp-config.test.ts index ac6c4e0..4ac4d1f 100644 --- a/ui/src/utils/__tests__/mcp-config.test.ts +++ b/ui/src/utils/__tests__/mcp-config.test.ts @@ -3,10 +3,50 @@ import { describe, expect, it } from 'vitest' import { mcpClientGuides } from '../mcp-config' describe('mcpClientGuides', () => { - it('writes the Codex token directly to a static authorization header', () => { - const codex = mcpClientGuides('hmcp_secret').find((guide) => guide.id === 'codex') + it('embeds no bearer material in any guide content or install URL', () => { + for (const guide of mcpClientGuides()) { + expect(guide.content).not.toContain('hmcp_') + expect(guide.content).not.toMatch(/Bearer (?!\$)/) + const installPayload = decodeURIComponent(guide.installUrl ?? '') + expect(installPayload).not.toContain('hmcp_') + expect(installPayload).not.toMatch(/Bearer (?!\$)/) + } + }) + + it('references the token from the environment for CLI and file-based clients', () => { + const guides = mcpClientGuides() + + const claudeCode = guides.find((guide) => guide.id === 'claude-code') + expect(claudeCode?.content).toContain("--header 'Authorization: Bearer ${HALO_MCP_TOKEN}'") + + const codex = guides.find((guide) => guide.id === 'codex') + expect(codex?.content).toContain('bearer_token_env_var = "HALO_MCP_TOKEN"') + expect(codex?.content).not.toContain('http_headers') + + const cursor = guides.find((guide) => guide.id === 'cursor') + expect(cursor?.content).toContain('Bearer ${env:HALO_MCP_TOKEN}') + }) + + it('keeps only endpoint metadata and an env reference in the Cursor install URL', () => { + const cursor = mcpClientGuides().find((guide) => guide.id === 'cursor') + + const config = JSON.parse(atob(cursor!.installUrl!.split('config=')[1]!)) + expect(config.type).toBe('http') + expect(config.url).toMatch(/\/mcp$/) + expect(config.headers.Authorization).toBe('Bearer ${env:HALO_MCP_TOKEN}') + }) + + it('lets VS Code prompt for the token through secure input storage', () => { + const vscode = mcpClientGuides().find((guide) => guide.id === 'vscode') + + expect(vscode?.content).toContain('${input:halo-mcp-token}') + expect(vscode?.content).toContain('"password": true') - expect(codex?.content).toContain('http_headers = { Authorization = "Bearer hmcp_secret" }') - expect(codex?.content).not.toContain('bearer_token_env_var') + const payload = JSON.parse( + decodeURIComponent(vscode!.installUrl!.slice('vscode:mcp/install?'.length)), + ) + expect(payload.headers.Authorization).toBe('Bearer ${input:halo-mcp-token}') + expect(payload.inputs).toHaveLength(1) + expect(payload.inputs[0]).toMatchObject({ type: 'promptString', password: true }) }) }) diff --git a/ui/src/utils/mcp-config.ts b/ui/src/utils/mcp-config.ts index 5347802..d0a6888 100644 --- a/ui/src/utils/mcp-config.ts +++ b/ui/src/utils/mcp-config.ts @@ -3,25 +3,8 @@ export function mcpEndpoint() { } const SERVER_NAME = 'halo' -const TOKEN_PLACEHOLDER = '$HALO_MCP_TOKEN' - -export function mcpHttpConfig(token = TOKEN_PLACEHOLDER) { - return JSON.stringify( - { - mcpServers: { - [SERVER_NAME]: { - type: 'http', - url: mcpEndpoint(), - headers: { - Authorization: `Bearer ${token}`, - }, - }, - }, - }, - null, - 2, - ) -} +const TOKEN_ENV_VAR = 'HALO_MCP_TOKEN' +const VSCODE_INPUT_ID = 'halo-mcp-token' export type McpClientId = 'claude-code' | 'codex' | 'cursor' | 'vscode' @@ -32,20 +15,39 @@ export interface McpClientGuide { installUrl?: string } -export function mcpClientGuides(token?: string): McpClientGuide[] { +/** + * Guides never receive the plaintext token: CLI and file-based clients read it from the + * HALO_MCP_TOKEN environment variable at runtime, and VS Code prompts for it through its + * secure input storage on first start. + */ +export function mcpClientGuides(): McpClientGuide[] { const endpoint = mcpEndpoint() - const bearer = `Bearer ${token ?? TOKEN_PLACEHOLDER}` - const serverConfig = { + + const cursorServer = { type: 'http', url: endpoint, - headers: { Authorization: bearer }, + headers: { Authorization: `Bearer \${env:${TOKEN_ENV_VAR}}` }, } - const guides: McpClientGuide[] = [ + const vscodeInputs = [ + { + type: 'promptString', + id: VSCODE_INPUT_ID, + description: 'Halo MCP 密钥', + password: true, + }, + ] + const vscodeServer = { + type: 'http', + url: endpoint, + headers: { Authorization: `Bearer \${input:${VSCODE_INPUT_ID}}` }, + } + + return [ { id: 'claude-code', label: 'Claude Code', - content: `claude mcp add --transport http ${SERVER_NAME} ${endpoint} --header "Authorization: ${bearer}"`, + content: `claude mcp add --transport http ${SERVER_NAME} ${endpoint} --header 'Authorization: Bearer \${${TOKEN_ENV_VAR}}'`, }, { id: 'codex', @@ -53,32 +55,25 @@ export function mcpClientGuides(token?: string): McpClientGuide[] { content: `# ~/.codex/config.toml [mcp_servers.${SERVER_NAME}] url = "${endpoint}" -http_headers = { Authorization = "${bearer}" }`, +bearer_token_env_var = "${TOKEN_ENV_VAR}"`, }, { id: 'cursor', label: 'Cursor', - content: mcpHttpConfig(token), + content: JSON.stringify({ mcpServers: { [SERVER_NAME]: cursorServer } }, null, 2), + installUrl: `cursor://anysphere.cursor-deeplink/mcp/install?name=${SERVER_NAME}&config=${window.btoa(JSON.stringify(cursorServer))}`, }, { id: 'vscode', label: 'VS Code', - content: JSON.stringify({ servers: { [SERVER_NAME]: serverConfig } }, null, 2), + content: JSON.stringify( + { inputs: vscodeInputs, servers: { [SERVER_NAME]: vscodeServer } }, + null, + 2, + ), + installUrl: `vscode:mcp/install?${encodeURIComponent( + JSON.stringify({ name: SERVER_NAME, ...vscodeServer, inputs: vscodeInputs }), + )}`, }, ] - - if (token) { - for (const guide of guides) { - if (guide.id === 'cursor') { - const config = window.btoa(JSON.stringify(serverConfig)) - guide.installUrl = `cursor://anysphere.cursor-deeplink/mcp/install?name=${SERVER_NAME}&config=${config}` - } - if (guide.id === 'vscode') { - const config = encodeURIComponent(JSON.stringify({ name: SERVER_NAME, ...serverConfig })) - guide.installUrl = `vscode:mcp/install?${config}` - } - } - } - - return guides }