codeant-ai-for-open-source[bot] commented on code in PR #45404:
URL: https://github.com/apache/superset/pull/45404#discussion_r4427616108
##########
superset-websocket/src/index.ts:
##########
@@ -715,72 +729,204 @@ export const readSocketIdentity = (
};
};
+/** An inbound `auth` frame: `{type: 'auth', token, tab_id?}`. */
+interface AuthFrame {
+ token: string;
+ tabId?: string;
+}
+
/**
- * WebSocket `connection` event handler. `identity` is the one already verified
- * during the HTTP upgrade (see `httpUpgrade`), so each connection's JWT is
- * verified exactly once.
+ * Parse a raw inbound frame into an `auth` frame, or `null` when it is
anything
+ * else (malformed JSON, another message type, or a missing token). An invalid
or
+ * absent `tab_id` is dropped rather than failing the frame, so the socket
falls
+ * back to principal-grain delivery.
*/
-export const wsConnection = (
- ws: WebSocket,
- identity: SocketIdentity,
- tabId?: string,
-) => {
- const { channel } = identity;
-
- // Refuse the connection if a configured connection limit has been reached,
- // before tracking it against the internal registries. The cap is keyed on
the
- // principal channel (which counts all of a principal's tabs); the per-tab
- // channel is a delivery refinement and is not separately limit-checked.
- const limitReason = connectionLimitReason(channel);
- if (limitReason) {
- statsd.increment('ws_connection_rejected');
- logger.warn(`Refusing connection on channel ${channel}: ${limitReason}`);
- ws.close(CONNECTION_LIMIT_CLOSE_CODE, limitReason);
- return;
+const parseAuthFrame = (raw: WebSocket.RawData): AuthFrame | null => {
+ let parsed: unknown;
+ try {
+ parsed = JSON.parse(raw.toString());
+ } catch {
+ return null;
}
+ if (!parsed || typeof parsed !== 'object') return null;
+ const {
+ type,
+ token,
+ tab_id: tabId,
+ } = parsed as {
+ type?: unknown;
+ token?: unknown;
+ tab_id?: unknown;
+ };
+ if (type !== 'auth' || typeof token !== 'string' || token.length === 0) {
+ return null;
+ }
+ return {
+ token,
+ tabId: typeof tabId === 'string' && isValidTabId(tabId) ? tabId :
undefined,
+ };
+};
- // When the browser advertised a tab id at connect, also register the socket
- // under a per-tab channel derived from its authorized principal channel
- // (`<channel>:<tabId>`). The tab id is client-supplied but always prefixed
by
- // the authorized principal channel, so it can never address another
- // principal's sockets. Task-status can then target one tab, while the
- // principal channel still reaches all of a principal's tabs.
- const tabChannel = tabId ? `${channel}:${tabId}` : undefined;
-
- const socketInstance: SocketInstance = {
- ws,
- channel,
- tabChannel,
- identity,
- pongTs: Date.now(),
+const sendAuthOk = (ws: WebSocket, identity: SocketIdentity): void => {
+ const message: OutboundMessage = {
+ topic: AUTH_REPLY_TOPIC,
+ payload: {
+ status: 'ok',
+ expires_in: Math.max(
+ 0,
+ Math.floor((identity.tokenExpiresAtMs - Date.now()) / 1000),
+ ),
+ },
};
+ ws.send(JSON.stringify(message));
+};
+
+/**
+ * Removes a socket id from a channel's socket list, dropping the channel when
it
+ * becomes empty.
+ */
+const untrackFromChannel = (channel: string, socketId: string): void => {
+ const entry = channels[channel];
+ if (!entry) return;
+ entry.sockets = entry.sockets.filter(id => id !== socketId);
+ if (entry.sockets.length === 0) delete channels[channel];
+};
- // Register once (mints a single socket id), then also index that same id
under
- // the per-tab channel - never call trackClient twice (it would mint a second
- // id and a duplicate registry entry).
- const socketId = trackClient(channel, socketInstance);
+/**
+ * Binds an already tracked socket to the per-tab channel derived from its
+ * authorized principal channel (`<channel>:<tabId>`), replacing a previous
+ * binding. The tab id is client-supplied but always prefixed by the authorized
+ * principal channel, so it can never address another principal's sockets.
+ * Task-status can then target one tab, while the principal channel still
reaches
+ * all of a principal's tabs.
+ */
+const bindTabChannel = (
+ socketId: string,
+ socketInstance: SocketInstance,
+ tabId: string | undefined,
+): void => {
+ const tabChannel = tabId ? `${socketInstance.channel}:${tabId}` : undefined;
+ if (tabChannel === socketInstance.tabChannel) return;
+ if (socketInstance.tabChannel) {
+ untrackFromChannel(socketInstance.tabChannel, socketId);
+ }
if (tabChannel) {
(channels[tabChannel] ??= { sockets: [] }).sockets.push(socketId);
}
- logger.debug(
- `socket ${socketId} connected on channel ${channel}` +
- (tabChannel ? ` (tab channel ${tabChannel})` : ''),
- );
+ socketInstance.tabChannel = tabChannel;
+};
+
+/**
+ * WebSocket `connection` event handler. The connection starts `pending_auth`:
+ * it is tracked in no registry and routed no message until its first `auth`
frame
+ * carries a valid ticket, at which point it is registered on its principal
channel
+ * (and tab channel). Later `auth` frames refresh the ticket on the open
socket.
+ */
+export const wsConnection = (ws: WebSocket) => {
+ // Set once the first `auth` frame is accepted; `undefined` while
pending_auth.
+ let socketId: string | undefined;
+
+ const authTimer = setTimeout(() => {
+ if (socketId === undefined) {
+ statsd.increment('ws_auth_timeout');
+ ws.close(AUTH_CLOSE_CODE, 'authentication timeout');
+ }
+ }, AUTH_TIMEOUT_MS);
Review Comment:
✅ **CodeAnt verified this suggestion was addressed in subsequent commits and
marked this thread resolved** as of `88ea23d`.
The upgrade handler checks connection limits and the pending-connection cap
before accepting a pending reservation. It increments `pendingConnections` and
releases the reservation on close, timeout, or authentication.
<sub>If that's not right, unresolve this thread and CodeAnt will leave it
open.</sub>
<!-- codeant-auto-resolve-reply -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
