fix(mitm): store certificate key in network.db

ulixee · Jan 21, 2022 · 749bf1e · 749bf1e
1 parent 395f7dd
commit 749bf1e
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 20 deletions.
diff --git a/mitm-socket/lib/CertificateGenerator.ts b/mitm-socket/lib/CertificateGenerator.ts
@@ -12,7 +12,7 @@ export default class CertificateGenerator extends BaseIpcHandler {
 
   private pendingCertsById = new Map<number, Resolvable<{ cert: string; expireDate: number }>>();
 
-  private privateKey: string;
+  private privateKey: Buffer;
   private waitForInit = new Resolvable<void>();
   private hasWaitForInitListeners = false;
 
@@ -26,12 +26,9 @@ export default class CertificateGenerator extends BaseIpcHandler {
     super({ ...options, mode: 'certs' });
   }
 
-  public async getPrivateKey(): Promise<string> {
-    await this.waitForInit;
-    return this.privateKey;
-  }
-
-  public async generateCerts(host: string): Promise<{ cert: string; expireDate: number }> {
+  public async generateCerts(
+    host: string,
+  ): Promise<{ cert: Buffer; expireDate: number; key: Buffer }> {
     await this.waitForConnected;
     certRequestId += 1;
     const id = certRequestId;
@@ -40,15 +37,16 @@ export default class CertificateGenerator extends BaseIpcHandler {
     this.pendingCertsById.set(id, resolvable);
 
     try {
+      await this.waitForInit;
       await this.sendIpcMessage({ id, host });
     } catch (error) {
       if (this.isClosing) return;
       throw error;
     }
 
     this.hasWaitForInitListeners = true;
-    await this.waitForInit;
-    return await resolvable.promise;
+    const { cert, expireDate } = await resolvable.promise;
+    return { cert: Buffer.from(cert), expireDate, key: this.privateKey };
   }
 
   protected onMessage(rawMessage: string): void {
@@ -65,7 +63,7 @@ export default class CertificateGenerator extends BaseIpcHandler {
     }
 
     if (message.status === 'init') {
-      this.privateKey = message.privateKey;
+      this.privateKey = Buffer.from(message.privateKey);
       this.waitForInit.resolve();
       return;
     }

diff --git a/mitm/lib/MitmProxy.ts b/mitm/lib/MitmProxy.ts
@@ -502,23 +502,22 @@ export default class MitmProxy {
     }
   }
 
-  private static async getCertificate(host: string): Promise<{ cert: string; key: string }> {
+  private static async getCertificate(host: string): Promise<{ cert: Buffer; key: Buffer }> {
     const { networkDb, certificateGenerator } = this;
 
     if (!certificateGenerator) return null;
 
     await certificateGenerator.waitForConnected;
-    const key = await certificateGenerator.getPrivateKey();
     const existing = networkDb.certificates.get(host);
     if (existing) {
       return {
-        key,
+        key: existing.key,
         cert: existing.pem,
       };
     }
     // if it doesn't exist, generate now
-    const { expireDate, cert } = await certificateGenerator.generateCerts(host);
-    networkDb.certificates.insert({ host, pem: cert, expireDate });
+    const { expireDate, cert, key } = await certificateGenerator.generateCerts(host);
+    networkDb.certificates.insert({ host, pem: cert, expireDate, key });
     return { key, cert };
   }
 

diff --git a/mitm/models/CertificatesTable.ts b/mitm/models/CertificatesTable.ts
@@ -7,10 +7,11 @@ export default class CertificatesTable extends SqliteTable<ICertificateRecord> {
   constructor(readonly db: SqliteDatabase) {
     super(
       db,
-      'CertificatesV2',
+      'CertificatesV3',
       [
         ['host', 'TEXT', 'NOT NULL PRIMARY KEY'],
-        ['pem', 'TEXT'],
+        ['key', 'BLOB'],
+        ['pem', 'BLOB'],
         ['expireDate', 'INTEGER'],
       ],
       true,
@@ -19,9 +20,9 @@ export default class CertificatesTable extends SqliteTable<ICertificateRecord> {
   }
 
   public insert(record: ICertificateRecord): void {
-    const { host, pem, expireDate } = record;
+    const { host, key, pem, expireDate } = record;
     this.pemByHost.set(host, record);
-    this.queuePendingInsert([host, pem, expireDate]);
+    this.queuePendingInsert([host, key, pem, expireDate]);
   }
 
   public get(host: string): ICertificateRecord {
@@ -43,6 +44,7 @@ export default class CertificatesTable extends SqliteTable<ICertificateRecord> {
 
 export interface ICertificateRecord {
   host: string;
-  pem: string;
+  key: Buffer;
+  pem: Buffer;
   expireDate: number;
 }