Nour/swift subscriber

sdk/package.json

@@ -56,6 +56,7 @@
 		"solana-bankrun": "0.3.1",
 		"strict-event-emitter-types": "2.0.0",
 		"tweetnacl": "1.0.3",
+		"tweetnacl-util": "0.15.1",
 		"uuid": "8.3.2",
 		"yargs": "17.7.2",
 		"zstddec": "0.1.0"

sdk/src/index.ts

@@ -86,6 +86,7 @@ export * from './oracles/pythClient';
 export * from './oracles/pythPullClient';
 export * from './oracles/pythLazerClient';
 export * from './oracles/switchboardOnDemandClient';
+export * from './swift/swiftOrderSubscriber';
 export * from './tx/fastSingleTxSender';
 export * from './tx/retryTxSender';
 export * from './tx/whileValidTxSender';

sdk/src/swift/swiftOrderSubscriber.ts

@@ -0,0 +1,211 @@
+import {
+	DevnetPerpMarkets,
+	DriftClient,
+	DriftEnv,
+	getUserAccountPublicKey,
+	getUserStatsAccountPublicKey,
+	MainnetPerpMarkets,
+	MarketType,
+	OptionalOrderParams,
+	PostOnlyParams,
+	SwiftOrderParamsMessage,
+	UserMap,
+} from '..';
+import { Keypair, PublicKey, TransactionInstruction } from '@solana/web3.js';
+import nacl from 'tweetnacl';
+import { decodeUTF8 } from 'tweetnacl-util';
+import WebSocket from 'ws';
+
+export type SwiftOrderSubscriberConfig = {
+	driftClient: DriftClient;
+	userMap: UserMap;
+	driftEnv: DriftEnv;
+	endpoint?: string;
+	marketIndexes: number[];
+	keypair: Keypair;
+};
+
+export class SwiftOrderSubscriber {
+	private heartbeatTimeout: NodeJS.Timeout | null = null;
+	private readonly heartbeatIntervalMs = 60000;
+	private ws: WebSocket | null = null;
+	private driftClient: DriftClient;
+	private userMap: UserMap;
+	subscribed = false;
+
+	constructor(
+		private config: SwiftOrderSubscriberConfig,
+		private onOrder: (
+			orderMessageRaw: any,
+			swiftOrderParamsMessage: SwiftOrderParamsMessage
+		) => Promise<void>
+	) {
+		this.driftClient = config.driftClient;
+		this.userMap = config.userMap;
+	}
+
+	getSymbolForMarketIndex(marketIndex: number): string {
+		const markets =
+			this.config.driftEnv === 'devnet'
+				? DevnetPerpMarkets
+				: MainnetPerpMarkets;
+		return markets[marketIndex].symbol;
+	}
+
+	generateChallengeResponse(nonce: string): string {
+		const messageBytes = decodeUTF8(nonce);
+		const signature = nacl.sign.detached(
+			messageBytes,
+			this.config.keypair.secretKey
+		);
+		const signatureBase64 = Buffer.from(signature).toString('base64');
+		return signatureBase64;
+	}
+
+	handleAuthMessage(message: any): void {
+		if (message['channel'] === 'auth' && message['nonce'] != null) {
+			const signatureBase64 = this.generateChallengeResponse(message['nonce']);
+			this.ws?.send(
+				JSON.stringify({
+					pubkey: this.config.keypair.publicKey.toBase58(),
+					signature: signatureBase64,
+				})
+			);
+		}
+
+		if (
+			message['channel'] === 'auth' &&
+			message['message']?.toLowerCase() === 'authenticated'
+		) {
+			this.subscribed = true;
+			this.config.marketIndexes.forEach(async (marketIndex) => {
+				this.ws?.send(
+					JSON.stringify({
+						action: 'subscribe',
+						market_type: 'perp',
+						market_name: this.getSymbolForMarketIndex(marketIndex),
+					})
+				);
+				await new Promise((resolve) => setTimeout(resolve, 100));
+			});
+		}
+	}
+
+	async subscribe(): Promise<void> {
+		const endpoint =
+			this.config.endpoint || this.config.driftEnv === 'devnet'
+				? 'wss://master.swift.drift.trade/ws'
+				: 'wss://swift.drift.trade/ws';
+		const ws = new WebSocket(
+			endpoint + '?pubkey=' + this.config.keypair.publicKey.toBase58()
+		);
+		this.ws = ws;
+		ws.on('open', async () => {
+			console.log('Connected to the server');
+
+			ws.on('message', async (data: WebSocket.Data) => {
+				const message = JSON.parse(data.toString());
+				this.startHeartbeatTimer();
+
+				if (message['channel'] === 'auth') {
+					this.handleAuthMessage(message);
+				}
+
+				if (message['order']) {
+					const order = JSON.parse(message['order']);
+					const swiftOrderParamsBuf = Buffer.from(
+						order['order_message'],
+						'base64'
+					);
+					const swiftOrderParamsMessage: SwiftOrderParamsMessage =
+						this.driftClient.program.coder.types.decode(
+							'SwiftOrderParamsMessage',
+							swiftOrderParamsBuf
+						);
+
+					if (!swiftOrderParamsMessage.swiftOrderParams.price) {
+						console.error(
+							`order has no price: ${JSON.stringify(
+								swiftOrderParamsMessage.swiftOrderParams
+							)}`
+						);
+						return;
+					}
+
+					this.onOrder(order, swiftOrderParamsMessage);
+				}
+			});
+
+			ws.on('close', () => {
+				console.log('Disconnected from the server');
+				this.reconnect();
+			});
+
+			ws.on('error', (error: Error) => {
+				console.error('WebSocket error:', error);
+				this.reconnect();
+			});
+		});
+	}
+
+	async getPlaceAndMakeSwiftOrderIxs(
+		orderMessageRaw: any,
+		swiftOrderParamsMessage: SwiftOrderParamsMessage,
+		makerOrderParams: OptionalOrderParams
+	): Promise<TransactionInstruction[]> {
+		const swiftOrderParamsBuf = Buffer.from(
+			orderMessageRaw['order_message'],
+			'base64'
+		);
+		const takerAuthority = new PublicKey(orderMessageRaw['taker_authority']);
+		const takerUserPubkey = await getUserAccountPublicKey(
+			this.driftClient.program.programId,
+			takerAuthority,
+			swiftOrderParamsMessage.subAccountId
+		);
+		const takerUserAccount = (
+			await this.userMap.mustGet(takerUserPubkey.toString())
+		).getUserAccount();
+		const ixs = await this.driftClient.getPlaceAndMakeSwiftPerpOrderIxs(
+			swiftOrderParamsBuf,
+			Buffer.from(orderMessageRaw['order_signature'], 'base64'),
+			decodeUTF8(orderMessageRaw['uuid']),
+			{
+				taker: takerUserPubkey,
+				takerUserAccount,
+				takerStats: getUserStatsAccountPublicKey(
+					this.driftClient.program.programId,
+					takerUserAccount.authority
+				),
+			},
+			Object.assign({}, makerOrderParams, {
+				postOnly: PostOnlyParams.MUST_POST_ONLY,
+				immediateOrCancel: true,
+				marketType: MarketType.PERP,
+			})
+		);
+		return ixs;
+	}
+
+	private startHeartbeatTimer() {
+		if (this.heartbeatTimeout) {
+			clearTimeout(this.heartbeatTimeout);
+		}
+		this.heartbeatTimeout = setTimeout(() => {
+			console.warn('No heartbeat received within 30 seconds, reconnecting...');
+			this.reconnect();
+		}, this.heartbeatIntervalMs);
+	}
+
+	private reconnect() {
+		if (this.ws) {
+			this.ws.removeAllListeners();
+			this.ws.terminate();
+		}
+
+		console.log('Reconnecting to WebSocket...');
+		setTimeout(() => {
+			this.subscribe();
+		}, 1000);
+	}
+}

sdk/yarn.lock

@@ -3178,6 +3178,11 @@ tsutils@^3.21.0:
   dependencies:
     tslib "^1.8.1"
 
+tweetnacl-util@^0.15.1:
+  version "0.15.1"
+  resolved "https://registry.yarnpkg.com/tweetnacl-util/-/tweetnacl-util-0.15.1.tgz#b80fcdb5c97bcc508be18c44a4be50f022eea00b"
+  integrity sha512-RKJBIj8lySrShN4w6i/BonWp2Z/uxwC3h4y7xsRrpP59ZboCd0GpEVsOnMDYLMmKBpYhb5TgHzZXy7wTfYFBRw==
+
 [email protected]:
   version "1.0.3"
   resolved "https://registry.yarnpkg.com/tweetnacl/-/tweetnacl-1.0.3.tgz#ac0af71680458d8a6378d0d0d050ab1407d35596"