Wphan/custom coder

sdk/src/decode/customCoder.ts

@@ -0,0 +1,114 @@
+import { Buffer } from 'buffer';
+import camelcase from 'camelcase';
+import { Idl, IdlTypeDef } from '@coral-xyz/anchor/dist/cjs/idl';
+import {
+	AccountsCoder,
+	BorshAccountsCoder,
+	BorshEventCoder,
+	BorshInstructionCoder,
+	Coder,
+} from '@coral-xyz/anchor/dist/cjs/coder';
+import { BorshTypesCoder } from '@coral-xyz/anchor/dist/cjs/coder/borsh/types';
+import { discriminator } from '@coral-xyz/anchor/dist/cjs/coder/borsh/discriminator';
+
+export class CustomBorshCoder<
+	A extends string = string,
+	T extends string = string,
+> implements Coder
+{
+	readonly idl: Idl;
+
+	/**
+	 * Instruction coder.
+	 */
+	readonly instruction: BorshInstructionCoder;
+
+	/**
+	 * Account coder.
+	 */
+	readonly accounts: CustomBorshAccountsCoder<A>;
+
+	/**
+	 * Coder for events.
+	 */
+	readonly events: BorshEventCoder;
+
+	/**
+	 * Coder for user-defined types.
+	 */
+	readonly types: BorshTypesCoder<T>;
+
+	constructor(idl: Idl) {
+		this.instruction = new BorshInstructionCoder(idl);
+		this.accounts = new CustomBorshAccountsCoder(idl);
+		this.events = new BorshEventCoder(idl);
+		this.types = new BorshTypesCoder(idl);
+		this.idl = idl;
+	}
+}
+
+/**
+ * Custom accounts coder that wraps BorshAccountsCoder to fix encode buffer sizing.
+ */
+export class CustomBorshAccountsCoder<A extends string = string>
+	implements AccountsCoder
+{
+	private baseCoder: BorshAccountsCoder<A>;
+	private idl: Idl;
+
+	public constructor(idl: Idl) {
+		this.baseCoder = new BorshAccountsCoder<A>(idl);
+		this.idl = idl;
+	}
+
+	public async encode<T = any>(accountName: A, account: T): Promise<Buffer> {
+		const idlAcc = this.idl.accounts?.find((acc) => acc.name === accountName);
+		if (!idlAcc) {
+			throw new Error(`Unknown account not found in idl: ${accountName}`);
+		}
+
+		const buffer = Buffer.alloc(this.size(idlAcc)); // fix encode issue - use proper size instead of fixed 1000
+		const layout = this.baseCoder['accountLayouts'].get(accountName);
+		if (!layout) {
+			throw new Error(`Unknown account: ${accountName}`);
+		}
+		const len = layout.encode(account, buffer);
+		const accountData = buffer.slice(0, len);
+		const discriminator = BorshAccountsCoder.accountDiscriminator(accountName);
+		return Buffer.concat([discriminator, accountData]);
+	}
+
+	// Delegate all other methods to the base coder
+	public decode<T = any>(accountName: A, data: Buffer): T {
+		return this.baseCoder.decode(accountName, data);
+	}
+
+	public decodeAny<T = any>(data: Buffer): T {
+		return this.baseCoder.decodeAny(data);
+	}
+
+	public decodeUnchecked<T = any>(accountName: A, ix: Buffer): T {
+		return this.baseCoder.decodeUnchecked(accountName, ix);
+	}
+
+	public memcmp(accountName: A, appendData?: Buffer): any {
+		return this.baseCoder.memcmp(accountName, appendData);
+	}
+
+	public size(idlAccount: IdlTypeDef): number {
+		return this.baseCoder.size(idlAccount);
+	}
+
+	/**
+	 * Calculates and returns a unique 8 byte discriminator prepended to all anchor accounts.
+	 *
+	 * @param name The name of the account to calculate the discriminator.
+	 */
+	public static accountDiscriminator(name: string): Buffer {
+		const discriminatorPreimage = `account:${camelcase(name, {
+			pascalCase: true,
+			preserveConsecutiveUppercase: true,
+		})}`;
+		return discriminator(discriminatorPreimage);
+	}
+}

