fix: onConnect retry

nodejs · Nov 12, 2024 · 4f8dff1 · 4f8dff1
1 parent 20b8026
commit 4f8dff1
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 17 deletions.
diff --git a/lib/api/api-connect.js b/lib/api/api-connect.js
@@ -28,6 +28,7 @@ class ConnectHandler extends AsyncResource {
     this.responseHeaders = responseHeaders || null
     this.callback = callback
     this.abort = null
+    this.upgraded = false
 
     addSignal(this, signal)
   }
@@ -49,6 +50,12 @@ class ConnectHandler extends AsyncResource {
   }
 
   onUpgrade (statusCode, rawHeaders, socket) {
+    if (this.upgraded) {
+      throw new Error('request already upgraded')
+    }
+
+    this.upgraded = true
+
     const { callback, opaque, context } = this
 
     removeSignal(this)

diff --git a/lib/api/api-pipeline.js b/lib/api/api-pipeline.js
@@ -93,6 +93,7 @@ class PipelineHandler extends AsyncResource {
     this.abort = null
     this.context = null
     this.onInfo = onInfo || null
+    this.headersSent = false
 
     this.req = new PipelineRequest().on('error', noop)
 
@@ -160,7 +161,16 @@ class PipelineHandler extends AsyncResource {
     this.context = context
   }
 
+  onUpgrade () {
+    throw new Error('bad upgrade')
+  }
+
   onHeaders (statusCode, rawHeaders, resume) {
+    if (this.headersSent) {
+      throw new Error('headers already sent')
+    }
+    this.headersSent = true
+
     const { opaque, handler, context } = this
 
     if (statusCode < 200) {

diff --git a/lib/api/api-request.js b/lib/api/api-request.js
@@ -58,6 +58,7 @@ class RequestHandler extends AsyncResource {
     this.highWaterMark = highWaterMark
     this.reason = null
     this.removeAbortListener = null
+    this.headersSent = false
 
     if (signal?.aborted) {
       this.reason = signal.reason ?? new RequestAbortedError()
@@ -85,7 +86,16 @@ class RequestHandler extends AsyncResource {
     this.context = context
   }
 
+  onUpgrade () {
+    throw new Error('bad upgrade')
+  }
+
   onHeaders (statusCode, rawHeaders, resume, statusMessage) {
+    if (this.headersSent) {
+      throw new Error('headers already sent')
+    }
+    this.headersSent = true
+
     const { callback, opaque, abort, context, responseHeaders, highWaterMark } = this
 
     const headers = responseHeaders === 'raw' ? util.parseRawHeaders(rawHeaders) : util.parseHeaders(rawHeaders)

diff --git a/lib/api/api-stream.js b/lib/api/api-stream.js
@@ -56,6 +56,7 @@ class StreamHandler extends AsyncResource {
     this.trailers = null
     this.body = body
     this.onInfo = onInfo || null
+    this.headersSent = false
 
     if (util.isStream(body)) {
       body.on('error', (err) => {
@@ -78,7 +79,16 @@ class StreamHandler extends AsyncResource {
     this.context = context
   }
 
+  onUpgrade () {
+    throw new Error('bad upgrade')
+  }
+
   onHeaders (statusCode, rawHeaders, resume, statusMessage) {
+    if (this.headersSent) {
+      throw new Error('headers already sent')
+    }
+    this.headersSent = true
+
     const { factory, opaque, context, responseHeaders } = this
 
     const headers = responseHeaders === 'raw' ? util.parseRawHeaders(rawHeaders) : util.parseHeaders(rawHeaders)

diff --git a/lib/api/api-upgrade.js b/lib/api/api-upgrade.js
@@ -29,6 +29,7 @@ class UpgradeHandler extends AsyncResource {
     this.callback = callback
     this.abort = null
     this.context = null
+    this.upgraded = false
 
     addSignal(this, signal)
   }
@@ -50,6 +51,11 @@ class UpgradeHandler extends AsyncResource {
   }
 
   onUpgrade (statusCode, rawHeaders, socket) {
+    if (this.upgraded) {
+      throw new Error('request already upgraded')
+    }
+    this.upgraded = true
+
     assert(statusCode === 101)
 
     const { callback, opaque, context } = this

diff --git a/lib/handler/retry-handler.js b/lib/handler/retry-handler.js
@@ -35,8 +35,6 @@ class RetryHandler {
     this.dispatch = handlers.dispatch
     this.handler = handlers.handler
     this.opts = { ...dispatchOpts, body: wrapRequestBody(opts.body) }
-    this.abort = null
-    this.aborted = false
     this.retryOpts = {
       retry: retryFn ?? RetryHandler[kRetryHandlerDefaultRetry],
       retryAfter: retryAfter ?? true,
@@ -62,22 +60,14 @@ class RetryHandler {
       ]
     }
 
+    this.abort = null
+    this.aborted = false
     this.retryCount = 0
     this.retryCountCheckpoint = 0
     this.start = 0
     this.end = null
     this.etag = null
     this.resume = null
-
-    // Handle possible onConnect duplication
-    this.handler.onConnect(reason => {
-      this.aborted = true
-      if (this.abort) {
-        this.abort(reason)
-      } else {
-        this.reason = reason
-      }
-    })
   }
 
   onRequestSent () {
@@ -93,11 +83,19 @@ class RetryHandler {
   }
 
   onConnect (abort) {
-    if (this.aborted) {
-      abort(this.reason)
-    } else {
-      this.abort = abort
-    }
+    this.abort = abort
+    this.aborted = false
+    this.retryCount = 0
+    this.retryCountCheckpoint = 0
+    this.start = 0
+    this.end = null
+    this.etag = null
+    this.resume = null
+
+    this.handler.onConnect(reason => {
+      this.aborted = true
+      this.abort(reason)
+    })
   }
 
   onBodySent (chunk) {