From a179b68e2dfd3667999c31b1dbd3870a0b92568c Mon Sep 17 00:00:00 2001
From: Josh Kelley <joshkel@gmail.com>
Date: Wed, 1 Nov 2023 11:11:33 -0400
Subject: [PATCH] fix: prevent one-time handlers from incorrectly marking
 themselves as used (#1822)

---
 src/core/handlers/RequestHandler.ts  |  13 ++-
 src/core/utils/handleRequest.test.ts | 134 +++++++++++++++++++++++++++
 2 files changed, 142 insertions(+), 5 deletions(-)

diff --git a/src/core/handlers/RequestHandler.ts b/src/core/handlers/RequestHandler.ts
index 917d6edc4..8f45a6e41 100644
--- a/src/core/handlers/RequestHandler.ts
+++ b/src/core/handlers/RequestHandler.ts
@@ -190,11 +190,6 @@ export abstract class RequestHandler<
     // and the response resolver so we can always read it for logging.
     const mainRequestRef = args.request.clone()
 
-    // Immediately mark the handler as used.
-    // Can't await the resolver to be resolved because it's potentially
-    // asynchronous, and there may be multiple requests hitting this handler.
-    this.isUsed = true
-
     const parsedResult = await this.parse({
       request: args.request,
       resolutionContext: args.resolutionContext,
@@ -209,6 +204,14 @@ export abstract class RequestHandler<
       return null
     }
 
+    // Re-check isUsed, in case another request hit this handler while we were
+    // asynchronously parsing the request.
+    if (this.isUsed && this.options?.once) {
+      return null
+    }
+
+    this.isUsed = true
+
     // Create a response extraction wrapper around the resolver
     // since it can be both an async function and a generator.
     const executeResolver = this.wrapResolver(this.resolver)
diff --git a/src/core/utils/handleRequest.test.ts b/src/core/utils/handleRequest.test.ts
index 2f49662ea..bff38a4a6 100644
--- a/src/core/utils/handleRequest.test.ts
+++ b/src/core/utils/handleRequest.test.ts
@@ -342,3 +342,137 @@ it('returns undefined without warning on a passthrough request', async () => {
   expect(callbacks.onPassthroughResponse).toHaveBeenNthCalledWith(1, request)
   expect(callbacks.onMockedResponse).not.toHaveBeenCalled()
 })
+
+it('marks the first matching one-time handler as used', async () => {
+  const { emitter } = setup()
+
+  const oneTimeHandler = http.get(
+    '/resource',
+    () => {
+      return HttpResponse.text('One-time')
+    },
+    { once: true },
+  )
+  const anotherHandler = http.get('/resource', () => {
+    return HttpResponse.text('Another')
+  })
+  const handlers: Array<RequestHandler> = [oneTimeHandler, anotherHandler]
+
+  const requestId = uuidv4()
+  const request = new Request('http://localhost/resource')
+  const firstResult = await handleRequest(
+    request,
+    requestId,
+    handlers,
+    options,
+    emitter,
+    callbacks,
+  )
+
+  expect(await firstResult?.text()).toBe('One-time')
+  expect(oneTimeHandler.isUsed).toBe(true)
+  expect(anotherHandler.isUsed).toBe(false)
+
+  const secondResult = await handleRequest(
+    request,
+    requestId,
+    handlers,
+    options,
+    emitter,
+    callbacks,
+  )
+
+  expect(await secondResult?.text()).toBe('Another')
+  expect(anotherHandler.isUsed).toBe(true)
+  expect(oneTimeHandler.isUsed).toBe(true)
+})
+
+it('does not mark non-matching one-time handlers as used', async () => {
+  const { emitter } = setup()
+
+  const oneTimeHandler = http.get(
+    '/resource',
+    () => {
+      return HttpResponse.text('One-time')
+    },
+    { once: true },
+  )
+  const anotherHandler = http.get(
+    '/another',
+    () => {
+      return HttpResponse.text('Another')
+    },
+    { once: true },
+  )
+  const handlers: Array<RequestHandler> = [oneTimeHandler, anotherHandler]
+
+  const requestId = uuidv4()
+  const firstResult = await handleRequest(
+    new Request('http://localhost/another'),
+    requestId,
+    handlers,
+    options,
+    emitter,
+    callbacks,
+  )
+
+  expect(await firstResult?.text()).toBe('Another')
+  expect(oneTimeHandler.isUsed).toBe(false)
+  expect(anotherHandler.isUsed).toBe(true)
+
+  const secondResult = await handleRequest(
+    new Request('http://localhost/resource'),
+    requestId,
+    handlers,
+    options,
+    emitter,
+    callbacks,
+  )
+
+  expect(await secondResult?.text()).toBe('One-time')
+  expect(anotherHandler.isUsed).toBe(true)
+  expect(oneTimeHandler.isUsed).toBe(true)
+})
+
+it('handles parallel requests with one-time handlers', async () => {
+  const { emitter } = setup()
+
+  const oneTimeHandler = http.get(
+    '/resource',
+    () => {
+      return HttpResponse.text('One-time')
+    },
+    { once: true },
+  )
+  const anotherHandler = http.get('/resource', () => {
+    return HttpResponse.text('Another')
+  })
+  const handlers: Array<RequestHandler> = [oneTimeHandler, anotherHandler]
+
+  const requestId = uuidv4()
+  const request = new Request('http://localhost/resource')
+  const firstResultPromise = handleRequest(
+    request,
+    requestId,
+    handlers,
+    options,
+    emitter,
+    callbacks,
+  )
+  const secondResultPromise = handleRequest(
+    request,
+    requestId,
+    handlers,
+    options,
+    emitter,
+    callbacks,
+  )
+
+  const firstResult = await firstResultPromise
+  const secondResult = await secondResultPromise
+
+  expect(await firstResult?.text()).toBe('One-time')
+  expect(await secondResult?.text()).toBe('Another')
+  expect(oneTimeHandler.isUsed).toBe(true)
+  expect(anotherHandler.isUsed).toBe(true)
+})