Add session ownership validation to Redis transport
- Implement setSessionOwner, getSessionOwner, and validation functions - Add isSessionOwnedBy check combining liveness and ownership - Simplify getShttpTransport by removing cleanup return - Move cleanup logic to transport onclose handler 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
75bcc06702
commit
5740a7da15
1 changed files with 24 additions and 33 deletions
|
|
@ -3,7 +3,6 @@ import { redisClient } from "../redis.js";
|
||||||
import { Transport, TransportSendOptions } from "@modelcontextprotocol/sdk/shared/transport.js";
|
import { Transport, TransportSendOptions } from "@modelcontextprotocol/sdk/shared/transport.js";
|
||||||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||||
import { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js";
|
import { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js";
|
||||||
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
|
|
||||||
|
|
||||||
let redisTransportCounter = 0;
|
let redisTransportCounter = 0;
|
||||||
const notificationStreamId = "__GET_stream";
|
const notificationStreamId = "__GET_stream";
|
||||||
|
|
@ -60,6 +59,27 @@ export async function isLive(sessionId: string): Promise<boolean> {
|
||||||
return numSubs > 0;
|
return numSubs > 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export async function setSessionOwner(sessionId: string, userId: string): Promise<void> {
|
||||||
|
await redisClient.set(`session:${sessionId}:owner`, userId);
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function getSessionOwner(sessionId: string): Promise<string | null> {
|
||||||
|
return await redisClient.get(`session:${sessionId}:owner`);
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function validateSessionOwnership(sessionId: string, userId: string): Promise<boolean> {
|
||||||
|
const owner = await getSessionOwner(sessionId);
|
||||||
|
return owner === userId;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function isSessionOwnedBy(sessionId: string, userId: string): Promise<boolean> {
|
||||||
|
const isLiveSession = await isLive(sessionId);
|
||||||
|
if (!isLiveSession) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
return await validateSessionOwnership(sessionId, userId);
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise<void> {
|
export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise<void> {
|
||||||
let redisCleanup: (() => Promise<void>) | undefined = undefined;
|
let redisCleanup: (() => Promise<void>) | undefined = undefined;
|
||||||
|
|
@ -182,45 +202,16 @@ export class ServerRedisTransport implements Transport {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function startServerListeningToRedis(serverWrapper: { server: Server; cleanup: () => void }, sessionId: string): Promise<ServerRedisTransport> {
|
export async function getShttpTransport(sessionId: string): Promise<StreamableHTTPServerTransport> {
|
||||||
const serverRedisTransport = new ServerRedisTransport(sessionId);
|
|
||||||
|
|
||||||
// Add cleanup callback to the transport
|
|
||||||
const originalClose = serverRedisTransport.close.bind(serverRedisTransport);
|
|
||||||
serverRedisTransport.close = async () => {
|
|
||||||
serverWrapper.cleanup();
|
|
||||||
await originalClose();
|
|
||||||
};
|
|
||||||
|
|
||||||
// The server.connect() will call start() on the transport
|
|
||||||
await serverWrapper.server.connect(serverRedisTransport)
|
|
||||||
|
|
||||||
return serverRedisTransport;
|
|
||||||
}
|
|
||||||
|
|
||||||
export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, cleanup: () => Promise<void>}> {
|
|
||||||
const shttpTransport = new StreamableHTTPServerTransport({
|
|
||||||
sessionIdGenerator: () => sessionId,
|
|
||||||
enableJsonResponse: true, // Enable JSON response mode
|
|
||||||
});
|
|
||||||
|
|
||||||
// Use the new request-id based relay approach
|
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, shttpTransport);
|
|
||||||
|
|
||||||
return { shttpTransport, cleanup };
|
|
||||||
}
|
|
||||||
|
|
||||||
export async function getShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, cleanup: () => Promise<void>}> {
|
|
||||||
// Giving undefined here and setting the sessionId means the
|
// Giving undefined here and setting the sessionId means the
|
||||||
// transport wont try to create a new session.
|
// transport wont try to create a new session.
|
||||||
const shttpTransport = new StreamableHTTPServerTransport({
|
const shttpTransport = new StreamableHTTPServerTransport({
|
||||||
sessionIdGenerator: undefined,
|
sessionIdGenerator: undefined,
|
||||||
enableJsonResponse: true, // Use JSON response mode for all requests
|
|
||||||
})
|
})
|
||||||
shttpTransport.sessionId = sessionId;
|
shttpTransport.sessionId = sessionId;
|
||||||
|
|
||||||
// Use the new request-id based relay approach
|
// Use the new request-id based relay approach
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, shttpTransport);
|
const cleanup = redisRelayToMcpServer(sessionId, shttpTransport);
|
||||||
|
shttpTransport.onclose = cleanup;
|
||||||
return { shttpTransport, cleanup };
|
return shttpTransport;
|
||||||
}
|
}
|
||||||
Loading…
Reference in a new issue