Include AuthInfo from bearer validation in transport request handler

This commit is contained in:
Allen Zhou 2025-02-28 15:20:55 -08:00
commit af025e26a7
No known key found for this signature in database
GPG key ID: 5817C145DE5C5556
3 changed files with 22 additions and 10 deletions

View file

@ -4,6 +4,7 @@ import { Transport } from "../shared/transport.js";
import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
import getRawBody from "raw-body";
import contentType from "content-type";
import { AuthInfo } from "./auth/types.js";
const MAXIMUM_MESSAGE_SIZE = "4mb";
@ -18,7 +19,7 @@ export class SSEServerTransport implements Transport {
onclose?: () => void;
onerror?: (error: Error) => void;
onmessage?: (message: JSONRPCMessage) => void;
onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
/**
* Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL identified by `_endpoint`.
@ -66,7 +67,7 @@ export class SSEServerTransport implements Transport {
* This should be called when a POST request is made to send a message to the server.
*/
async handlePostMessage(
req: IncomingMessage,
req: IncomingMessage & { auth?: AuthInfo },
res: ServerResponse,
parsedBody?: unknown,
): Promise<void> {
@ -75,6 +76,7 @@ export class SSEServerTransport implements Transport {
res.writeHead(500).end(message);
throw new Error(message);
}
const authInfo: AuthInfo | undefined = req.auth;
let body: string | unknown;
try {
@ -94,7 +96,7 @@ export class SSEServerTransport implements Transport {
}
try {
await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body);
await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, authInfo);
} catch {
res.writeHead(400).end(`Invalid message: ${body}`);
return;
@ -106,7 +108,7 @@ export class SSEServerTransport implements Transport {
/**
* Handle a client message, regardless of how it arrived. This can be used to inform the server of messages that arrive via a means different than HTTP POST.
*/
async handleMessage(message: unknown): Promise<void> {
async handleMessage(message: unknown, authInfo?: AuthInfo): Promise<void> {
let parsedMessage: JSONRPCMessage;
try {
parsedMessage = JSONRPCMessageSchema.parse(message);
@ -115,7 +117,7 @@ export class SSEServerTransport implements Transport {
throw error;
}
this.onmessage?.(parsedMessage);
this.onmessage?.(parsedMessage, authInfo);
}
async close(): Promise<void> {

View file

@ -19,6 +19,7 @@ import {
ServerCapabilities,
} from "../types.js";
import { Transport } from "./transport.js";
import { AuthInfo } from "../server/auth/types.js";
/**
* Callback for progress notifications.
@ -88,6 +89,11 @@ export type RequestHandlerExtra = {
* An abort signal used to communicate if the request was cancelled from the sender's side.
*/
signal: AbortSignal;
/**
* Information about a validated access token, provided to request handlers.
*/
authInfo?: AuthInfo;
};
/**
@ -232,11 +238,11 @@ export abstract class Protocol<
this._onerror(error);
};
this._transport.onmessage = (message) => {
this._transport.onmessage = (message, authInfo) => {
if (!("method" in message)) {
this._onresponse(message);
} else if ("id" in message) {
this._onrequest(message);
this._onrequest(message, authInfo);
} else {
this._onnotification(message);
}
@ -282,7 +288,7 @@ export abstract class Protocol<
);
}
private _onrequest(request: JSONRPCRequest): void {
private _onrequest(request: JSONRPCRequest, authInfo?: AuthInfo): void {
const handler =
this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler;
@ -309,7 +315,7 @@ export abstract class Protocol<
// Starting with Promise.resolve() puts any synchronous errors into the monad as well.
Promise.resolve()
.then(() => handler(request, { signal: abortController.signal }))
.then(() => handler(request, { signal: abortController.signal, authInfo }))
.then(
(result) => {
if (abortController.signal.aborted) {

View file

@ -1,3 +1,4 @@
import { AuthInfo } from "../server/auth/types.js";
import { JSONRPCMessage } from "../types.js";
/**
@ -39,6 +40,9 @@ export interface Transport {
/**
* Callback for when a message (request or response) is received over the connection.
*
* Includes the authInfo if the transport is authenticated.
*
*/
onmessage?: (message: JSONRPCMessage) => void;
onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
}