feat(dei): enforce JWT expiry & revocation; add kickIfInvalid; tests#939
feat(dei): enforce JWT expiry & revocation; add kickIfInvalid; tests#939Rhayxz wants to merge 1 commit intomiddleapi:mainfrom
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
dosubot
Bot
added
the
size:L
WalkthroughThe call handler in the durable event iterator is made async with an early token validation step. The WebSocket manager adds kickIfInvalid logic, close codes, an optional shouldKickWs hook, and makes publishEvent/sendEventsAfter async. Tests are updated for dynamic token timing, async behavior, and new expiry scenarios. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client
participant DO as Durable Object Handler
participant WM as WebSocket Manager
Client->>DO: RPC call via WebSocket
DO->>WM: kickIfInvalid(ws)
alt Token invalid
WM-->>DO: true
DO-->>Client: ORPCError(UNAUTHORIZED)
else Token valid
WM-->>DO: false
DO->>DO: Resolve allowed methods from token
DO->>Service: Invoke requested method
Service-->>DO: Result/Stream
DO-->>Client: Response/Events
end
sequenceDiagram
autonumber
participant Publisher as Event Publisher
participant WM as WebSocket Manager
participant WS* as Connected WebSockets
rect rgba(230,245,255,0.5)
note right of WM: publishEvent
Publisher->>WM: publishEvent(payload)
loop For each WS with hibernationId
WM->>WM: kickIfInvalid(ws)?
alt Invalid (expired/revoked)
WM-->>WS*: Close (4001/4003)
else Valid
WM-->>WS*: Send hibernation event
end
end
end
sequenceDiagram
autonumber
participant WM as WebSocket Manager
participant WS as WebSocket
note right of WM: sendEventsAfter
WM->>WM: kickIfInvalid(ws)?
alt Invalid
WM-->>WS: Close (4001/4003)
else Valid
WM->>WM: Load events after cursor
WM-->>WS: Send events
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @Rhayxz, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces robust JWT (JSON Web Token) expiry and revocation enforcement within the Durable Event Iterator (DEI) system. It ensures that WebSocket connections are validated against their JWTs both when events are published and during inbound RPC calls, enhancing security by automatically disconnecting clients with expired or revoked tokens.
Highlights
- JWT Expiry Enforcement: WebSocket connections are now automatically closed with a 4001 status code if their associated JWT has expired.
- Optional JWT Revocation Hook: A new shouldKickWs option allows for custom logic to revoke JWTs, closing connections with a 4003 status code.
- Inbound RPC Call Protection: The durableEventIteratorRouter now validates the JWT of the current WebSocket connection before processing any inbound RPC calls, preventing unauthorized access.
- Unified Token Validation Logic: A new kickIfInvalid method centralizes the logic for checking token validity (expiry and revocation) and closing invalid WebSocket connections.
- Comprehensive Test Coverage: New and updated tests ensure the correct behavior of JWT expiry and revocation, including the proper closing of WebSocket connections.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces important security enhancements by enforcing JWT expiration and adding a mechanism for token revocation. The changes ensure that tokens are validated before processing RPC calls and when publishing events to websockets. New tests have been added to cover these scenarios, including checks for expired tokens. My main feedback is regarding error handling in the new token revocation hook, where silently swallowing errors could pose a security risk.
| try { | ||
| isRevoked = await this.options?.shouldKickWs?.(payload) | ||
| } | ||
| catch {} |
There was a problem hiding this comment.
The try...catch block around the shouldKickWs call silently swallows any errors. This creates a "fail-open" security posture, where a failure in the revocation check (e.g., a database error) results in the token being treated as valid. This could allow users with revoked tokens to maintain access if the check mechanism fails.
At a minimum, the error should be logged to aid in debugging. For enhanced security, consider a "fail-closed" approach where an error during the revocation check causes the token to be treated as invalid.
The optional chaining (?.) on this.options?.shouldKickWs within the try block is also redundant given the if (this.options.shouldKickWs) check on the line above.
try {
isRevoked = await this.options.shouldKickWs(payload)
}
catch (e) {
console.error('Error in `shouldKickWs` revocation check. Defaulting to not revoked.', e)
}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
packages/durable-event-iterator/src/durable-object/handler.test.ts (1)
10-16: Fix JWT time units in test (use seconds, not ms).The router subscribes using
new Date((payload.iat - 1) * 1000), assumingiatis seconds (JWT NumericDate). Heredate = Date.now()makesiat/expmilliseconds, producing an incorrect “after” date. Align with seconds for consistency withkickIfInvalidand other tests.Apply:
- const date = Date.now() + const dateSec = Math.floor(Date.now() / 1000) const tokenPayload = { att: { some: 'attachment' }, - iat: date, - exp: date + 1000, + iat: dateSec, + exp: dateSec + 1000, chn: 'test-channel', rpc: ['someMethod'], } ... - new Date((date - 1) * 1000), + new Date((dateSec - 1) * 1000),Optionally freeze time with fake timers for determinism.
Also applies to: 49-53
packages/durable-event-iterator/src/durable-object/websocket-manager.ts (2)
122-131: Possible null deref if no prior attachment.
old = deserializeAttachment(ws)may benull/undefined(see tests’ stub). Accessingold[...]would throw. Default to{}and use optional chaining.- serializeAttachment(ws: WebSocket, attachment: TWsAttachment): void { - const old = this.deserializeAttachment(ws) + serializeAttachment(ws: WebSocket, attachment: TWsAttachment): void { + const old = this.deserializeAttachment(ws) ?? ({} as any) ws.serializeAttachment({ ...attachment, - [DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY]: old[DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY], - [DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY]: old[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY], + [DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY]: old?.[DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY], + [DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY]: old?.[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY], }) }
132-140: Object spread over possibly-nulloldcan throw.Spreading
null/undefinedresults in a TypeError. Defaultoldto{}.- serializeInternalAttachment(ws: WebSocket, attachment: Partial<DurableEventIteratorObjectWebsocketInternalAttachment<TTokenAttachment>>): void { - const old = this.deserializeAttachment(ws) + serializeInternalAttachment(ws: WebSocket, attachment: Partial<DurableEventIteratorObjectWebsocketInternalAttachment<TTokenAttachment>>): void { + const old = this.deserializeAttachment(ws) ?? ({} as any) ws.serializeAttachment({ - ...old, + ...(old as any), [DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY]: attachment[DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY] ?? old?.[DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY], [DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY]: attachment[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY] ?? old?.[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY], }) }
🧹 Nitpick comments (5)
packages/durable-event-iterator/src/durable-object/handler.ts (2)
43-48: Early invalid-token guard looks good.Async
kickIfInvalidgate before dispatch is correct and maps invalids to UNAUTHORIZED. Consider distinguishing expired vs revoked in the error message in future, if useful.
27-41: Guard against missing token payload insubscribehandlerIt’s possible that
deserializeAttachment(...)[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY]comes backundefined, causingpayload.iatto throw. Add a quick null check and early return before using the payload:subscribe: base.subscribe.handler(({ context, lastEventId }) => { return new HibernationEventIterator<any>((hibernationId) => { context.websocketManager.serializeInternalAttachment(context.currentWebsocket, { [DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY]: hibernationId, }) - const payload = context.websocketManager.deserializeAttachment(context.currentWebsocket)[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY] + const att = context.websocketManager.deserializeAttachment(context.currentWebsocket) + const payload = att?.[DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY] + if (!payload) { + // No token payload—nothing to replay + return + } context.websocketManager.sendEventsAfter( context.currentWebsocket, hibernationId, lastEventId !== undefined ? lastEventId : new Date((payload.iat - 1) * 1000), ) }) }),– All existing uses of
iat/exphave been confirmed to be in seconds per schema and tests; no further changes needed there.packages/durable-event-iterator/src/durable-object/websocket-manager.test.ts (1)
1-271: Add a revocation test (4003) via shouldKickWs.You added the hook; add one targeted test to lock behavior.
+it('publishEvent closes sockets when shouldKickWs revokes (4003)', async () => { + const ctx = createDurableObjectState() + const storage = new DurableEventIteratorObjectEventStorage(ctx) + const websocket = createCloudflareWebsocket() + ;(websocket as any).close = vi.fn() + const manager = new DurableEventIteratorObjectWebsocketManager<any, any, any>( + ctx, storage, { customJsonSerializers: [], shouldKickWs: () => true } + ) + const nowSec = Math.floor(Date.now() / 1000) + manager.serializeInternalAttachment(websocket, { + [DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY]: { att: { att: true }, chn: 'chn', iat: nowSec, exp: nowSec + 3600, rpc: ['x'] }, + [DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY]: 'hid', + }) + await manager.publishEvent([websocket], { test: 'event' }) + expect((websocket as any).close).toHaveBeenCalledWith(4003, 'token revoked') + expect(websocket.send).not.toHaveBeenCalled() +})packages/durable-event-iterator/src/durable-object/websocket-manager.ts (2)
85-99: Shield broadcast loops fromws.sendexceptions.A single bad socket send could abort the loop. Wrap sends to continue.
- ws.send(encodeHibernationRPCEvent(hibernationEventIteratorId, payload, this.options)) + try { + ws.send(encodeHibernationRPCEvent(hibernationEventIteratorId, payload, this.options)) + } catch {} ... - ws.send(encodeHibernationRPCEvent(hibernationId, event, this.options)) + try { + ws.send(encodeHibernationRPCEvent(hibernationId, event, this.options)) + } catch {}Also applies to: 109-120
8-11: Type hint: bind hook to the generic attachment type.Minor: you could tie
shouldKickWsto the manager’s genericTTokenAttachmentto provide stronger typing foratt.-export interface DurableEventIteratorObjectWebsocketManagerOptions extends StandardRPCJsonSerializerOptions { - shouldKickWs?: (payload: DurableEventIteratorTokenPayload) => boolean | Promise<boolean> +export interface DurableEventIteratorObjectWebsocketManagerOptions<TAtt = unknown> extends StandardRPCJsonSerializerOptions { + shouldKickWs?: (payload: DurableEventIteratorTokenPayload & { att: TAtt }) => boolean | Promise<boolean> }…and pass
TTokenAttachmentwhen constructing options. Low priority.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
packages/durable-event-iterator/src/durable-object/handler.test.ts(1 hunks)packages/durable-event-iterator/src/durable-object/handler.ts(1 hunks)packages/durable-event-iterator/src/durable-object/websocket-manager.test.ts(10 hunks)packages/durable-event-iterator/src/durable-object/websocket-manager.ts(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
packages/durable-event-iterator/src/durable-object/websocket-manager.test.ts (4)
packages/durable-event-iterator/src/durable-object/consts.ts (2)
DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY(2-2)DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY(1-1)packages/durable-event-iterator/tests/shared.ts (2)
createDurableObjectState(3-24)createCloudflareWebsocket(26-34)packages/durable-event-iterator/src/durable-object/event-storage.ts (1)
DurableEventIteratorObjectEventStorage(23-145)packages/durable-event-iterator/src/durable-object/websocket-manager.ts (1)
DurableEventIteratorObjectWebsocketManager(36-145)
packages/durable-event-iterator/src/durable-object/websocket-manager.ts (2)
packages/durable-event-iterator/src/schemas.ts (1)
DurableEventIteratorTokenPayload(3-3)packages/durable-event-iterator/src/durable-object/consts.ts (2)
DURABLE_EVENT_ITERATOR_TOKEN_PAYLOAD_KEY(2-2)DURABLE_EVENT_ITERATOR_HIBERNATION_ID_KEY(1-1)
🔇 Additional comments (5)
packages/durable-event-iterator/src/durable-object/websocket-manager.test.ts (4)
8-10: Mocking hibernation encoder is appropriate.
Keeps assertions focused on manager behavior.
24-26: Good switch to seconds-based iat/exp.
Matches runtime checks inkickIfInvalid.Also applies to: 76-83, 140-149
166-225: Expiry path coverage (publishEvent) is solid.Asserts 4001 close and skip-send. Nice use of fake timers.
227-271: Expiry path coverage (sendEventsAfter) is solid.Verifies early return and 4001 close without sends.
packages/durable-event-iterator/src/durable-object/websocket-manager.ts (1)
47-83: kickIfInvalid implementation is correct and side-effect-safe.Checks exp first, then optional revocation, and returns a boolean gate.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
More templates
@orpc/arktype
@orpc/client
@orpc/contract
@orpc/experimental-durable-event-iterator
@orpc/hey-api
@orpc/interop
@orpc/json-schema
@orpc/nest
@orpc/openapi
@orpc/openapi-client
@orpc/otel
@orpc/react
@orpc/react-query
@orpc/experimental-react-swr
@orpc/server
@orpc/shared
@orpc/solid-query
@orpc/standard-server
@orpc/standard-server-aws-lambda
@orpc/standard-server-fetch
@orpc/standard-server-node
@orpc/standard-server-peer
@orpc/svelte-query
@orpc/tanstack-query
@orpc/trpc
@orpc/valibot
@orpc/vue-colada
@orpc/vue-query
@orpc/zod
commit: |
|
Your approach still missing some cases. For example, a procedure can return an Event Iterator. Currently, you only check it on the first call, while an Event Iterator can live for a long time, especially with a Hibernation Event Iterator. Maybe I can work on this soon to improve the DX with a proxy, so every time a WebSocket message is sent or received, it will be checked |
Yes please, I tried to touch some event iterator code, but ended up running into a few issues as I'm quite unfamiliar how the ORPC event iterators work under the hood. |
|
I've created #965 with a different approach that better addresses this issue. |
2 participants
Ensures JWT is checked on send and on inbound RPC. Adds optional revocation hook, closes WS with 4001/4003, updates router to guard .call via kickIfInvalid, and adds expiry tests.
Summary by CodeRabbit