feat: #1989: subaccount filter in ownedPositionIds

.changeset/lucky-snails-clean.md

@@ -0,0 +1,13 @@
+---
+'@penumbra-zone/storage': major
+'@penumbra-zone/protobuf': minor
+'@penumbra-zone/services': minor
+'@penumbra-zone/types': minor
+'@penumbra-zone/wasm': minor
+---
+
+- storage: add subaccount filter to `getOwnedPositionIds` method
+- protobuf: sync latest changes in penumbra protobufs
+- services: add subaccount filter to `ownedPositionIds` method in ViewService
+- types: change `isControlledAddress` wasm method to `getIndexByAddress`
+- wasm: change `isControlledAddress` method to `getIndexByAddress`

packages/protobuf/package.json

@@ -16,7 +16,7 @@
     "gen:ibc": "buf generate buf.build/cosmos/ibc:7ab44ae956a0488ea04e04511efa5f70",
     "gen:ics23": "buf generate buf.build/cosmos/ics23:55085f7c710a45f58fa09947208eb70b",
     "gen:noble": "buf generate buf.build/noble-assets/forwarding:5a8609a6772d417584a9c60cd8b80881",
-    "gen:penumbra": "buf generate buf.build/penumbra-zone/penumbra:0a56a4f32c244e7eb277e02f6e85afbd",
+    "gen:penumbra": "buf generate buf.build/penumbra-zone/penumbra:649f1d61327144cb9a7e15c7ad210dcb",
     "lint": "eslint src",
     "lint:fix": "eslint src --fix",
     "lint:strict": "tsc --noEmit && eslint src --max-warnings 0",

packages/services/src/view-service/owned-position-ids.ts

@@ -9,7 +9,8 @@ export const ownedPositionIds: Impl['ownedPositionIds'] = async function* (req,
   for await (const positionId of indexedDb.getOwnedPositionIds(
     req.positionState,
     req.tradingPair,
+    req.subaccount,
   )) {
-    yield { positionId: positionId };
+    yield { positionId: positionId, subaccount: req.subaccount };
   }
 };

packages/storage/src/indexed-db/index.ts

