From b60a0819db6c8a067f728007e9f2a59df1841778 Mon Sep 17 00:00:00 2001
From: KaKa <23028015+climba03003@users.noreply.github.com>
Date: Wed, 3 Jul 2024 22:36:37 +0800
Subject: [PATCH] fix: throw on retry when payload is consume by downstream
 (#3389)

* fix: throw on retry when payload is consumed

* fixup

* fixup

(cherry picked from commit 71121e5b67fc3dba79888269957198f639de8df4)
---
 lib/handler/retry-handler.js | 14 +++++++--
 test/issue-3356.js           | 57 ++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 test/issue-3356.js

diff --git a/lib/handler/retry-handler.js b/lib/handler/retry-handler.js
index f7dedfa4bac..e00e82faf08 100644
--- a/lib/handler/retry-handler.js
+++ b/lib/handler/retry-handler.js
@@ -192,8 +192,18 @@ class RetryHandler {
     if (this.resume != null) {
       this.resume = null
 
-      if (statusCode !== 206) {
-        return true
+      // Only Partial Content 206 supposed to provide Content-Range,
+      // any other status code that partially consumed the payload
+      // should not be retry because it would result in downstream
+      // wrongly concatanete multiple responses.
+      if (statusCode !== 206 && (this.start > 0 || statusCode !== 200)) {
+        this.abort(
+          new RequestRetryError('server does not support the range header and the payload was partially consumed', statusCode, {
+            headers,
+            data: { count: this.retryCount }
+          })
+        )
+        return false
       }
 
       const contentRange = parseRangeHeader(headers['content-range'])
diff --git a/test/issue-3356.js b/test/issue-3356.js
new file mode 100644
index 00000000000..927208583a9
--- /dev/null
+++ b/test/issue-3356.js
@@ -0,0 +1,57 @@
+'use strict'
+
+const { tspl } = require('@matteo.collina/tspl')
+const { test, after } = require('node:test')
+const { createServer } = require('node:http')
+const { once } = require('node:events')
+
+const { fetch, Agent, RetryAgent } = require('..')
+
+test('https://github.com/nodejs/undici/issues/3356', async (t) => {
+  t = tspl(t, { plan: 3 })
+
+  let shouldRetry = true
+  const server = createServer()
+  server.on('request', (req, res) => {
+    res.writeHead(200, { 'content-type': 'text/plain' })
+    if (shouldRetry) {
+      shouldRetry = false
+
+      res.flushHeaders()
+      res.write('h')
+      setTimeout(() => { res.end('ello world!') }, 100)
+    } else {
+      res.end('hello world!')
+    }
+  })
+
+  server.listen(0)
+
+  await once(server, 'listening')
+
+  after(async () => {
+    server.close()
+
+    await once(server, 'close')
+  })
+
+  const agent = new RetryAgent(new Agent({ bodyTimeout: 50 }), {
+    errorCodes: ['UND_ERR_BODY_TIMEOUT']
+  })
+
+  const response = await fetch(`http://localhost:${server.address().port}`, {
+    dispatcher: agent
+  })
+  setTimeout(async () => {
+    try {
+      t.equal(response.status, 200)
+      // consume response
+      await response.text()
+    } catch (err) {
+      t.equal(err.name, 'TypeError')
+      t.equal(err.cause.code, 'UND_ERR_REQ_RETRY')
+    }
+  }, 200)
+
+  await t.completed
+})