From 63bf370cc055bdfc2713500c58f539ef6bf18ce5 Mon Sep 17 00:00:00 2001
From: Jeremy Apthorp <nornagon@nornagon.net>
Date: Fri, 25 Jan 2019 10:57:26 -0800
Subject: [PATCH] fix: registerStreamProtocol callback with large chunks
 (#16532)

---
 atom/browser/net/url_request_stream_job.cc |  3 ++-
 spec/api-protocol-spec.js                  | 24 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/atom/browser/net/url_request_stream_job.cc b/atom/browser/net/url_request_stream_job.cc
index 4f2fe0524418..5fc4114d0029 100644
--- a/atom/browser/net/url_request_stream_job.cc
+++ b/atom/browser/net/url_request_stream_job.cc
@@ -141,6 +141,7 @@ void URLRequestStreamJob::StartAsync(
 }
 
 void URLRequestStreamJob::OnData(std::vector<char>&& buffer) {  // NOLINT
+  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
   if (write_buffer_.empty()) {
     // Quick branch without copying.
     write_buffer_ = std::move(buffer);
@@ -175,7 +176,7 @@ void URLRequestStreamJob::OnError(int error) {
 int URLRequestStreamJob::ReadRawData(net::IOBuffer* dest, int dest_size) {
   response_start_time_ = base::TimeTicks::Now();
 
-  if (ended_)
+  if (ended_ && write_buffer_.empty())
     return 0;
 
   // When write_buffer_ is empty, there is no data valable yet, we have to save
diff --git a/spec/api-protocol-spec.js b/spec/api-protocol-spec.js
index cedc87380c82..f78b6235deb0 100644
--- a/spec/api-protocol-spec.js
+++ b/spec/api-protocol-spec.js
@@ -654,6 +654,30 @@ describe('protocol module', () => {
         })
       })
     })
+
+    it('can handle large responses', async () => {
+      const data = Buffer.alloc(128 * 1024)
+      const handler = (request, callback) => {
+        callback(getStream(data.length, data))
+      }
+      await new Promise((resolve, reject) => {
+        protocol.registerStreamProtocol(protocolName, handler, err => {
+          if (err) return reject(err)
+          resolve()
+        })
+      })
+      const r = await new Promise((resolve, reject) => {
+        $.ajax({
+          url: protocolName + '://fake-host',
+          cache: false,
+          success: resolve,
+          error: (xhr, errorType, error) => {
+            reject(error || new Error(`Request failed: ${xhr.status}`))
+          }
+        })
+      })
+      assert.strictEqual(r.length, data.length)
+    })
   })
 
   describe('protocol.isProtocolHandled', (done) => {