feat: Introduce whitelisted logic for prividium mode (#4190)

## What ❔
- Added `whitelisted_wallets` to the YAML configuration, allowing for
flexible wallet authorization.
- Updated the `Authorizer` class to handle the new `whitelisted_wallets`
property, supporting both 'all' and specific wallet addresses.
- Enhanced `YamlParser` to parse and validate the `whitelisted_wallets`
field.
- Implemented a new route in `usersRoutes` to check if a user's address
is whitelisted based on the updated logic.

## Why ❔

<!-- Why are these changes done? What goal do they contribute to? What
are the principles behind them? -->
<!-- The `Why` has to be clear to non-Matter Labs entities running their
own ZK Chain -->
<!-- Example: PR templates ensure PR reviewers, observers, and future
iterators are in context about the evolution of repos. -->

## Is this a breaking change?
- [ ] Yes
- [ ] No

## Operational changes
<!-- Any config changes? Any new flags? Any changes to any scripts? -->
<!-- Please add anything that non-Matter Labs entities running their own
ZK Chain may need to know -->

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zkstack dev fmt` and `zkstack dev
lint`.

Author: Jrigada

core/tests/ts-integration/src/private-rpc-permissions-editor.ts

@@ -10,7 +10,11 @@ type YamlMethod = {
 };
 
 type YamlContract = { address: string; methods: YamlMethod[] };
-type YamlRoot = { groups: any[]; contracts: YamlContract[] };
+type YamlRoot = {
+    whitelisted_wallets?: string | string[];
+    groups: any[];
+    contracts: YamlContract[];
+};
 const mutex = new Mutex();
 /**
  * Adds a method with public read/write permissions to the YAML file.
@@ -62,3 +66,29 @@ export async function injectPermissionsToFile(
     await fs.writeFile(filePath, dumped, 'utf8');
     mutex.release();
 }
+
+/**
+ * Sets the whitelisted wallets in the YAML file.
+ *
+ * @param filePath - absolute or relative path to the YAML file
+ * @param wallets - either "all" to allow all wallets, or an array of wallet addresses
+ */
+export async function setWhitelistedWallets(filePath: string, wallets: string | string[]): Promise<void> {
+    await mutex.acquire();
+    try {
+        // --- load & parse ---------------------------------------------------------
+        const raw = await fs.readFile(filePath, 'utf8');
+        const data = yaml.load(raw) as YamlRoot;
+
+        if (!data) throw new Error('Invalid YAML structure');
+
+        // --- set whitelisted wallets ----------------------------------------------
+        data.whitelisted_wallets = wallets;
+
+        // --- dump & save ----------------------------------------------------------
+        const dumped = yaml.dump(data, { lineWidth: 120, sortKeys: false, quotingType: '"' });
+        await fs.writeFile(filePath, dumped, 'utf8');
+    } finally {
+        mutex.release();
+    }
+}

core/tests/ts-integration/tests/prividium.test.ts

@@ -6,7 +6,7 @@ import { RetryableWallet } from '../src/retry-provider';
 import * as zksync from 'zksync-ethers';
 import { sleep } from 'utils';
 import { shouldLoadConfigFromFile } from 'utils/build/file-configs';
-import { injectPermissionsToFile } from '../src/private-rpc-permissions-editor';
+import { injectPermissionsToFile, setWhitelistedWallets } from '../src/private-rpc-permissions-editor';
 import path from 'path';
 import { shouldChangeTokenBalances, shouldOnlyTakeFee } from '../src/modifiers/balance-checker';
 import { Token } from '../src/types';
@@ -76,6 +76,12 @@ describe('Tests for the private rpc', () => {
         const runCommand = `zkstack private-rpc run --verbose --chain ${chainName}`;
 
         await executeCommandWithLogs(initCommand, await logsPath('private-rpc-init.log'));
+
+        // Set whitelisted wallets to "all" for integration tests
+        const pathToHome = path.join(__dirname, '../../../..');
+        const permissionsPath = path.join(pathToHome, `chains/${chainName}/configs/private-rpc-permissions.yaml`);
+        await setWhitelistedWallets(permissionsPath, 'all');
+
         executeCommandWithLogs(runCommand, await logsPath('private-rpc-run.log'));
 
         await waitForHealth(rpcUrl());

private-rpc/README.md

@@ -30,6 +30,33 @@ zkstack private-rpc reset-db
 The permissions can be modified by updating the `private-rpc-permissions.yaml` file. The exact path of this file will be
 printed during the init command.
 
+The file has two main sections: `whitelisted_wallets` and `contracts`.
+
+### Wallet Whitelisting
+
+The `whitelisted_wallets` key controls which wallet addresses are allowed to connect to the RPC. This is the first layer
+of security. It has two modes:
+
+1. **Allow all wallets**: To disable whitelisting and allow any address to connect, use the literal string `"all"`.
+
+   ```yaml
+   whitelisted_wallets: 'all'
+   ```
+
+2. **Allow specific wallets**: To restrict access, provide a list of authorized wallet addresses. The list cannot be
+   empty.
+
+   ```yaml
+   whitelisted_wallets:
+     - '0xAb5801a7D398351b8bE11C439e05C5B3259aeC9B'
+     - '0x...another address...'
+   ```
+
+### Contract & Method Permissions
+
+The `contracts` section allows for fine-grained control over which addresses can call specific methods on specific
+contracts. This acts as a second layer of security after the initial wallet whitelist check.
+
 ## Creating access tokens
 
 ```bash
@@ -41,6 +68,11 @@ curl -X POST http://localhost:4041/users \
          }'
 ```
 
+The server will respond with a JSON object:
+
+- `{"authorized":true}` if the address is allowed.
+- `{"authorized":false}` if the address is not on the whitelist.
+
 ## Using the rpc
 
 ```bash

