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

* feat: Overhaul the 'privateKey' loading mechanism * chore: Remove duplicated 'blockDelay' configuration
catalystdao · Jul 24, 2024 · 6880909 · 6880909
1 parent 5395d37
commit 6880909
Show file tree

Hide file tree

Showing 10 changed files with 162 additions and 24 deletions.
diff --git a/config.example.yaml b/config.example.yaml
@@ -1,6 +1,15 @@
 # Global Underwriter configuration
 global:
-  privateKey: '' # The privateKey of the account that will be submitting the underwrites
+  # ! The 'privateKey' of the account that will be submitting the underwrites is by default
+  # ! loaded from the environment variable 'UNDERWRITER_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",
@@ -83,10 +83,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" },
@@ -95,10 +92,28 @@ const GLOBAL_SCHEMA = {
         expirer: { $ref: "expirer-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",
@@ -357,6 +372,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(LISTENER_SCHEMA);
     ajv.addSchema(RELAY_DELIVERY_COSTS_SCHEMA);

diff --git a/src/config/config.service.ts b/src/config/config.service.ts
@@ -4,7 +4,7 @@ import * as yaml from 'js-yaml';
 import dotenv from 'dotenv';
 import { getConfigValidator } from './config.schema';
 import { GlobalConfig, ChainConfig, AMBConfig, MonitorGlobalConfig, ListenerGlobalConfig, UnderwriterGlobalConfig, ExpirerGlobalConfig, WalletGlobalConfig, MonitorConfig, ListenerConfig, UnderwriterConfig, WalletConfig, ExpirerConfig, TokensConfig, EndpointConfig, VaultTemplateConfig, RelayDeliveryCosts } from './config.types';
-
+import { loadPrivateKeyLoader } from './privateKeyLoaders/privateKeyLoader';
 
 @Injectable()
 export class ConfigService {
@@ -80,6 +80,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'];
 
@@ -91,7 +106,7 @@ export class ConfigService {
 
         return {
             port: parseInt(process.env['UNDERWRITER_PORT']),
-            privateKey: rawGlobalConfig.privateKey,
+            privateKey: this.loadPrivateKey(rawGlobalConfig.privateKey),
             logLevel: rawGlobalConfig.logLevel,
             monitor: this.formatMonitorGlobalConfig(rawGlobalConfig.monitor),
             listener: this.formatListenerGlobalConfig(rawGlobalConfig.listener),

diff --git a/src/config/config.types.ts b/src/config/config.types.ts
@@ -1,7 +1,7 @@
 
 export interface GlobalConfig {
     port: number;
-    privateKey: string;
+    privateKey: Promise<string>;
     logLevel?: string;
     monitor: MonitorGlobalConfig;
     listener: ListenerGlobalConfig;
@@ -11,6 +11,11 @@ export interface GlobalConfig {
 }
 
 
+export type PrivateKeyConfig = string | {
+  loader: string;
+}
+
+
 export interface MonitorGlobalConfig {
     blockDelay?: number;
     retryInterval?: 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 = 'UNDERWRITER_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/expirer/expirer.service.ts b/src/expirer/expirer.service.ts
@@ -198,7 +198,7 @@ export class ExpirerService implements OnModuleInit {
                 ?? defaultConfig.expireBlocksMargin,
             minUnderwriteDuration,
 
-            underwriterPublicKey: this.walletService.publicKey,
+            underwriterPublicKey: await this.walletService.publicKey,
             monitorPort: await this.monitorService.attachToMonitor(chainId),
             walletPort: await this.walletService.attachToWallet(),
             loggerOptions: this.loggerService.loggerOptions

diff --git a/src/main.ts b/src/main.ts
@@ -26,6 +26,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;
+
     await logLoadedOptions(configService, loggerService);
 
     await app.listen(configService.globalConfig.port);

diff --git a/src/underwriter/underwriter.service.ts b/src/underwriter/underwriter.service.ts
@@ -107,7 +107,7 @@ export class UnderwriterService implements OnModuleInit {
     }
 
     private async initializeWorkers(): Promise<void> {
-        const defaultWorkerConfig = this.loadDefaultWorkerConfig();
+        const defaultWorkerConfig = await this.loadDefaultWorkerConfig();
 
         const ambs = Object.fromEntries(this.configService.ambsConfig.entries());
 
@@ -148,7 +148,7 @@ export class UnderwriterService implements OnModuleInit {
         }
     }
 
-    private loadDefaultWorkerConfig(): DefaultUnderwriterWorkerData {
+    private async loadDefaultWorkerConfig(): Promise<DefaultUnderwriterWorkerData> {
         const globalUnderwriterConfig = this.configService.globalConfig.underwriter;
 
         const enabled = globalUnderwriterConfig.enabled != false;
@@ -168,7 +168,7 @@ export class UnderwriterService implements OnModuleInit {
         const profitabilityFactor = globalUnderwriterConfig.profitabilityFactor ?? DEFAULT_UNDERWRITER_PROFITABILITY_FACTOR;
         const lowTokenBalanceWarning = globalUnderwriterConfig.lowTokenBalanceWarning;
         const tokenBalanceUpdateInterval = globalUnderwriterConfig.tokenBalanceUpdateInterval ?? DEFAULT_UNDERWRITER_TOKEN_BALANCE_UPDATE_INTERVAL;
-        const walletPublicKey = (new Wallet(this.configService.globalConfig.privateKey)).address;
+        const walletPublicKey = (new Wallet(await this.configService.globalConfig.privateKey)).address;
 
         const relayDeliveryCosts: RelayDeliveryCosts = globalUnderwriterConfig.relayDeliveryCosts ?? {
             gasUsage: DEFAULT_UNDERWRITER_RELAY_DELIVERY_GAS_USAGE

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,