From 390cf34816e0c213550ccb3e4e085700c8409de1 Mon Sep 17 00:00:00 2001
From: Christian Schmidt <github@chsc.dk>
Date: Wed, 4 Sep 2024 21:09:35 +0200
Subject: [PATCH] Do not consume body before yielding

---
 app/lib/request.rb       |  9 +--------
 spec/lib/request_spec.rb | 41 +++++++++++++++++++++++++++++++++++-----
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/app/lib/request.rb b/app/lib/request.rb
index d7da9fe63c1..fb07ecca08b 100644
--- a/app/lib/request.rb
+++ b/app/lib/request.rb
@@ -111,16 +111,9 @@ class Request
     end
 
     begin
-      # If we are using a persistent connection, we have to
-      # read every response to be able to move forward at all.
-      # However, simply calling #to_s or #flush may not be safe,
-      # as the response body, if malicious, could be too big
-      # for our memory. So we use the #body_with_limit method
-      response.body_with_limit if http_client.persistent?
-
       yield response if block_given?
     ensure
-      http_client.close unless http_client.persistent?
+      http_client.close unless http_client.persistent? && response.connection.finished_request?
     end
   end
 
diff --git a/spec/lib/request_spec.rb b/spec/lib/request_spec.rb
index c600a48ee24..69aa028c48c 100644
--- a/spec/lib/request_spec.rb
+++ b/spec/lib/request_spec.rb
@@ -4,7 +4,9 @@ require 'rails_helper'
 require 'securerandom'
 
 RSpec.describe Request do
-  subject { described_class.new(:get, 'http://example.com') }
+  subject { described_class.new(:get, 'http://example.com', **options) }
+
+  let(:options) { {} }
 
   describe '#headers' do
     it 'returns user agent' do
@@ -39,8 +41,8 @@ RSpec.describe Request do
   end
 
   describe '#perform' do
-    context 'with valid host' do
-      before { stub_request(:get, 'http://example.com') }
+    context 'with valid host and non-persistent connection' do
+      before { stub_request(:get, 'http://example.com').to_return(body: 'lorem ipsum') }
 
       it 'executes a HTTP request' do
         expect { |block| subject.perform(&block) }.to yield_control
@@ -71,9 +73,9 @@ RSpec.describe Request do
         expect(subject.send(:http_client)).to have_received(:close)
       end
 
-      it 'returns response which implements body_with_limit' do
+      it 'yields response' do
         subject.perform do |response|
-          expect(response).to respond_to :body_with_limit
+          expect(response.body_with_limit).to eq 'lorem ipsum'
         end
       end
     end
@@ -95,6 +97,35 @@ RSpec.describe Request do
         expect { subject.perform }.to raise_error Mastodon::ValidationError
       end
     end
+
+    context 'with persistent connection' do
+      before { stub_request(:get, 'http://example.com').to_return(body: SecureRandom.random_bytes(100.kilobytes)) }
+
+      let(:http_client) { described_class.http_client.persistent('http://example.com') }
+      let(:options) { { http_client: http_client } }
+
+      it 'leaves connection open after complete response' do
+        allow(http_client).to receive(:close)
+
+        subject.perform(&:truncated_body)
+
+        expect(http_client).to_not have_received(:close)
+      end
+
+      it 'closes connection after aborted response' do
+        allow(http_client).to receive(:close)
+
+        subject.perform { |response| response.truncated_body(1.kilobyte) }
+
+        expect(http_client).to have_received(:close)
+      end
+
+      it 'yields response' do
+        subject.perform do |response|
+          expect(response.body_with_limit.size).to eq 100.kilobytes
+        end
+      end
+    end
   end
 
   describe "response's body_with_limit method" do