From 40e5d2303ba1edc51beae66cc15263675980106a Mon Sep 17 00:00:00 2001
From: Akihiko Odaki <akihiko.odaki.4i@stu.hosei.ac.jp>
Date: Mon, 26 Mar 2018 21:02:10 +0900
Subject: 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.
---
 app/lib/request.rb | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

(limited to 'app/lib/request.rb')

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
-- 
cgit