Skip to content

Commit

Permalink
feat(WebsocketManager): retroactive token setting (#10418)
Browse files Browse the repository at this point in the history
Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
  • Loading branch information
didinele and kodiakhq[bot] authored Jul 31, 2024
1 parent 8f97d2b commit de94eaf
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ export interface IContextFetchingStrategy {
}

export async function managerToFetchingStrategyOptions(manager: WebSocketManager): Promise<FetchingStrategyOptions> {
/* eslint-disable @typescript-eslint/unbound-method */
const {
buildIdentifyThrottler,
buildStrategy,
Expand All @@ -44,10 +43,10 @@ export async function managerToFetchingStrategyOptions(manager: WebSocketManager
rest,
...managerOptions
} = manager.options;
/* eslint-enable @typescript-eslint/unbound-method */

return {
...managerOptions,
token: manager.token,
gatewayInformation: await manager.fetchGatewayInformation(),
shardCount: await manager.getShardCount(),
};
Expand Down
2 changes: 1 addition & 1 deletion packages/ws/src/utils/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ export const DefaultWebSocketManagerOptions = {
handshakeTimeout: 30_000,
helloTimeout: 60_000,
readyTimeout: 15_000,
} as const satisfies OptionalWebSocketManagerOptions;
} as const satisfies Omit<OptionalWebSocketManagerOptions, 'token'>;

export const ImportantGatewayOpcodes = new Set([
GatewayOpcodes.Heartbeat,
Expand Down
38 changes: 33 additions & 5 deletions packages/ws/src/ws/WebSocketManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -67,10 +67,6 @@ export interface RequiredWebSocketManagerOptions {
* The REST instance to use for fetching gateway information
*/
rest: REST;
/**
* The token to use for identifying with the gateway
*/
token: string;
}

/**
Expand Down Expand Up @@ -172,6 +168,12 @@ export interface OptionalWebSocketManagerOptions {
* ```
*/
shardIds: number[] | ShardRange | null;
/**
* The token to use for identifying with the gateway
*
* If not provided, the token must be set using {@link WebSocketManager.setToken}
*/
token: string;
/**
* Function used to store session information for a given shard
*/
Expand Down Expand Up @@ -211,10 +213,12 @@ export interface ManagerShardEventsMap {
}

export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> implements AsyncDisposable {
#token: string | null = null;

/**
* The options being used by this manager
*/
public readonly options: WebSocketManagerOptions;
public readonly options: Omit<WebSocketManagerOptions, 'token'>;

/**
* Internal cache for a GET /gateway/bot result
Expand All @@ -236,10 +240,26 @@ export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> i
*/
private readonly strategy: IShardingStrategy;

/**
* Gets the token set for this manager. If no token is set, an error is thrown.
* To set the token, use {@link WebSocketManager.setToken} or pass it in the options.
*
* @remarks
* This getter is mostly used to pass the token to the sharding strategy internally, there's not much reason to use it.
*/
public get token(): string {
if (!this.#token) {
throw new Error('Token has not been set');
}

return this.#token;
}

public constructor(options: CreateWebSocketManagerOptions) {
super();
this.options = { ...DefaultWebSocketManagerOptions, ...options };
this.strategy = this.options.buildStrategy(this);
this.#token = options.token ?? null;
}

/**
Expand Down Expand Up @@ -334,6 +354,14 @@ export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> i
await this.strategy.connect();
}

public setToken(token: string): void {
if (this.#token) {
throw new Error('Token has already been set');
}

this.#token = token;
}

public destroy(options?: Omit<WebSocketShardDestroyOptions, 'recover'>) {
return this.strategy.destroy(options);
}
Expand Down

0 comments on commit de94eaf

Please sign in to comment.