Validate HTTP response length while receiving (#6891)

to_s method of HTTP::Response keeps blocking while it receives the whole
content, no matter how it is big. This means it may waste time to receive
unacceptably large files. It may also consume memory and disk in the
process. This solves the inefficency by checking response length while
receiving.

1 files changed, 29 insertions, 2 deletions

diff --git a/app/lib/request.rb b/app/lib/request.rb
index 8a127c65f..dca93a6e9 100644
--- a/app/lib/request.rb
+++ b/app/lib/request.rb
@@ -40,7 +40,7 @@ class Request
     end
 
     begin
-      yield response
+      yield response.extend(ClientLimit)
     ensure
       http_client.close
     end
@@ -99,6 +99,33 @@ class Request
     @http_client ||= HTTP.timeout(:per_operation, timeout).follow(max_hops: 2)
   end
 
+  module ClientLimit
+    def body_with_limit(limit = 1.megabyte)
+      raise Mastodon::LengthValidationError if content_length.present? && content_length > limit
+
+      if charset.nil?
+        encoding = Encoding::BINARY
+      else
+        begin
+          encoding = Encoding.find(charset)
+        rescue ArgumentError
+          encoding = Encoding::BINARY
+        end
+      end
+
+      contents = String.new(encoding: encoding)
+
+      while (chunk = readpartial)
+        contents << chunk
+        chunk.clear
+
+        raise Mastodon::LengthValidationError if contents.bytesize > limit
+      end
+
+      contents
+    end
+  end
+
   class Socket < TCPSocket
     class << self
       def open(host, *args)
@@ -118,5 +145,5 @@ class Request
     end
   end
 
-  private_constant :Socket
+  private_constant :ClientLimit, :Socket
 end