Respect CLIENT_STATE_V2 capability

5 files changed, 46 insertions, 6 deletions

diff --git a/src/gateway/opcodes/Identify.ts b/src/gateway/opcodes/Identify.ts
index 6aa152d3..f035fe15 100644
--- a/src/gateway/opcodes/Identify.ts
+++ b/src/gateway/opcodes/Identify.ts
@@ -16,7 +16,15 @@
 	along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
-import { WebSocket, Payload, setupListener } from "@spacebar/gateway";
+import {
+	WebSocket,
+	Payload,
+	setupListener,
+	Capabilities,
+	CLOSECODES,
+	OPCODES,
+	Send,
+} from "@spacebar/gateway";
 import {
 	checkToken,
 	Intents,
@@ -45,11 +53,9 @@ import {
 	Permissions,
 	DMChannel,
 	GuildOrUnavailable,
+	Recipient,
 } from "@spacebar/util";
-import { Send } from "../util/Send";
-import { CLOSECODES, OPCODES } from "../util/Constants";
 import { check } from "./instanceOf";
-import { Recipient } from "@spacebar/util";
 
 // TODO: user sharding
 // TODO: check privileged intents, if defined in the config
@@ -73,6 +79,8 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 	check.call(this, IdentifySchema, data.d);
 	const identify: IdentifySchema = data.d;
 
+	this.capabilities = new Capabilities(identify.capabilities || 0);
+
 	// Check auth
 	// TODO: the checkToken call will fetch user, and then we have to refetch with different select
 	// checkToken should be able to select what we want
@@ -374,7 +382,9 @@ export async function onIdentify(this: WebSocket, data: Payload) {
 			: undefined,
 		user: user.toPrivateUser(),
 		user_settings: user.settings,
-		guilds: guilds.map((x) => new ReadyGuildDTO(x).toJSON()),
+		guilds: this.capabilities.has(Capabilities.FLAGS.CLIENT_STATE_V2)
+			? guilds.map((x) => new ReadyGuildDTO(x).toJSON())
+			: guilds,
 		relationships: user.relationships.map((x) => x.toPublicRelationship()),
 		read_state: {
 			entries: read_states,
diff --git a/src/gateway/util/Capabilities.ts b/src/gateway/util/Capabilities.ts
new file mode 100644
index 00000000..6c94bb45
--- /dev/null
+++ b/src/gateway/util/Capabilities.ts
@@ -0,0 +1,26 @@
+import { BitField, BitFieldResolvable, BitFlag } from "@spacebar/util";
+
+export type CapabilityResolvable = BitFieldResolvable | CapabilityString;
+type CapabilityString = keyof typeof Capabilities.FLAGS;
+
+export class Capabilities extends BitField {
+	static FLAGS = {
+		// Thanks, Opencord!
+		// https://github.com/MateriiApps/OpenCord/blob/master/app/src/main/java/com/xinto/opencord/gateway/io/Capabilities.kt
+		LAZY_USER_NOTES: BitFlag(0),
+		NO_AFFINE_USER_IDS: BitFlag(1),
+		VERSIONED_READ_STATES: BitFlag(2),
+		VERSIONED_USER_GUILD_SETTINGS: BitFlag(3),
+		DEDUPLICATE_USER_OBJECTS: BitFlag(4),
+		PRIORITIZED_READY_PAYLOAD: BitFlag(5),
+		MULTIPLE_GUILD_EXPERIMENT_POPULATIONS: BitFlag(6),
+		NON_CHANNEL_READ_STATES: BitFlag(7),
+		AUTH_TOKEN_REFRESH: BitFlag(8),
+		USER_SETTINGS_PROTO: BitFlag(9),
+		CLIENT_STATE_V2: BitFlag(10),
+		PASSIVE_GUILD_UPDATE: BitFlag(11),
+	};
+
+	any = (capability: CapabilityResolvable) => super.any(capability);
+	has = (capability: CapabilityResolvable) => super.has(capability);
+}
diff --git a/src/gateway/util/WebSocket.ts b/src/gateway/util/WebSocket.ts
index 972129c7..833756ff 100644
--- a/src/gateway/util/WebSocket.ts
+++ b/src/gateway/util/WebSocket.ts
@@ -19,6 +19,7 @@
 import { Intents, ListenEventOpts, Permissions } from "@spacebar/util";
 import WS from "ws";
 import { Deflate, Inflate } from "fast-zlib";
+import { Capabilities } from "./Capabilities";
 // import { Client } from "@spacebar/webrtc";
 
 export interface WebSocket extends WS {
@@ -40,5 +41,6 @@ export interface WebSocket extends WS {
 	events: Record<string, undefined | (() => unknown)>;
 	member_events: Record<string, () => unknown>;
 	listen_options: ListenEventOpts;
+	capabilities?: Capabilities;
 	// client?: Client;
 }
diff --git a/src/gateway/util/index.ts b/src/gateway/util/index.ts
index 627f12b2..6ef694d9 100644
--- a/src/gateway/util/index.ts
+++ b/src/gateway/util/index.ts
@@ -21,3 +21,4 @@ export * from "./Send";
 export * from "./SessionUtils";
 export * from "./Heartbeat";
 export * from "./WebSocket";
+export * from "./Capabilities";
diff --git a/src/util/interfaces/Event.ts b/src/util/interfaces/Event.ts
index 21c34c6c..16df48aa 100644
--- a/src/util/interfaces/Event.ts
+++ b/src/util/interfaces/Event.ts
@@ -42,6 +42,7 @@ import {
 	UserPrivate,
 	ReadyUserGuildSettingsEntries,
 	ReadyPrivateChannel,
+	GuildOrUnavailable,
 } from "@spacebar/util";
 
 export interface Event {
@@ -73,7 +74,7 @@ export interface ReadyEventData {
 	user: UserPrivate;
 	private_channels: ReadyPrivateChannel[]; // this will be empty for bots
 	session_id: string; // resuming
-	guilds: IReadyGuildDTO[];
+	guilds: IReadyGuildDTO[] | GuildOrUnavailable[]; // depends on capability
 	analytics_token?: string;
 	connected_accounts?: ConnectedAccount[];
 	consents?: {