diff --git a/sdk/package.json b/sdk/package.json index 865e72866..d85b10a58 100644 --- a/sdk/package.json +++ b/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" diff --git a/sdk/src/index.ts b/sdk/src/index.ts index 7d141def0..f61687ece 100644 --- a/sdk/src/index.ts +++ b/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'; diff --git a/sdk/src/swift/swiftOrderSubscriber.ts b/sdk/src/swift/swiftOrderSubscriber.ts new file mode 100644 index 000000000..7b9bcdc22 --- /dev/null +++ b/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 + ) { + 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 { + 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 { + 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); + } +} diff --git a/sdk/yarn.lock b/sdk/yarn.lock index 85fb9e30f..a17e43e5a 100644 --- a/sdk/yarn.lock +++ b/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== + tweetnacl@1.0.3: version "1.0.3" resolved "https://registry.yarnpkg.com/tweetnacl/-/tweetnacl-1.0.3.tgz#ac0af71680458d8a6378d0d0d050ab1407d35596"