fix(kb-open): assert session belongs to path kbId + cleanup

- requireSessionOwnership now also checks session.kbId() == path kbId (404 on
  mismatch), so a research session started under one KB cannot be addressed via
  another KB path even when the caller's key is bound to both — defense-in-depth
  on top of the keyId ownership check.
- Drop internal "R7" / "review #446" markers from the controller Javadoc in
  favour of functional wording.
- Import Set/Map/concurrent types and static any() instead of inline FQNs in the
  new kb-open research/auth tests, per code style.
This commit is contained in:
matevip 2026-07-02 17:51:25 +08:00
parent 20c681a7c8
commit 2e3dd071a5
3 changed files with 23 additions and 13 deletions

View File

@ -36,11 +36,11 @@ import java.util.concurrent.Executors;
* pipeline) with SSE progress. The start endpoint returns a sessionId; the
* caller subscribes to SSE for progress, or polls status for the final result.
*
* <p>R7: the SSE endpoint uses {@code ?token=} query param because browser
* <p>The SSE endpoint uses a {@code ?token=} query param because browser
* EventSource cannot set Authorization headers. The {@code KbOpenApiAuthFilter}
* already falls back to query param tokens.
*
* <h3>Cost &amp; lifecycle controls (review #446)</h3>
* <h3>Cost &amp; lifecycle controls</h3>
* <ul>
* <li>Cancel is <b>cooperative</b>: it calls {@link ChatStreamTracker#requestStop}
* so {@link WikiResearchService} bails at the next stage boundary the
@ -128,7 +128,7 @@ public class KbOpenResearchController {
@PathVariable Long kbId,
@PathVariable String sessionId,
HttpServletRequest request) {
requireSessionOwnership(request, sessionId);
requireSessionOwnership(request, kbId, sessionId);
SseEmitter emitter = new Utf8SseEmitter(10 * 60 * 1000L);
boolean attached = streamTracker.attach(sessionId, emitter);
@ -154,7 +154,7 @@ public class KbOpenResearchController {
@PathVariable Long kbId,
@PathVariable String sessionId,
HttpServletRequest request) {
Session session = requireSessionOwnership(request, sessionId);
Session session = requireSessionOwnership(request, kbId, sessionId);
Status status = session.status();
ResearchResult result = session.result();
@ -181,7 +181,7 @@ public class KbOpenResearchController {
@PathVariable Long kbId,
@PathVariable String sessionId,
HttpServletRequest request) {
Session session = requireSessionOwnership(request, sessionId);
Session session = requireSessionOwnership(request, kbId, sessionId);
if (session.status() != Status.RUNNING) {
throw new MateClawException(409, "Session is not running (status: " + session.status() + ")");
}
@ -209,7 +209,7 @@ public class KbOpenResearchController {
return ctx;
}
private Session requireSessionOwnership(HttpServletRequest request, String sessionId) {
private Session requireSessionOwnership(HttpServletRequest request, Long kbId, String sessionId) {
KbApiKeyContext ctx = requireContext(request);
Optional<Session> session = sessionRegistry.get(sessionId);
if (session.isEmpty()) {
@ -219,6 +219,12 @@ public class KbOpenResearchController {
if (!session.get().keyId().equals(ctx.keyId())) {
throw new MateClawException(403, "Session does not belong to this API key");
}
// The session must also belong to the KB named in the path, so a session
// started under one KB cannot be addressed via another even when the
// caller's key happens to be bound to both.
if (!session.get().kbId().equals(kbId)) {
throw new MateClawException(404, "Research session not found: " + sessionId);
}
return session.get();
}
}

View File

@ -13,6 +13,7 @@ import java.util.Optional;
import java.util.Set;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
@ -48,7 +49,7 @@ class KbOpenApiAuthFilterTest {
rateLimiter = mock(KbApiKeyRateLimiter.class);
filter = new KbOpenApiAuthFilter(keyService, rateLimiter);
when(keyService.authenticate(KEY)).thenReturn(Optional.of(new AuthResult(CTX, LocalDateTime.now())));
when(rateLimiter.tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any())).thenReturn(true);
when(rateLimiter.tryAcquire(anyLong(), anyInt(), any())).thenReturn(true);
}
private MockHttpServletRequest startRequest() {
@ -113,7 +114,7 @@ class KbOpenApiAuthFilterTest {
void sseSkipsRateLimit() throws Exception {
run(sseRequest());
verify(rateLimiter, never())
.tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any());
.tryAcquire(anyLong(), anyInt(), any());
}
@Test
@ -121,6 +122,6 @@ class KbOpenApiAuthFilterTest {
void nonSseHitsRateLimit() throws Exception {
run(startRequest());
verify(rateLimiter)
.tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any());
.tryAcquire(anyLong(), anyInt(), any());
}
}

View File

@ -9,8 +9,11 @@ import vip.mate.wiki.service.WikiResearchService.ResearchResult;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@ -187,10 +190,10 @@ class KbResearchSessionRegistryTest {
int cap = 3;
KbResearchSessionRegistry registry = new KbResearchSessionRegistry(cap, Duration.ofMinutes(30));
int threads = cap * 4; // far more contenders than slots
java.util.concurrent.CountDownLatch start = new java.util.concurrent.CountDownLatch(1);
java.util.concurrent.atomic.AtomicInteger admitted = new java.util.concurrent.atomic.AtomicInteger();
java.util.concurrent.atomic.AtomicInteger rejected = new java.util.concurrent.atomic.AtomicInteger();
java.util.List<Thread> workers = new java.util.ArrayList<>();
CountDownLatch start = new CountDownLatch(1);
AtomicInteger admitted = new AtomicInteger();
AtomicInteger rejected = new AtomicInteger();
List<Thread> workers = new ArrayList<>();
for (int i = 0; i < threads; i++) {
String sid = "concurrent-" + i;