Merge pull request #200 from badaitech/fix/context-keypair-race-condition

Fix race condition by saving promise

Author: yakud

packages/chaingraph-types/src/execution/__tests__/execution-context.test.ts

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2025 BadLabs
+ *
+ * Use of this software is governed by the Business Source License 1.1 included in the file LICENSE.txt.
+ *
+ * As of the Change Date specified in that file, in accordance with the Business Source License, use of this software will be governed by the Apache License, version 2.0.
+ */
+
+import { subtle } from 'node:crypto'
+import { setTimeout as sleep } from 'node:timers/promises'
+import { describe, expect, it, vi } from 'vitest'
+import { ExecutionContext } from '../execution-context'
+
+describe('execution context', () => {
+  it('has the same ECDH keypair after multiple calls', async () => {
+    const abortController = new AbortController()
+    const ctx = new ExecutionContext('', abortController)
+
+    const expected = await ctx.getECDHKeyPair()
+    const actual = await ctx.getECDHKeyPair()
+
+    expect(actual).toStrictEqual(expected)
+  })
+
+  it('has the same ECDH keypair for concurrent calls', async () => {
+    const originalGenerateKey = subtle.generateKey.bind(subtle)
+
+    vi.spyOn(subtle, 'generateKey').mockImplementation(async (...args) => {
+      // Introduce a delay to simulate async operation and increase the chance of race condition
+      await sleep(10)
+      return originalGenerateKey(...args)
+    })
+
+    try {
+      const abortController = new AbortController()
+      const ctx = new ExecutionContext('', abortController)
+
+      const results = await Promise.all([
+        ctx.getECDHKeyPair(),
+        ctx.getECDHKeyPair(),
+        ctx.getECDHKeyPair(),
+        ctx.getECDHKeyPair(),
+        ctx.getECDHKeyPair(),
+      ])
+
+      const expected = await subtle.exportKey('raw', results[0].publicKey)
+      for (let i = 1; i < results.length; i++) {
+        const actual = await subtle.exportKey('raw', results[i].publicKey)
+        expect(actual).toStrictEqual(expected)
+      }
+    } finally {
+      vi.spyOn(subtle, 'generateKey').mockRestore()
+    }
+  })
+})

packages/chaingraph-types/src/execution/execution-context.ts

@@ -34,7 +34,7 @@ export class ExecutionContext {
   public readonly metadata: Record<string, unknown>
   public readonly abortController: AbortController
 
-  private ecdhKeyPair?: CryptoKeyPair | null = null
+  private ecdhKeyPairPromise: Promise<CryptoKeyPair> | null = null
 
   // debug log events queue
   private readonly eventsQueue: EventQueue<ExecutionEvent>
@@ -96,19 +96,15 @@ export class ExecutionContext {
     return this.abortController.signal
   }
 
-  private async generateECDHKeyPair() {
-    this.ecdhKeyPair = await subtle.generateKey({
-      name: 'ECDH',
-      namedCurve: 'P-256',
-    }, false, ['deriveKey'])
-  }
-
   async getECDHKeyPair(): Promise<CryptoKeyPair> {
-    if (this.ecdhKeyPair === null || this.ecdhKeyPair === undefined) {
-      await this.generateECDHKeyPair()
+    if (!this.ecdhKeyPairPromise) {
+      this.ecdhKeyPairPromise = subtle.generateKey({
+        name: 'ECDH',
+        namedCurve: 'P-256',
+      }, false, ['deriveKey'])
     }
 
-    return this.ecdhKeyPair!
+    return this.ecdhKeyPairPromise
   }
 
   async sendEvent(event: ExecutionEvent): Promise<void> {
@@ -167,7 +163,7 @@ export class ExecutionContext {
       eventData,
       true, // Is child execution
     )
-    ctx.ecdhKeyPair = this.ecdhKeyPair // Share key pair if needed
+    ctx.ecdhKeyPairPromise = this.ecdhKeyPairPromise // Share key pair promise
     return ctx
   }
 }