private-rpc/example-permissions.yaml

@@ -1,3 +1,10 @@
+# To allow any wallet to connect, use the literal string "all".
+# Example: whitelisted_wallets: "all"
+
+# To restrict access to specific wallets, provide a non-empty list of addresses.
+whitelisted_wallets:
+  - "0x742d35Cc6634C0532925a3b8D69C7F16F6d34d2c"  # Example wallet address
+  
 groups:
   - name: "group1"
     members:

private-rpc/src/permissions/authorizer.ts

@@ -8,10 +8,12 @@ import { env } from '@/env';
 export class Authorizer {
     permissions: Map<string, AccessRule>;
     postReadFilters: Map<string, ResponseFilter>;
+    whitelistedWallets: Set<Address> | 'all';
 
     constructor() {
         this.permissions = new Map();
         this.postReadFilters = new Map();
+        this.whitelistedWallets = new Set();
     }
 
     addReadRule(address: Address, method: Hex, rule: AccessRule): void {
@@ -45,12 +47,27 @@ export class Authorizer {
         return this.postReadFilters.get(`${address}:${method}`) || null;
     }
 
+    isAddressWhitelisted(address: Address): boolean {
+        if (this.whitelistedWallets === 'all') {
+            return true;
+        }
+        return this.whitelistedWallets.has(address);
+    }
+
     reloadFromEnv(): Authorizer {
         const filePath = env.PERMISSIONS_YAML_PATH;
         console.log(`loading permissions from ${filePath}`);
         this.permissions = new Map();
         this.postReadFilters = new Map();
-        new YamlParser(filePath).load_rules(this);
+        const parser = new YamlParser(filePath);
+        parser.load_rules(this);
+
+        const wallets = parser.getWhitelistedWallets();
+        if (wallets === 'all') {
+            this.whitelistedWallets = 'all';
+        } else {
+            this.whitelistedWallets = new Set(wallets);
+        }
         return this;
     }
 }

private-rpc/src/permissions/yaml-parser.ts

@@ -18,7 +18,7 @@ import YAML from 'yaml';
 
 const publicSchema = z.object({ type: z.literal('public') });
 const closedSchema = z.object({ type: z.literal('closed') });
-const groupSchema = z.object({
+const groupSchemaRaw = z.object({
     type: z.literal('group'),
     groups: z.array(z.string())
 });
@@ -28,10 +28,10 @@ const checkArgumentSchema = z.object({
 });
 const oneOfSchema = z.object({
     type: z.literal('oneOf'),
-    rules: z.array(z.union([publicSchema, groupSchema, checkArgumentSchema]))
+    rules: z.array(z.union([publicSchema, groupSchemaRaw, checkArgumentSchema]))
 });
 
-const ruleSchema = z.union([publicSchema, closedSchema, groupSchema, checkArgumentSchema, oneOfSchema]);
+const ruleSchema = z.union([publicSchema, closedSchema, groupSchemaRaw, checkArgumentSchema, oneOfSchema]);
 type Rule = z.infer<typeof ruleSchema>;
 
 const methodSchema = z.object({
@@ -42,20 +42,26 @@ const methodSchema = z.object({
 });
 type RawMethod = z.infer<typeof methodSchema>;
 
+const groupSchema = z.object({
+    name: z.string(),
+    members: z.array(addressSchema)
+});
+type RawGroup = z.infer<typeof groupSchema>;
+
+const contractSchema = z.object({
+    address: addressSchema,
+    methods: z.array(methodSchema)
+});
+
 const yamlSchema = z.object({
-    groups: z.array(
-        z.object({
-            name: z.string(),
-            members: z.array(addressSchema)
-        })
-    ),
-
-    contracts: z.array(
-        z.object({
-            address: addressSchema,
-            methods: z.array(methodSchema)
-        })
-    )
+    whitelisted_wallets: z.union([
+        z.array(addressSchema).nonempty({
+            message: 'whitelisted_wallets array cannot be empty. To allow all wallets, use the literal "all".'
+        }),
+        z.literal('all')
+    ]),
+    groups: z.array(groupSchema),
+    contracts: z.array(contractSchema)
 });
 
 export class YamlParser {
@@ -66,7 +72,11 @@ export class YamlParser {
         const buf = readFileSync(filePath);
         const raw = YAML.parse(buf.toString());
         this.yaml = yamlSchema.parse(raw);
-        this.groups = this.yaml.groups.map(({ name, members }) => new Group(name, members));
+        this.groups = this.yaml.groups.map(({ name, members }: RawGroup) => new Group(name, members));
+    }
+
+    getWhitelistedWallets(): Address[] | 'all' {
+        return this.yaml.whitelisted_wallets;
     }
 
     private extractSelector(method: RawMethod): Hex {

private-rpc/src/routes/rpc-routes.ts

@@ -13,6 +13,11 @@ export function rpcRoutes(app: WebServer) {
             maybe.expect(new HttpError('Unauthorized', 401))
         );
 
+        // Check if user is still whitelisted before processing any RPC calls
+        if (!app.context.authorizer.isAddressWhitelisted(user.address)) {
+            throw new HttpError('Forbidden: Address not whitelisted', 403);
+        }
+
         const handler = new RpcCallHandler(allHandlers, {
             currentUser: user.address,
             targetRpcUrl: app.context.targetRpc,

private-rpc/src/routes/users-routes.ts

@@ -18,13 +18,19 @@ const createUserSchema = {
 export function usersRoutes(app: WebServer) {
     app.post('/', createUserSchema, async (req, reply) => {
         const { address, secret } = req.body;
+        const { authorizer, db, createTokenSecret } = app.context;
+
+        if (!authorizer.isAddressWhitelisted(address)) {
+            throw new HttpError('Forbidden: Address not whitelisted', 403);
+        }
+
         const token = nanoid(32);
 
-        if (secret !== app.context.createTokenSecret) {
+        if (secret !== createTokenSecret) {
             throw new HttpError('forbidden', 403);
         }
 
-        await app.context.db.transaction(async (tx) => {
+        await db.transaction(async (tx) => {
             await tx.delete(usersTable).where(eq(usersTable.address, address));
 
             await tx.insert(usersTable).values({
@@ -35,4 +41,18 @@ export function usersRoutes(app: WebServer) {
 
         return reply.send({ ok: true, token });
     });
+
+    const getUserSchema = {
+        schema: {
+            params: z.object({
+                address: addressSchema
+            })
+        }
+    };
+
+    app.get('/:address', getUserSchema, (req, reply) => {
+        const { authorizer } = app.context;
+        const authorized = authorizer.isAddressWhitelisted(req.params.address);
+        return reply.send({ authorized });
+    });
 }