Merge pull request #177 from ar-io/PE-6516-ant-registry

feat(PE-6516): ant registry class

src/common/ant-registry.ts

@@ -0,0 +1,118 @@
+/**
+ * Copyright (C) 2022-2024 Permanent Data Solutions, Inc. All Rights Reserved.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+import { ANT_REGISTRY_ID } from '../constants.js';
+import {
+  AoANTRegistryRead,
+  AoANTRegistryWrite,
+  AoMessageResult,
+  AoSigner,
+  OptionalSigner,
+  ProcessConfiguration,
+  WithSigner,
+  isProcessConfiguration,
+  isProcessIdConfiguration,
+} from '../types.js';
+import { createAoSigner } from '../utils/ao.js';
+import { AOProcess, InvalidContractConfigurationError } from './index.js';
+
+export class ANTRegistry {
+  static init(): AoANTRegistryRead;
+  static init(
+    config: Required<ProcessConfiguration> & { signer?: undefined },
+  ): AoANTRegistryRead;
+  static init({
+    signer,
+    ...config
+  }: WithSigner<Required<ProcessConfiguration>>): AoANTRegistryRead;
+  static init(
+    config?: OptionalSigner<ProcessConfiguration>,
+  ): AoANTRegistryRead | AoANTRegistryWrite {
+    if (config && config.signer) {
+      const { signer, ...rest } = config;
+      return new AoANTRegistryWriteable({
+        ...rest,
+        signer,
+      });
+    }
+    return new AoANTRegistryReadable(config);
+  }
+}
+
+export class AoANTRegistryReadable implements AoANTRegistryRead {
+  protected process: AOProcess;
+
+  constructor(config?: ProcessConfiguration) {
+    if (
+      config &&
+      (isProcessIdConfiguration(config) || isProcessConfiguration(config))
+    ) {
+      if (isProcessConfiguration(config)) {
+        this.process = config.process;
+      } else if (isProcessIdConfiguration(config)) {
+        this.process = new AOProcess({
+          processId: config.processId,
+        });
+      } else {
+        throw new InvalidContractConfigurationError();
+      }
+    } else {
+      this.process = new AOProcess({
+        processId: ANT_REGISTRY_ID,
+      });
+    }
+  }
+
+  // Should we rename this to "getANTsByAddress"? seems more clear, though not same as handler name
+  async accessControlList({
+    address,
+  }: {
+    address: string;
+  }): Promise<{ Owned: string[]; Controlled: string[] }> {
+    return this.process.read({
+      tags: [
+        { name: 'Action', value: 'Access-Control-List' },
+        { name: 'Address', value: address },
+      ],
+    });
+  }
+}
+
+export class AoANTRegistryWriteable
+  extends AoANTRegistryReadable
+  implements AoANTRegistryWrite
+{
+  private signer: AoSigner;
+
+  constructor({ signer, ...config }: WithSigner<ProcessConfiguration>) {
+    super(config);
+    this.signer = createAoSigner(signer);
+  }
+
+  async register({
+    processId,
+  }: {
+    processId: string;
+  }): Promise<AoMessageResult> {
+    return this.process.send({
+      tags: [
+        { name: 'Action', value: 'Register' },
+        { name: 'Process-Id', value: processId },
+      ],
+      signer: this.signer,
+    });
+  }
+}

src/common/index.ts

@@ -17,6 +17,7 @@
 export * from './error.js';
 export * from './logger.js';
 export * from './ant.js';
+export * from './ant-registry.js';
 
 // ao
 export * from './io.js';

src/constants.ts

@@ -37,6 +37,8 @@ export const IO_DEVNET_PROCESS_ID =
 export const ioDevnetProcessId = IO_DEVNET_PROCESS_ID;
 export const IO_TESTNET_PROCESS_ID =
   'agYcCFJtrMG6cqMuZfskIkFTGvUPddICmtQSBIoPdiA';
+
+export const ANT_REGISTRY_ID = 'i_le_yKKPVstLTDSmkHRqf-wYphMnwB9OhleiTgMkWc';
 export const MIO_PER_IO = 1_000_000;
 export const AOS_MODULE_ID = 'cbn0KKrBZH7hdNkNokuXLtGryrWM--PjSTBqIzw9Kkk';
 export const ANT_LUA_ID = 'Flwio4Lr08g6s6uim6lEJNnVGD9ylvz0_aafvpiL8FI';

src/io.ts

@@ -325,6 +325,16 @@ export interface AoANTWrite extends AoANTRead {
   setName({ name }): Promise<AoMessageResult>;
 }
 
+export interface AoANTRegistryRead {
+  accessControlList(params: {
+    address: string;
+  }): Promise<{ Owned: string[]; Controlled: string[] }>;
+}
+
+export interface AoANTRegistryWrite extends AoANTRegistryRead {
+  register(params: { processId: string }): Promise<AoMessageResult>;
+}
+
 // AO Contract types
 export interface AoIOState {
   GatewayRegistry: Record<WalletAddress, AoGateway>;

src/utils/ao.ts

@@ -21,6 +21,7 @@ import { defaultArweave } from '../common/arweave.js';
 import { AOProcess } from '../common/index.js';
 import {
   ANT_LUA_ID,
+  ANT_REGISTRY_ID,
   AOS_MODULE_ID,
   DEFAULT_SCHEDULER_ID,
 } from '../constants.js';
@@ -35,6 +36,7 @@ export async function spawnANT({
   scheduler = DEFAULT_SCHEDULER_ID,
   state,
   stateContractTxId,
+  antRegistryId = ANT_REGISTRY_ID,
 }: {
   signer: AoSigner;
   module?: string;
@@ -43,6 +45,7 @@ export async function spawnANT({
   scheduler?: string;
   state?: ANTState;
   stateContractTxId?: string;
+  antRegistryId?: string;
 }): Promise<string> {
   //TODO: cache locally and only fetch if not cached
   const luaString = (await defaultArweave.transactions.getData(luaCodeTxId, {
@@ -54,6 +57,12 @@ export async function spawnANT({
     module,
     scheduler,
     signer,
+    tags: [
+      {
+        name: 'ANT-Registry-Id',
+        value: antRegistryId,
+      },
+    ],
   });
 
   const aosClient = new AOProcess({
@@ -124,7 +133,7 @@ export async function evolveANT({
 
 export function createAoSigner(signer: ContractSigner): AoSigner {
   if (!('publicKey' in signer)) {
-    return createDataItemSigner(signer) as any;
+    return createDataItemSigner(signer) as AoSigner;
   }
 
   const aoSigner = async ({ data, tags, target, anchor }) => {

src/utils/json.ts

@@ -14,6 +14,8 @@
  * You should have received a copy of the GNU Affero General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function safeDecode<R = unknown>(data: any): R {
   try {
     return JSON.parse(data);

src/utils/processes.ts

@@ -17,59 +17,32 @@
 import { EventEmitter } from 'eventemitter3';
 import { pLimit } from 'plimit-lit';
 
+import { ANTRegistry } from '../common/ant-registry.js';
 import { ANT } from '../common/ant.js';
 import { IO } from '../common/io.js';
 import { ILogger, Logger } from '../common/logger.js';
 import { IO_TESTNET_PROCESS_ID } from '../constants.js';
 import {
+  AoANTRegistryRead,
   AoANTState,
   AoArNSNameData,
   AoIORead,
   ProcessId,
   WalletAddress,
 } from '../types.js';
 
+/**
+ * @beta This API is in beta and may change in the future.
+ */
 export const getANTProcessesOwnedByWallet = async ({
   address,
-  contract = IO.init({
-    processId: IO_TESTNET_PROCESS_ID,
-  }),
+  registry = ANTRegistry.init(),
 }: {
   address: WalletAddress;
-  contract?: AoIORead;
+  registry?: AoANTRegistryRead;
 }): Promise<ProcessId[]> => {
-  const throttle = pLimit(50);
-  // get the record names of the registry - TODO: this may need to be paginated
-  const records: Record<string, AoArNSNameData> = await fetchAllArNSRecords({
-    contract: contract,
-  });
-  const uniqueContractProcessIds = Object.values(records)
-    .filter((record) => record.processId !== undefined)
-    .map((record) => record.processId);
-
-  // check the contract owner and controllers
-  const ownedOrControlledByWallet = await Promise.all(
-    uniqueContractProcessIds.map(async (processId) =>
-      throttle(async () => {
-        const ant = ANT.init({
-          processId,
-        });
-        const { Owner, Controllers } = await ant.getState();
-
-        if (Owner === address || Controllers.includes(address)) {
-          return processId;
-        }
-        return;
-      }),
-    ),
-  );
-
-  if (ownedOrControlledByWallet.length === 0) {
-    return [];
-  }
-
-  // TODO: insert gql query to find ANT processes owned by wallet given wallet not currently in the registry
-  return [...new Set(ownedOrControlledByWallet)] as string[];
+  const res = await registry.accessControlList({ address });
+  return [...new Set([...res.Owned, ...res.Controlled])];
 };
 
 function timeout(ms: number, promise) {
@@ -118,9 +91,11 @@ export class ArNSEventEmitter extends EventEmitter {
   async fetchProcessesOwnedByWallet({
     address,
     pageSize,
+    antRegistry = ANTRegistry.init(),
   }: {
     address: WalletAddress;
     pageSize?: number;
+    antRegistry?: AoANTRegistryRead;
   }) {
     const uniqueContractProcessIds: Record<
       string,
@@ -129,7 +104,8 @@ export class ArNSEventEmitter extends EventEmitter {
         names: Record<string, AoArNSNameData>;
       }
     > = {};
-
+    const antIdRes = await antRegistry.accessControlList({ address });
+    const antIds = new Set([...antIdRes.Owned, ...antIdRes.Controlled]);
     await timeout(
       this.timeoutMs,
       fetchAllArNSRecords({ contract: this.contract, emitter: this, pageSize }),
@@ -142,19 +118,18 @@ export class ArNSEventEmitter extends EventEmitter {
         });
         return {};
       })
-      .then((records) => {
-        if (!records) return;
-        Object.entries(records).forEach(([name, record]) => {
-          if (record.processId === undefined) {
-            return;
-          }
-          if (uniqueContractProcessIds[record.processId] === undefined) {
-            uniqueContractProcessIds[record.processId] = {
-              state: undefined,
-              names: {},
-            };
+      .then((records: Record<string, AoArNSNameData>) => {
+        Object.entries(records).forEach(([name, arnsRecord]) => {
+          if (antIds.has(arnsRecord.processId)) {
+            if (uniqueContractProcessIds[arnsRecord.processId] == undefined) {
+              uniqueContractProcessIds[arnsRecord.processId] = {
+                state: undefined,
+                names: {},
+              };
+            }
+            uniqueContractProcessIds[arnsRecord.processId].names[name] =
+              arnsRecord;
           }
-          uniqueContractProcessIds[record.processId].names[name] = record;
         });
       });
 
@@ -220,6 +195,7 @@ export const fetchAllArNSRecords = async ({
   do {
     const pageResult = await contract
       .getArNSRecords({ cursor, limit: pageSize })
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
       .catch((e: any) => {
         logger?.error(`Error getting ArNS records`, {
           message: e?.message,

tests/e2e/cjs/index.test.js

@@ -4,7 +4,7 @@ const assert = require('node:assert/strict');
  * Ensure that npm link has been ran prior to running these tests
  * (simply running npm run test:integration will ensure npm link is ran)
  */
-const { IO, ioDevnetProcessId } = require('@ar.io/sdk');
+const { IO, ioDevnetProcessId, ANTRegistry } = require('@ar.io/sdk');
 
 const io = IO.init({
   processId: ioDevnetProcessId,
@@ -266,3 +266,14 @@ describe('IO', async () => {
     assert.ok(tokenCost);
   });
 });
+
+describe('ANTRegistry', async () => {
+  const registry = ANTRegistry.init();
+  const address = '7waR8v4STuwPnTck1zFVkQqJh5K9q9Zik4Y5-5dV7nk';
+
+  it('should retrieve ids from registry', async () => {
+    const affiliatedAnts = await registry.accessControlList({ address });
+    assert(Array.isArray(affiliatedAnts.Owned));
+    assert(Array.isArray(affiliatedAnts.Controlled));
+  });
+});

tests/e2e/esm/index.test.js

@@ -1,4 +1,4 @@
-import { IO, ioDevnetProcessId } from '@ar.io/sdk';
+import { ANTRegistry, IO, ioDevnetProcessId } from '@ar.io/sdk';
 import { strict as assert } from 'node:assert';
 import { describe, it } from 'node:test';
 
@@ -267,3 +267,14 @@ describe('IO', async () => {
     assert.ok(tokenCost);
   });
 });
+
+describe('ANTRegistry', async () => {
+  const registry = ANTRegistry.init();
+  const address = '7waR8v4STuwPnTck1zFVkQqJh5K9q9Zik4Y5-5dV7nk';
+
+  it('should retrieve ids from registry', async () => {
+    const affiliatedAnts = await registry.accessControlList({ address });
+    assert(Array.isArray(affiliatedAnts.Owned));
+    assert(Array.isArray(affiliatedAnts.Controlled));
+  });
+});