@@ -570,6 +570,7 @@ export class IndexedDb implements IndexedDbInterface {
   async *getOwnedPositionIds(
     positionState: PositionState | undefined,
     tradingPair: TradingPair | undefined,
+    subaccount: AddressIndex | undefined,
   ) {
     yield* new ReadableStream({
       start: async cont => {
@@ -578,7 +579,10 @@ export class IndexedDb implements IndexedDbInterface {
           const position = Position.fromJson(cursor.value.position);
           if (
             (!positionState || positionState.equals(position.state)) &&
-            (!tradingPair || tradingPair.equals(position.phi?.pair))
+            (!tradingPair || tradingPair.equals(position.phi?.pair)) &&
+            (!subaccount ||
+              (cursor.value.subaccount &&
+                subaccount.equals(AddressIndex.fromJson(cursor.value.subaccount))))
           ) {
             cont.enqueue(PositionId.fromJson(cursor.value.id));
           }
@@ -589,16 +593,25 @@ export class IndexedDb implements IndexedDbInterface {
     });
   }
 
-  async addPosition(positionId: PositionId, position: Position): Promise<void> {
+  async addPosition(
+    positionId: PositionId,
+    position: Position,
+    subaccount?: AddressIndex,
+  ): Promise<void> {
     assertPositionId(positionId);
     const positionRecord = {
       id: positionId.toJson() as Jsonified<PositionId>,
       position: position.toJson() as Jsonified<Position>,
+      subaccount: subaccount && (subaccount.toJson() as Jsonified<AddressIndex>),
     };
     await this.u.update({ table: 'POSITIONS', value: positionRecord });
   }
 
-  async updatePosition(positionId: PositionId, newState: PositionState): Promise<void> {
+  async updatePosition(
+    positionId: PositionId,
+    newState: PositionState,
+    subaccount?: AddressIndex,
+  ): Promise<void> {
     assertPositionId(positionId);
     const key = uint8ArrayToBase64(positionId.inner);
     const positionRecord = await this.db.get('POSITIONS', key);
@@ -615,6 +628,7 @@ export class IndexedDb implements IndexedDbInterface {
       value: {
         id: positionId.toJson() as Jsonified<PositionId>,
         position: position.toJson() as Jsonified<Position>,
+        subaccount: subaccount ? (subaccount.toJson() as Jsonified<AddressIndex>) : undefined,
       },
     });
   }

packages/storage/src/indexed-db/indexed-db.test-data.ts

@@ -11,6 +11,7 @@ import { Transaction } from '@penumbra-zone/protobuf/penumbra/core/transaction/v
 import type { ScanBlockResult } from '@penumbra-zone/types/state-commitment-tree';
 import { base64ToUint8Array } from '@penumbra-zone/types/base64';
 import { StateCommitment } from '@penumbra-zone/protobuf/penumbra/crypto/tct/v1/tct_pb';
+import { AddressIndex } from '@penumbra-zone/protobuf/penumbra/core/keys/v1/keys_pb';
 
 const hash3312332298 = base64ToUint8Array('JbOzRkf0VKm4eIM0DS27N5igX8jxvPhAMpBWSr2bj/Q=');
 
@@ -717,3 +718,6 @@ export const epoch3 = new Epoch({
   index: 3n,
   startHeight: 300n,
 });
+
+export const mainAccount = new AddressIndex({ account: 0 });
+export const firstSubaccount = new AddressIndex({ account: 1 });

packages/storage/src/indexed-db/indexed-db.test.ts

@@ -28,6 +28,8 @@ import {
   tradingPairGmGn,
   transaction,
   transactionId,
+  mainAccount,
+  firstSubaccount,
 } from './indexed-db.test-data.js';
 import { AddressIndex, WalletId } from '@penumbra-zone/protobuf/penumbra/core/keys/v1/keys_pb';
 import {
@@ -496,13 +498,13 @@ describe('IndexedDb', () => {
     it('position should be added and their state should change', async () => {
       const db = await IndexedDb.initialize({ ...generateInitialProps() });
 
-      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy);
+      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy, mainAccount);
       await db.updatePosition(
         positionIdGmPenumbraBuy,
         new PositionState({ state: PositionState_PositionStateEnum.CLOSED }),
       );
       const ownedPositions: PositionId[] = [];
-      for await (const positionId of db.getOwnedPositionIds(undefined, undefined)) {
+      for await (const positionId of db.getOwnedPositionIds(undefined, undefined, undefined)) {
         ownedPositions.push(positionId as PositionId);
       }
       expect(ownedPositions.length).toBe(1);
@@ -521,27 +523,28 @@ describe('IndexedDb', () => {
 
     it('should get all position ids', async () => {
       const db = await IndexedDb.initialize({ ...generateInitialProps() });
-      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy);
-      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell);
-      await db.addPosition(positionIdGmGnSell, positionGmGnSell);
+      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy, mainAccount);
+      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell, mainAccount);
+      await db.addPosition(positionIdGmGnSell, positionGmGnSell, firstSubaccount);
 
       const ownedPositions: PositionId[] = [];
-      for await (const positionId of db.getOwnedPositionIds(undefined, undefined)) {
+      for await (const positionId of db.getOwnedPositionIds(undefined, undefined, undefined)) {
         ownedPositions.push(positionId as PositionId);
       }
       expect(ownedPositions.length).toBe(3);
     });
 
     it('should get all position with given position state', async () => {
       const db = await IndexedDb.initialize({ ...generateInitialProps() });
-      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy);
-      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell);
-      await db.addPosition(positionIdGmGnSell, positionGmGnSell);
+      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy, mainAccount);
+      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell, mainAccount);
+      await db.addPosition(positionIdGmGnSell, positionGmGnSell, firstSubaccount);
 
       const ownedPositions: PositionId[] = [];
       for await (const positionId of db.getOwnedPositionIds(
         new PositionState({ state: PositionState_PositionStateEnum.CLOSED }),
         undefined,
+        undefined,
       )) {
         ownedPositions.push(positionId as PositionId);
       }
@@ -550,16 +553,33 @@ describe('IndexedDb', () => {
 
     it('should get all position with given trading pair', async () => {
       const db = await IndexedDb.initialize({ ...generateInitialProps() });
-      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy);
-      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell);
-      await db.addPosition(positionIdGmGnSell, positionGmGnSell);
+      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy, mainAccount);
+      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell, mainAccount);
+      await db.addPosition(positionIdGmGnSell, positionGmGnSell, firstSubaccount);
 
       const ownedPositions: PositionId[] = [];
-      for await (const positionId of db.getOwnedPositionIds(undefined, tradingPairGmGn)) {
+      for await (const positionId of db.getOwnedPositionIds(
+        undefined,
+        tradingPairGmGn,
+        undefined,
+      )) {
         ownedPositions.push(positionId as PositionId);
       }
       expect(ownedPositions.length).toBe(1);
     });
+
+    it('should get all position with given subaccount index', async () => {
+      const db = await IndexedDb.initialize({ ...generateInitialProps() });
+      await db.addPosition(positionIdGmPenumbraBuy, positionGmPenumbraBuy, mainAccount);
+      await db.addPosition(positionIdGnPenumbraSell, positionGnPenumbraSell, mainAccount);
+      await db.addPosition(positionIdGmGnSell, positionGmGnSell, firstSubaccount);
+
+      const ownedPositions: PositionId[] = [];
+      for await (const positionId of db.getOwnedPositionIds(undefined, undefined, mainAccount)) {
+        ownedPositions.push(positionId as PositionId);
+      }
+      expect(ownedPositions.length).toBe(2);
+    });
   });
 
   describe('prices', () => {

packages/types/src/indexed-db.ts

@@ -97,9 +97,14 @@ export interface IndexedDbInterface {
   getOwnedPositionIds(
     positionState: PositionState | undefined,
     tradingPair: TradingPair | undefined,
+    subaccount: AddressIndex | undefined,
   ): AsyncGenerator<PositionId, void>;
-  addPosition(positionId: PositionId, position: Position): Promise<void>;
-  updatePosition(positionId: PositionId, newState: PositionState): Promise<void>;
+  addPosition(positionId: PositionId, position: Position, subaccount?: AddressIndex): Promise<void>;
+  updatePosition(
+    positionId: PositionId,
+    newState: PositionState,
+    subaccount?: AddressIndex,
+  ): Promise<void>;
   addEpoch(startHeight: bigint): Promise<void>;
   getEpochByHeight(height: bigint): Promise<Epoch | undefined>;
   upsertValidatorInfo(validatorInfo: ValidatorInfo): Promise<void>;
@@ -297,6 +302,7 @@ export interface PenumbraDb extends DBSchema {
 export interface PositionRecord {
   id: Jsonified<PositionId>; // PositionId (must be JsonValue because ['id']['inner'] is a key )
   position: Jsonified<Position>; // Position
+  subaccount?: Jsonified<AddressIndex>; // Position AddressIndex
 }
 
 export type Tables = Record<string, StoreNames<PenumbraDb>>;

packages/types/src/servers.ts

@@ -1,7 +1,7 @@
 import { ScanBlockResult } from './state-commitment-tree.js';
 import { CompactBlock } from '@penumbra-zone/protobuf/penumbra/core/component/compact_block/v1/compact_block_pb';
 import { MerkleRoot } from '@penumbra-zone/protobuf/penumbra/crypto/tct/v1/tct_pb';
-import { Address } from '@penumbra-zone/protobuf/penumbra/core/keys/v1/keys_pb';
+import { Address, AddressIndex } from '@penumbra-zone/protobuf/penumbra/core/keys/v1/keys_pb';
 
 export interface ViewServerInterface {
   scanBlock(compactBlock: CompactBlock, skipTrialDecrypt: boolean): Promise<boolean>;
@@ -12,5 +12,5 @@ export interface ViewServerInterface {
 
   getSctRoot(): MerkleRoot;
 
-  isControlledAddress(address: Address): boolean;
+  getIndexByAddress(address: Address): AddressIndex | undefined;
 }

packages/wasm/crate/src/keys.rs

@@ -117,7 +117,7 @@ pub fn get_index_by_address(full_viewing_key: &[u8], address: &[u8]) -> WasmResu
 
     let address: Address = Address::decode(address)?;
     let fvk: FullViewingKey = FullViewingKey::decode(full_viewing_key)?;
-    let index: Option<pb::AddressIndex> = fvk.address_index(&address).map(Into::into);
+    let index: Option<pb::AddressIndex> = controlled_by_index(&fvk, &address).map(Into::into);
     let result = serde_wasm_bindgen::to_value(&index)?;
     Ok(result)
 }
@@ -129,11 +129,11 @@ pub fn is_controlled_address(full_viewing_key: &[u8], address: &[u8]) -> WasmRes
 
     let address: Address = Address::decode(address)?;
     let fvk: FullViewingKey = FullViewingKey::decode(full_viewing_key)?;
-    Ok(is_controlled_inner(&fvk, &address))
+    Ok(controlled_by_index(&fvk, &address).is_some())
 }
 
-pub fn is_controlled_inner(fvk: &FullViewingKey, address: &Address) -> bool {
-    fvk.address_index(address).is_some()
+pub fn controlled_by_index(fvk: &FullViewingKey, address: &Address) -> Option<AddressIndex> {
+    fvk.address_index(address)
 }
 
 #[wasm_bindgen(getter_with_clone)]

packages/wasm/crate/src/view_server.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 
 use indexed_db_futures::IdbDatabase;
 use penumbra_compact_block::{CompactBlock, StatePayload};
+use penumbra_keys::keys::AddressIndex;
 use penumbra_keys::{Address, FullViewingKey};
 use penumbra_proto::DomainType;
 use penumbra_sct::Nullifier;
@@ -16,7 +17,7 @@ use wasm_bindgen::prelude::wasm_bindgen;
 use wasm_bindgen::JsValue;
 
 use crate::error::WasmResult;
-use crate::keys::is_controlled_inner;
+use crate::keys::controlled_by_index;
 use crate::note_record::SpendableNoteRecord;
 use crate::storage::{init_idb_storage, Storage};
 use crate::swap_record::SwapRecord;
@@ -307,13 +308,17 @@ impl ViewServer {
         Ok(root.encode_to_vec())
     }
 
-    /// Checks if address is controlled by view server full viewing key
+    /// Get the address index by provided address
+    /// Returns: `Uint8Array representing inner Address`
     #[wasm_bindgen]
-    pub fn is_controlled_address(&self, address: &[u8]) -> WasmResult<bool> {
+    pub fn get_index_by_address(full_viewing_key: &[u8], address: &[u8]) -> WasmResult<JsValue> {
         utils::set_panic_hook();
 
         let address: Address = Address::decode(address)?;
-        Ok(is_controlled_inner(&self.fvk, &address))
+        let fvk: FullViewingKey = FullViewingKey::decode(full_viewing_key)?;
+        let index: Option<AddressIndex> = controlled_by_index(&fvk, &address).map(Into::into);
+        let result = serde_wasm_bindgen::to_value(&index)?;
+        Ok(result)
     }
 }
 

packages/wasm/src/view-server.ts

@@ -10,8 +10,12 @@ import {
 } from '@penumbra-zone/types/state-commitment-tree';
 import type { IdbConstants } from '@penumbra-zone/types/indexed-db';
 import type { ViewServerInterface } from '@penumbra-zone/types/servers';
-import { Address, FullViewingKey } from '@penumbra-zone/protobuf/penumbra/core/keys/v1/keys_pb';
-import { isControlledAddress } from './address.js';
+import {
+  Address,
+  AddressIndex,
+  FullViewingKey,
+} from '@penumbra-zone/protobuf/penumbra/core/keys/v1/keys_pb';
+import { getAddressIndexByAddress } from './address.js';
 
 declare global {
   // eslint-disable-next-line no-var -- TODO: explain
@@ -89,7 +93,7 @@ export class ViewServer implements ViewServerInterface {
     };
   }
 
-  isControlledAddress(address: Address): boolean {
-    return isControlledAddress(this.fullViewingKey, address);
+  getIndexByAddress(address: Address): AddressIndex | undefined {
+    return getAddressIndexByAddress(this.fullViewingKey, address);
   }
 }