[feat]: Overhaul private key management (#27)

* feat: Overhaul the 'privateKey' loading mechanism * chore: Minor code refactor * chore: Undo the removal of the wallet service 'publicKey' property * Merge branch 'testnet' into jsanmi/overhaul-private-key-management
catalystdao · Jul 24, 2024 · 2736f9b · 2736f9b
1 parent 03a2e38
commit 2736f9b
Show file tree

Hide file tree

Showing 9 changed files with 161 additions and 23 deletions.
diff --git a/config.example.yaml b/config.example.yaml
@@ -1,6 +1,15 @@
 # Global Relayer configuration
 global:
-  privateKey: ''                        # The privateKey of the account that will be submitting the packet relays
+  # ! The 'privateKey' of the account that will be submitting the packet relays is by default
+  # ! loaded from the environment variable 'RELAYER_PRIVATE_KEY'. Alternatively, the privateKey
+  # ! may be specified here (not recommended).
+  # privateKey: ''
+  # ! Optionally, custom privateKey loaders may be implemented and specified (NOTE: the 'env' 
+  # ! loader is used if no privateKey configuration is specified):
+  # privateKey:
+  #   loader: 'env'                     # The privateKey loader name (must match the implementation on src/config/privateKeyLoaders/<loader>.ts).
+  #   customLoaderConfig: ''            # Custom loader configs may be specified.
+
   logLevel: 'info'
 
   monitor:

diff --git a/src/config/config.schema.ts b/src/config/config.schema.ts
@@ -4,8 +4,8 @@ import { AnyValidateFunction } from "ajv/dist/core"
 const MIN_PROCESSING_INTERVAL = 1;
 const MAX_PROCESSING_INTERVAL = 500;
 
-const EVM_ADDRESS_EXPR = '^0x[0-9a-fA-F]{40}$';  // '0x' + 20 bytes (40 chars)
-const BYTES_32_HEX_EXPR = '^0x[0-9a-fA-F]{64}$';  // '0x' + 32 bytes (64 chars)
+export const EVM_ADDRESS_EXPR = '^0x[0-9a-fA-F]{40}$';  // '0x' + 20 bytes (40 chars)
+export const BYTES_32_HEX_EXPR = '^0x[0-9a-fA-F]{64}$';  // '0x' + 32 bytes (64 chars)
 
 const POSITIVE_NUMBER_SCHEMA = {
     $id: "positive-number-schema",
@@ -73,10 +73,7 @@ const GLOBAL_SCHEMA = {
     $id: "global-schema",
     type: "object",
     properties: {
-        privateKey: {
-            type: "string",
-            pattern: BYTES_32_HEX_EXPR,
-        },
+        privateKey: { $ref: "private-key-schema" },
         logLevel: { $ref: "non-empty-string-schema" },
 
         monitor: { $ref: "monitor-schema" },
@@ -87,10 +84,28 @@ const GLOBAL_SCHEMA = {
         persister: { $ref: "persister-schema" },
         wallet: { $ref: "wallet-schema" },
     },
-    required: ["privateKey"],
+    required: [],
     additionalProperties: false
 }
 
+const PRIVATE_KEY_SCHEMA = {
+    $id: "private-key-schema",
+    "anyOf": [
+        {
+            type: "string",
+            pattern: BYTES_32_HEX_EXPR,
+        },
+        {
+            type: "object",
+            properties: {
+                loader: { $ref: "non-empty-string-schema" },
+            },
+            required: ["loader"],
+            additionalProperties: true,
+        }
+    ]
+}
+
 const MONITOR_SCHEMA = {
     $id: "monitor-schema",
     type: "object",
@@ -275,6 +290,7 @@ export function getConfigValidator(): AnyValidateFunction<unknown> {
     ajv.addSchema(PROCESSING_INTERVAL_SCHEMA);
     ajv.addSchema(CONFIG_SCHEMA);
     ajv.addSchema(GLOBAL_SCHEMA);
+    ajv.addSchema(PRIVATE_KEY_SCHEMA);
     ajv.addSchema(MONITOR_SCHEMA);
     ajv.addSchema(GETTER_SCHEMA);
     ajv.addSchema(PRICING_SCHEMA);

diff --git a/src/config/config.service.ts b/src/config/config.service.ts
@@ -3,8 +3,9 @@ import { readFileSync } from 'fs';
 import * as yaml from 'js-yaml';
 import dotenv from 'dotenv';
 import { PRICING_SCHEMA, getConfigValidator } from './config.schema';
-import { GlobalConfig, ChainConfig, AMBConfig, GetterGlobalConfig, SubmitterGlobalConfig, PersisterConfig, WalletGlobalConfig, GetterConfig, SubmitterConfig, WalletConfig, MonitorConfig, MonitorGlobalConfig, PricingConfig, PricingGlobalConfig, EvaluatorGlobalConfig, EvaluatorConfig } from './config.types';
+import { GlobalConfig, ChainConfig, AMBConfig, GetterGlobalConfig, SubmitterGlobalConfig, PersisterConfig, WalletGlobalConfig, GetterConfig, SubmitterConfig, WalletConfig, MonitorConfig, MonitorGlobalConfig, PricingGlobalConfig, EvaluatorGlobalConfig, PricingConfig, EvaluatorConfig } from './config.types';
 import { JsonRpcProvider } from 'ethers6';
+import { loadPrivateKeyLoader } from './privateKeyLoaders/privateKeyLoader';
 
 @Injectable()
 export class ConfigService {
@@ -85,6 +86,21 @@ export class ConfigService {
         }
     }
 
+    private async loadPrivateKey(rawPrivateKeyConfig: any): Promise<string> {
+        if (typeof rawPrivateKeyConfig === "string") {
+            //NOTE: Using 'console.warn' as the logger is not available at this point.  //TODO use logger
+            console.warn('WARNING: the privateKey has been loaded from the configuration file. Consider storing the privateKey using an alternative safer method.')
+            return rawPrivateKeyConfig;
+        }
+
+        const privateKeyLoader = loadPrivateKeyLoader(
+            rawPrivateKeyConfig?.['loader'] ?? null,
+            rawPrivateKeyConfig ?? {},
+        );
+
+        return privateKeyLoader.load();
+    }
+
     private loadGlobalConfig(): GlobalConfig {
         const rawGlobalConfig = this.rawConfig['global'];
 
@@ -96,7 +112,7 @@ export class ConfigService {
 
         return {
             port: parseInt(process.env['RELAYER_PORT']),
-            privateKey: rawGlobalConfig.privateKey,
+            privateKey: this.loadPrivateKey(rawGlobalConfig.privateKey),
             logLevel: rawGlobalConfig.logLevel,
             monitor: this.formatMonitorGlobalConfig(rawGlobalConfig.monitor),
             getter: this.formatGetterGlobalConfig(rawGlobalConfig.getter),

diff --git a/src/config/config.types.ts b/src/config/config.types.ts
@@ -1,6 +1,6 @@
 export interface GlobalConfig {
   port: number;
-  privateKey: string;
+  privateKey: Promise<string>;
   logLevel?: string;
   monitor: MonitorGlobalConfig;
   getter: GetterGlobalConfig;
@@ -11,6 +11,10 @@ export interface GlobalConfig {
   wallet: WalletGlobalConfig;
 }
 
+export type PrivateKeyConfig = string | {
+  loader: string;
+}
+
 export interface MonitorGlobalConfig {
   interval?: number;
   blockDelay?: number;

diff --git a/src/config/privateKeyLoaders/env.ts b/src/config/privateKeyLoaders/env.ts
@@ -0,0 +1,35 @@
+import { PrivateKeyLoader, PrivateKeyLoaderConfig } from "./privateKeyLoader";
+
+export const PRIVATE_KEY_LOADER_TYPE_ENVIRONMENT_VARIABLE = 'env';
+const DEFAULT_ENV_VARIABLE_NAME = 'RELAYER_PRIVATE_KEY';
+
+export interface EnvPrivateKeyLoaderConfig extends PrivateKeyLoaderConfig {
+    envVariableName?: string,
+}
+
+export class EnvPrivateKeyLoader extends PrivateKeyLoader {
+    override loaderType: string = PRIVATE_KEY_LOADER_TYPE_ENVIRONMENT_VARIABLE;
+    private readonly envVariableName: string;
+
+    constructor(
+        protected override readonly config: EnvPrivateKeyLoaderConfig,
+    ) {
+        super(config);
+
+        this.envVariableName = config.envVariableName ?? DEFAULT_ENV_VARIABLE_NAME;
+    }
+
+    override async loadPrivateKey(): Promise<string> {
+        const privateKey = process.env[this.envVariableName];
+
+        if (privateKey == undefined) {
+            throw new Error(
+                `Failed to load privateKey from enviornment variable '${this.envVariableName}'.`,
+            );
+        }
+
+        return privateKey;
+    }
+}
+
+export default EnvPrivateKeyLoader;
diff --git a/src/config/privateKeyLoaders/privateKeyLoader.ts b/src/config/privateKeyLoaders/privateKeyLoader.ts
@@ -0,0 +1,51 @@
+import { BYTES_32_HEX_EXPR } from "../config.schema";
+
+export const PRIVATE_KEY_LOADER_TYPE_BASE = 'base';
+
+const DEFAULT_PRIVATE_KEY_LOADER = 'env';
+
+export interface PrivateKeyLoaderConfig {
+}
+
+export function loadPrivateKeyLoader(
+    loader: string | null,
+    config: PrivateKeyLoaderConfig
+): BasePrivateKeyLoader {
+
+    // eslint-disable-next-line @typescript-eslint/no-var-requires
+    const module = require(`./${loader ?? DEFAULT_PRIVATE_KEY_LOADER}`);
+    const loaderClass: typeof BasePrivateKeyLoader = module.default;
+
+    return new loaderClass(
+        config,
+    )
+}
+
+export abstract class PrivateKeyLoader {
+    abstract readonly loaderType: string;
+
+    constructor(
+        protected readonly config: PrivateKeyLoaderConfig,
+    ) {}
+
+    abstract loadPrivateKey(): Promise<string>;
+
+    async load(): Promise<string> {
+        const privateKey = await this.loadPrivateKey();
+
+        if (!new RegExp(BYTES_32_HEX_EXPR).test(privateKey)) {
+            throw new Error('Invalid loaded privateKey format.')
+        }
+
+        return privateKey;
+    }
+}
+
+
+// ! 'BasePrivateKeyLoader' should only be used as a type.
+export class BasePrivateKeyLoader extends PrivateKeyLoader {
+    override loaderType: string = PRIVATE_KEY_LOADER_TYPE_BASE;
+    override loadPrivateKey(): Promise<string> {
+        throw new Error("Method not implemented.");
+    }
+}
diff --git a/src/main.ts b/src/main.ts
@@ -27,6 +27,9 @@ async function bootstrap() {
     const configService = app.get(ConfigService);
     const loggerService = app.get(LoggerService);
 
+    // Wait for the privateKey to be ready
+    await configService.globalConfig.privateKey;
+
     logLoadedOptions(configService, loggerService);
 
     await configService.isReady;

diff --git a/src/submitter/submitter.service.ts b/src/submitter/submitter.service.ts
@@ -63,7 +63,7 @@ export class SubmitterService {
     async onModuleInit(): Promise<void> {
         this.loggerService.info(`Starting the submitter on all chains...`);
 
-        const globalSubmitterConfig = this.loadGlobalSubmitterConfig();
+        const globalSubmitterConfig = await this.loadGlobalSubmitterConfig();
 
         // check if the submitter has been disabled.
         if (!globalSubmitterConfig.enabled) {
@@ -109,7 +109,7 @@ export class SubmitterService {
         await new Promise((r) => setTimeout(r, 5000));
     }
 
-    private loadGlobalSubmitterConfig(): GlobalSubmitterConfig {
+    private async loadGlobalSubmitterConfig(): Promise<GlobalSubmitterConfig> {
         const submitterConfig = this.configService.globalConfig.submitter;
 
         const enabled = submitterConfig['enabled'] ?? true;
@@ -129,7 +129,7 @@ export class SubmitterService {
         const maxEvaluationDuration =
             submitterConfig.maxEvaluationDuration ?? MAX_EVALUATION_DURATION_DEFAULT;
 
-        const walletPublicKey = (new Wallet(this.configService.globalConfig.privateKey)).address;
+        const walletPublicKey = (new Wallet(await this.configService.globalConfig.privateKey)).address;
 
         return {
             enabled,

diff --git a/src/wallet/wallet.service.ts b/src/wallet/wallet.service.ts
@@ -67,14 +67,14 @@ export class WalletService implements OnModuleInit {
 
     private readonly queuedMessages: Record<string, WalletServiceRoutingData[]> = {};
 
-    readonly publicKey: string;
+    readonly publicKey: Promise<string>;
 
     constructor(
         private readonly configService: ConfigService,
         private readonly loggerService: LoggerService,
     ) {
         this.defaultWorkerConfig = this.loadDefaultWorkerConfig();
-        this.publicKey = (new Wallet(this.configService.globalConfig.privateKey)).address;
+        this.publicKey = this.loadPublicKey();
     }
 
     async onModuleInit() {
@@ -85,10 +85,14 @@ export class WalletService implements OnModuleInit {
         this.initiateIntervalStatusLog();
     }
 
+    private async loadPublicKey(): Promise<string> {
+        return (new Wallet(await this.configService.globalConfig.privateKey)).address;
+    }
+
     private async initializeWorkers(): Promise<void> {
 
         for (const [chainId,] of this.configService.chainsConfig) {
-            this.spawnWorker(chainId);
+            await this.spawnWorker(chainId);
         }
 
         // Add a small delay to wait for the workers to be initialized
@@ -134,9 +138,9 @@ export class WalletService implements OnModuleInit {
         }
     }
 
-    private loadWorkerConfig(
+    private async loadWorkerConfig(
         chainId: string,
-    ): WalletWorkerData {
+    ): Promise<WalletWorkerData> {
 
         const defaultConfig = this.defaultWorkerConfig;
 
@@ -170,7 +174,7 @@ export class WalletService implements OnModuleInit {
                 chainWalletConfig.gasBalanceUpdateInterval ??
                 defaultConfig.balanceUpdateInterval,
 
-            privateKey: this.configService.globalConfig.privateKey,
+            privateKey: await this.configService.globalConfig.privateKey,
 
             maxFeePerGas:
                 chainWalletConfig.maxFeePerGas ??
@@ -200,10 +204,10 @@ export class WalletService implements OnModuleInit {
         };
     }
 
-    private spawnWorker(
+    private async spawnWorker(
         chainId: string
-    ): void {
-        const workerData = this.loadWorkerConfig(chainId);
+    ): Promise<void> {
+        const workerData = await this.loadWorkerConfig(chainId);
         this.loggerService.info(
             {
                 chainId,