Skip to content

Commit 54a9b5c

Browse files
wphanwphan
authored
Wphan/custom coder (#1682)
* sdk: allow custom coder * remove unused accounts coder * linter * move customCoder into sdk, lint * update test helpers * update testhelpers.ts
1 parent 32e7915 commit 54a9b5c

File tree

8 files changed

+456
-5
lines changed

8 files changed

+456
-5
lines changed

‎sdk/src/decode/customCoder.ts‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import { Buffer } from 'buffer';
2+
import camelcase from 'camelcase';
3+
import { Idl, IdlTypeDef } from '@coral-xyz/anchor/dist/cjs/idl';
4+
import {
5+
AccountsCoder,
6+
BorshAccountsCoder,
7+
BorshEventCoder,
8+
BorshInstructionCoder,
9+
Coder,
10+
} from '@coral-xyz/anchor/dist/cjs/coder';
11+
import { BorshTypesCoder } from '@coral-xyz/anchor/dist/cjs/coder/borsh/types';
12+
import { discriminator } from '@coral-xyz/anchor/dist/cjs/coder/borsh/discriminator';
13+
14+
export class CustomBorshCoder<
15+
A extends string = string,
16+
T extends string = string,
17+
> implements Coder
18+
{
19+
readonly idl: Idl;
20+
21+
/**
22+
* Instruction coder.
23+
*/
24+
readonly instruction: BorshInstructionCoder;
25+
26+
/**
27+
* Account coder.
28+
*/
29+
readonly accounts: CustomBorshAccountsCoder<A>;
30+
31+
/**
32+
* Coder for events.
33+
*/
34+
readonly events: BorshEventCoder;
35+
36+
/**
37+
* Coder for user-defined types.
38+
*/
39+
readonly types: BorshTypesCoder<T>;
40+
41+
constructor(idl: Idl) {
42+
this.instruction = new BorshInstructionCoder(idl);
43+
this.accounts = new CustomBorshAccountsCoder(idl);
44+
this.events = new BorshEventCoder(idl);
45+
this.types = new BorshTypesCoder(idl);
46+
this.idl = idl;
47+
}
48+
}
49+
50+
/**
51+
* Custom accounts coder that wraps BorshAccountsCoder to fix encode buffer sizing.
52+
*/
53+
export class CustomBorshAccountsCoder<A extends string = string>
54+
implements AccountsCoder
55+
{
56+
private baseCoder: BorshAccountsCoder<A>;
57+
private idl: Idl;
58+
59+
public constructor(idl: Idl) {
60+
this.baseCoder = new BorshAccountsCoder<A>(idl);
61+
this.idl = idl;
62+
}
63+
64+
public async encode<T = any>(accountName: A, account: T): Promise<Buffer> {
65+
const idlAcc = this.idl.accounts?.find((acc) => acc.name === accountName);
66+
if (!idlAcc) {
67+
throw new Error(`Unknown account not found in idl: ${accountName}`);
68+
}
69+
70+
const buffer = Buffer.alloc(this.size(idlAcc)); // fix encode issue - use proper size instead of fixed 1000
71+
const layout = this.baseCoder['accountLayouts'].get(accountName);
72+
if (!layout) {
73+
throw new Error(`Unknown account: ${accountName}`);
74+
}
75+
const len = layout.encode(account, buffer);
76+
const accountData = buffer.slice(0, len);
77+
const discriminator = BorshAccountsCoder.accountDiscriminator(accountName);
78+
return Buffer.concat([discriminator, accountData]);
79+
}
80+
81+
// Delegate all other methods to the base coder
82+
public decode<T = any>(accountName: A, data: Buffer): T {
83+
return this.baseCoder.decode(accountName, data);
84+
}
85+
86+
public decodeAny<T = any>(data: Buffer): T {
87+
return this.baseCoder.decodeAny(data);
88+
}
89+
90+
public decodeUnchecked<T = any>(accountName: A, ix: Buffer): T {
91+
return this.baseCoder.decodeUnchecked(accountName, ix);
92+
}
93+
94+
public memcmp(accountName: A, appendData?: Buffer): any {
95+
return this.baseCoder.memcmp(accountName, appendData);
96+
}
97+
98+
public size(idlAccount: IdlTypeDef): number {
99+
return this.baseCoder.size(idlAccount);
100+
}
101+
102+
/**
103+
* Calculates and returns a unique 8 byte discriminator prepended to all anchor accounts.
104+
*
105+
* @param name The name of the account to calculate the discriminator.
106+
*/
107+
public static accountDiscriminator(name: string): Buffer {
108+
const discriminatorPreimage = `account:${camelcase(name, {
109+
pascalCase: true,
110+
preserveConsecutiveUppercase: true,
111+
})}`;
112+
return discriminator(discriminatorPreimage);
113+
}
114+
}

‎sdk/src/driftClient.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,8 @@ export class DriftClient {
273273
this.program = new Program(
274274
driftIDL as Idl,
275275
config.programID ?? new PublicKey(DRIFT_PROGRAM_ID),
276-
this.provider
276+
this.provider,
277+
config.coder
277278
);
278279

279280
this.authority = config.authority ?? this.wallet.publicKey;

‎sdk/src/driftClientConfig.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { DriftEnv } from './config';
1212
import { TxSender } from './tx/types';
1313
import { TxHandler, TxHandlerConfig } from './tx/txHandler';
1414
import { DelistedMarketSetting, GrpcConfigs } from './accounts/types';
15+
import { Coder } from '@coral-xyz/anchor';
1516

1617
export type DriftClientConfig = {
1718
connection: Connection;
@@ -41,6 +42,7 @@ export type DriftClientConfig = {
4142
txHandlerConfig?: TxHandlerConfig;
4243
delistedMarketSetting?: DelistedMarketSetting;
4344
useHotWalletAdmin?: boolean;
45+
coder?: Coder;
4446
};
4547

4648
export type DriftClientSubscriptionConfig =

‎sdk/tests/accounts/customizedCadenceBulkAccountLoader.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ describe('CustomizedCadenceBulkAccountLoader', () => {
5454
const customFrequency = 500;
5555
const callback = () => {};
5656

57-
await loader.addAccount(pubkey, callback, customFrequency);
58-
loader.removeAccount(pubkey);
57+
const cid = await loader.addAccount(pubkey, callback, customFrequency);
58+
loader.removeAccount(pubkey, cid);
5959

6060
expect(loader.getAccountCadence(pubkey)).to.equal(null);
6161
});

‎sdk/tests/dlob/helpers.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,9 @@ export const mockAMM: AMM = {
146146
netUnsettledFundingPnl: new BN(0),
147147
quoteAssetAmountWithUnsettledLp: new BN(0),
148148
referencePriceOffset: 0,
149+
150+
takerSpeedBumpOverride: 0,
151+
ammSpreadAdjustment: 0,
149152
};
150153

151154
export const mockPerpMarkets: Array<PerpMarketAccount> = [

‎test-scripts/single-anchor-test.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ fi
66

77
export ANCHOR_WALLET=~/.config/solana/id.json
88

9-
test_files=(adminDeposit.ts)
9+
test_files=(overwritePerpAccounts.ts)
1010

1111
for test_file in ${test_files[@]}; do
1212
ts-mocha -t 300000 ./tests/${test_file}

‎tests/overwritePerpAccounts.ts‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
import * as anchor from '@coral-xyz/anchor';
2+
import { expect } from 'chai';
3+
import { Program } from '@coral-xyz/anchor';
4+
import { Keypair, LAMPORTS_PER_SOL, PublicKey } from '@solana/web3.js';
5+
import {
6+
BN,
7+
TestClient,
8+
QUOTE_PRECISION,
9+
PRICE_PRECISION,
10+
PEG_PRECISION,
11+
OracleSource,
12+
} from '../sdk/src';
13+
import {
14+
initializeQuoteSpotMarket,
15+
mockUSDCMint,
16+
mockUserUSDCAccount,
17+
mockOracleNoProgram,
18+
overWritePerpMarket,
19+
} from './testHelpers';
20+
import { startAnchor } from 'solana-bankrun';
21+
import { TestBulkAccountLoader } from '../sdk/src/accounts/testBulkAccountLoader';
22+
import { BankrunContextWrapper } from '../sdk/src/bankrun/bankrunConnection';
23+
import dotenv from 'dotenv';
24+
import {
25+
CustomBorshAccountsCoder,
26+
CustomBorshCoder,
27+
} from '../sdk/src/decode/customCoder';
28+
dotenv.config();
29+
30+
describe('Bankrun Overwrite Accounts', () => {
31+
const program = anchor.workspace.Drift as Program;
32+
// @ts-ignore
33+
program.coder.accounts = new CustomBorshAccountsCoder(program.idl);
34+
let bankrunContextWrapper: BankrunContextWrapper;
35+
let bulkAccountLoader: TestBulkAccountLoader;
36+
37+
let adminClient: TestClient;
38+
let usdcMint: Keypair;
39+
let spotMarketOracle: PublicKey;
40+
41+
const mantissaSqrtScale = new BN(Math.sqrt(PRICE_PRECISION.toNumber()));
42+
const ammInitialQuoteAssetReserve = new anchor.BN(10 * 10 ** 13).mul(
43+
mantissaSqrtScale
44+
);
45+
const ammInitialBaseAssetReserve = new anchor.BN(10 * 10 ** 13).mul(
46+
mantissaSqrtScale
47+
);
48+
49+
let userUSDCAccount: Keypair;
50+
51+
before(async () => {
52+
const context = await startAnchor(
53+
'',
54+
[
55+
{
56+
name: 'serum_dex',
57+
programId: new PublicKey(
58+
'srmqPvymJeFKQ4zGQed1GFppgkRHL9kaELCbyksJtPX'
59+
),
60+
},
61+
],
62+
[]
63+
);
64+
65+
// @ts-ignore
66+
bankrunContextWrapper = new BankrunContextWrapper(context);
67+
68+
bulkAccountLoader = new TestBulkAccountLoader(
69+
bankrunContextWrapper.connection,
70+
'processed',
71+
1
72+
);
73+
74+
usdcMint = await mockUSDCMint(bankrunContextWrapper);
75+
spotMarketOracle = await mockOracleNoProgram(bankrunContextWrapper, 200.1);
76+
77+
const keypair = new Keypair();
78+
await bankrunContextWrapper.fundKeypair(keypair, 50 * LAMPORTS_PER_SOL);
79+
80+
adminClient = new TestClient({
81+
connection: bankrunContextWrapper.connection.toConnection(),
82+
wallet: new anchor.Wallet(keypair),
83+
programID: program.programId,
84+
opts: {
85+
commitment: 'confirmed',
86+
},
87+
activeSubAccountId: 0,
88+
subAccountIds: [],
89+
perpMarketIndexes: [0, 1],
90+
spotMarketIndexes: [0, 1, 2],
91+
oracleInfos: [
92+
{
93+
publicKey: spotMarketOracle,
94+
source: OracleSource.PYTH,
95+
},
96+
],
97+
accountSubscription: {
98+
type: 'polling',
99+
accountLoader: bulkAccountLoader,
100+
},
101+
coder: new CustomBorshCoder(program.idl),
102+
});
103+
await adminClient.initialize(usdcMint.publicKey, true);
104+
await adminClient.subscribe();
105+
await initializeQuoteSpotMarket(adminClient, usdcMint.publicKey);
106+
107+
userUSDCAccount = await mockUserUSDCAccount(
108+
usdcMint,
109+
new BN(10).mul(QUOTE_PRECISION),
110+
bankrunContextWrapper,
111+
keypair.publicKey
112+
);
113+
114+
await adminClient.initializeUserAccountAndDepositCollateral(
115+
new BN(10).mul(QUOTE_PRECISION),
116+
userUSDCAccount.publicKey
117+
);
118+
119+
const periodicity = new BN(0);
120+
121+
await adminClient.initializePerpMarket(
122+
0,
123+
spotMarketOracle,
124+
ammInitialBaseAssetReserve,
125+
ammInitialQuoteAssetReserve,
126+
periodicity,
127+
new BN(224 * PEG_PRECISION.toNumber())
128+
);
129+
});
130+
131+
after(async () => {
132+
await adminClient.unsubscribe();
133+
});
134+
135+
it('should overwrite perp accounts', async () => {
136+
const perpMarket = adminClient.getPerpMarketAccount(0);
137+
138+
const newBaseAssetAmount = new BN(123);
139+
perpMarket.amm.baseAssetAmountWithAmm = newBaseAssetAmount;
140+
141+
await overWritePerpMarket(
142+
adminClient,
143+
bankrunContextWrapper,
144+
perpMarket.pubkey,
145+
perpMarket
146+
);
147+
148+
const updatedPerpMarket = adminClient.getPerpMarketAccount(0);
149+
expect(updatedPerpMarket.amm.baseAssetAmountWithAmm.eq(newBaseAssetAmount))
150+
.to.be.true;
151+
});
152+
});

0 commit comments

Comments
 (0)