sdk/src/driftClient.ts

@@ -273,7 +273,8 @@ export class DriftClient {
 		this.program = new Program(
 			driftIDL as Idl,
 			config.programID ?? new PublicKey(DRIFT_PROGRAM_ID),
-			this.provider
+			this.provider,
+			config.coder
 		);
 
 		this.authority = config.authority ?? this.wallet.publicKey;

sdk/src/driftClientConfig.ts

@@ -12,6 +12,7 @@ import { DriftEnv } from './config';
 import { TxSender } from './tx/types';
 import { TxHandler, TxHandlerConfig } from './tx/txHandler';
 import { DelistedMarketSetting, GrpcConfigs } from './accounts/types';
+import { Coder } from '@coral-xyz/anchor';
 
 export type DriftClientConfig = {
 	connection: Connection;
@@ -41,6 +42,7 @@ export type DriftClientConfig = {
 	txHandlerConfig?: TxHandlerConfig;
 	delistedMarketSetting?: DelistedMarketSetting;
 	useHotWalletAdmin?: boolean;
+	coder?: Coder;
 };
 
 export type DriftClientSubscriptionConfig =

sdk/tests/accounts/customizedCadenceBulkAccountLoader.test.ts

@@ -54,8 +54,8 @@ describe('CustomizedCadenceBulkAccountLoader', () => {
 		const customFrequency = 500;
 		const callback = () => {};
 
-		await loader.addAccount(pubkey, callback, customFrequency);
-		loader.removeAccount(pubkey);
+		const cid = await loader.addAccount(pubkey, callback, customFrequency);
+		loader.removeAccount(pubkey, cid);
 
 		expect(loader.getAccountCadence(pubkey)).to.equal(null);
 	});

sdk/tests/dlob/helpers.ts

@@ -146,6 +146,9 @@ export const mockAMM: AMM = {
 	netUnsettledFundingPnl: new BN(0),
 	quoteAssetAmountWithUnsettledLp: new BN(0),
 	referencePriceOffset: 0,
+
+	takerSpeedBumpOverride: 0,
+	ammSpreadAdjustment: 0,
 };
 
 export const mockPerpMarkets: Array<PerpMarketAccount> = [

test-scripts/single-anchor-test.sh

@@ -6,7 +6,7 @@ fi
 
 export ANCHOR_WALLET=~/.config/solana/id.json
 
-test_files=(adminDeposit.ts)
+test_files=(overwritePerpAccounts.ts)
 
 for test_file in ${test_files[@]}; do
   ts-mocha -t 300000 ./tests/${test_file}

tests/overwritePerpAccounts.ts

@@ -0,0 +1,152 @@
+import * as anchor from '@coral-xyz/anchor';
+import { expect } from 'chai';
+import { Program } from '@coral-xyz/anchor';
+import { Keypair, LAMPORTS_PER_SOL, PublicKey } from '@solana/web3.js';
+import {
+	BN,
+	TestClient,
+	QUOTE_PRECISION,
+	PRICE_PRECISION,
+	PEG_PRECISION,
+	OracleSource,
+} from '../sdk/src';
+import {
+	initializeQuoteSpotMarket,
+	mockUSDCMint,
+	mockUserUSDCAccount,
+	mockOracleNoProgram,
+	overWritePerpMarket,
+} from './testHelpers';
+import { startAnchor } from 'solana-bankrun';
+import { TestBulkAccountLoader } from '../sdk/src/accounts/testBulkAccountLoader';
+import { BankrunContextWrapper } from '../sdk/src/bankrun/bankrunConnection';
+import dotenv from 'dotenv';
+import {
+	CustomBorshAccountsCoder,
+	CustomBorshCoder,
+} from '../sdk/src/decode/customCoder';
+dotenv.config();
+
+describe('Bankrun Overwrite Accounts', () => {
+	const program = anchor.workspace.Drift as Program;
+	// @ts-ignore
+	program.coder.accounts = new CustomBorshAccountsCoder(program.idl);
+	let bankrunContextWrapper: BankrunContextWrapper;
+	let bulkAccountLoader: TestBulkAccountLoader;
+
+	let adminClient: TestClient;
+	let usdcMint: Keypair;
+	let spotMarketOracle: PublicKey;
+
+	const mantissaSqrtScale = new BN(Math.sqrt(PRICE_PRECISION.toNumber()));
+	const ammInitialQuoteAssetReserve = new anchor.BN(10 * 10 ** 13).mul(
+		mantissaSqrtScale
+	);
+	const ammInitialBaseAssetReserve = new anchor.BN(10 * 10 ** 13).mul(
+		mantissaSqrtScale
+	);
+
+	let userUSDCAccount: Keypair;
+
+	before(async () => {
+		const context = await startAnchor(
+			'',
+			[
+				{
+					name: 'serum_dex',
+					programId: new PublicKey(
+						'srmqPvymJeFKQ4zGQed1GFppgkRHL9kaELCbyksJtPX'
+					),
+				},
+			],
+			[]
+		);
+
+		// @ts-ignore
+		bankrunContextWrapper = new BankrunContextWrapper(context);
+
+		bulkAccountLoader = new TestBulkAccountLoader(
+			bankrunContextWrapper.connection,
+			'processed',
+			1
+		);
+
+		usdcMint = await mockUSDCMint(bankrunContextWrapper);
+		spotMarketOracle = await mockOracleNoProgram(bankrunContextWrapper, 200.1);
+
+		const keypair = new Keypair();
+		await bankrunContextWrapper.fundKeypair(keypair, 50 * LAMPORTS_PER_SOL);
+
+		adminClient = new TestClient({
+			connection: bankrunContextWrapper.connection.toConnection(),
+			wallet: new anchor.Wallet(keypair),
+			programID: program.programId,
+			opts: {
+				commitment: 'confirmed',
+			},
+			activeSubAccountId: 0,
+			subAccountIds: [],
+			perpMarketIndexes: [0, 1],
+			spotMarketIndexes: [0, 1, 2],
+			oracleInfos: [
+				{
+					publicKey: spotMarketOracle,
+					source: OracleSource.PYTH,
+				},
+			],
+			accountSubscription: {
+				type: 'polling',
+				accountLoader: bulkAccountLoader,
+			},
+			coder: new CustomBorshCoder(program.idl),
+		});
+		await adminClient.initialize(usdcMint.publicKey, true);
+		await adminClient.subscribe();
+		await initializeQuoteSpotMarket(adminClient, usdcMint.publicKey);
+
+		userUSDCAccount = await mockUserUSDCAccount(
+			usdcMint,
+			new BN(10).mul(QUOTE_PRECISION),
+			bankrunContextWrapper,
+			keypair.publicKey
+		);
+
+		await adminClient.initializeUserAccountAndDepositCollateral(
+			new BN(10).mul(QUOTE_PRECISION),
+			userUSDCAccount.publicKey
+		);
+
+		const periodicity = new BN(0);
+
+		await adminClient.initializePerpMarket(
+			0,
+			spotMarketOracle,
+			ammInitialBaseAssetReserve,
+			ammInitialQuoteAssetReserve,
+			periodicity,
+			new BN(224 * PEG_PRECISION.toNumber())
+		);
+	});
+
+	after(async () => {
+		await adminClient.unsubscribe();
+	});
+
+	it('should overwrite perp accounts', async () => {
+		const perpMarket = adminClient.getPerpMarketAccount(0);
+
+		const newBaseAssetAmount = new BN(123);
+		perpMarket.amm.baseAssetAmountWithAmm = newBaseAssetAmount;
+
+		await overWritePerpMarket(
+			adminClient,
+			bankrunContextWrapper,
+			perpMarket.pubkey,
+			perpMarket
+		);
+
+		const updatedPerpMarket = adminClient.getPerpMarketAccount(0);
+		expect(updatedPerpMarket.amm.baseAssetAmountWithAmm.eq(newBaseAssetAmount))
+			.to.be.true;
+	});
+});