1 files changed, 77 insertions, 32 deletions

diff --git a/app/lib/request.rb b/app/lib/request.rb
index e555ae6a1..5f7075a3c 100644
--- a/app/lib/request.rb
+++ b/app/lib/request.rb
@@ -17,15 +17,22 @@ end
 class Request
   REQUEST_TARGET = '(request-target)'
 
+  # We enforce a 5s timeout on DNS resolving, 5s timeout on socket opening
+  # and 5s timeout on the TLS handshake, meaning the worst case should take
+  # about 15s in total
+  TIMEOUT = { connect: 5, read: 10, write: 10 }.freeze
+
   include RoutingHelper
 
   def initialize(verb, url, **options)
     raise ArgumentError if url.blank?
 
-    @verb    = verb
-    @url     = Addressable::URI.parse(url).normalize
-    @options = options.merge(use_proxy? ? Rails.configuration.x.http_client_proxy : { socket_class: Socket })
-    @headers = {}
+    @verb        = verb
+    @url         = Addressable::URI.parse(url).normalize
+    @http_client = options.delete(:http_client)
+    @options     = options.merge(socket_class: use_proxy? ? ProxySocket : Socket)
+    @options     = @options.merge(Rails.configuration.x.http_client_proxy) if use_proxy?
+    @headers     = {}
 
     raise Mastodon::HostValidationError, 'Instance does not support hidden service connections' if block_hidden_service?
 
@@ -50,15 +57,24 @@ class Request
 
   def perform
     begin
-      response = http_client.headers(headers).public_send(@verb, @url.to_s, @options)
+      response = http_client.public_send(@verb, @url.to_s, @options.merge(headers: headers))
     rescue => e
       raise e.class, "#{e.message} on #{@url}", e.backtrace[0]
     end
 
     begin
-      yield response.extend(ClientLimit) if block_given?
+      response = response.extend(ClientLimit)
+
+      # 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
+      http_client.close unless http_client.persistent?
     end
   end
 
@@ -76,6 +92,10 @@ class Request
 
       %w(http https).include?(parsed_url.scheme) && parsed_url.host.present?
     end
+
+    def http_client
+      HTTP.use(:auto_inflate).timeout(:per_operation, TIMEOUT.dup).follow(max_hops: 2)
+    end
   end
 
   private
@@ -116,16 +136,8 @@ class Request
     end
   end
 
-  def timeout
-    # We enforce a 1s timeout on DNS resolving, 10s timeout on socket opening
-    # and 5s timeout on the TLS handshake, meaning the worst case should take
-    # about 16s in total
-
-    { connect: 5, read: 10, write: 10 }
-  end
-
   def http_client
-    @http_client ||= HTTP.use(:auto_inflate).timeout(:per_operation, timeout).follow(max_hops: 2)
+    @http_client ||= Request.http_client
   end
 
   def use_proxy?
@@ -166,26 +178,49 @@ class Request
   class Socket < TCPSocket
     class << self
       def open(host, *args)
-        return super(host, *args) if thru_hidden_service?(host)
-
         outer_e = nil
+        port    = args.first
 
-        Resolv::DNS.open do |dns|
-          dns.timeouts = 5
+        addresses = []
+        begin
+          addresses = [IPAddr.new(host)]
+        rescue IPAddr::InvalidAddressError
+          Resolv::DNS.open do |dns|
+            dns.timeouts = 5
+            addresses = dns.getaddresses(host).take(2)
+          end
+        end
 
-          addresses = dns.getaddresses(host).take(2)
-          time_slot = 10.0 / addresses.size
+        addresses.each do |address|
+          begin
+            check_private_address(address)
 
-          addresses.each do |address|
-            begin
-              raise Mastodon::HostValidationError if PrivateAddressCheck.private_address?(IPAddr.new(address.to_s))
+            sock     = ::Socket.new(address.is_a?(Resolv::IPv6) ? ::Socket::AF_INET6 : ::Socket::AF_INET, ::Socket::SOCK_STREAM, 0)
+            sockaddr = ::Socket.pack_sockaddr_in(port, address.to_s)
+
+            sock.setsockopt(::Socket::IPPROTO_TCP, ::Socket::TCP_NODELAY, 1)
 
-              ::Timeout.timeout(time_slot, HTTP::TimeoutError) do
-                return super(address.to_s, *args)
+            begin
+              sock.connect_nonblock(sockaddr)
+            rescue IO::WaitWritable
+              if IO.select(nil, [sock], nil, Request::TIMEOUT[:connect])
+                begin
+                  sock.connect_nonblock(sockaddr)
+                rescue Errno::EISCONN
+                  # Yippee!
+                rescue
+                  sock.close
+                  raise
+                end
+              else
+                sock.close
+                raise HTTP::TimeoutError, "Connect timed out after #{Request::TIMEOUT[:connect]} seconds"
               end
-            rescue => e
-              outer_e = e
             end
+
+            return sock
+          rescue => e
+            outer_e = e
           end
         end
 
@@ -198,11 +233,21 @@ class Request
 
       alias new open
 
-      def thru_hidden_service?(host)
-        Rails.configuration.x.access_to_hidden_service && /\.(onion|i2p)$/.match(host)
+      def check_private_address(address)
+        raise Mastodon::HostValidationError if PrivateAddressCheck.private_address?(IPAddr.new(address.to_s))
+      end
+    end
+  end
+
+  class ProxySocket < Socket
+    class << self
+      def check_private_address(_address)
+        # Accept connections to private addresses as HTTP proxies will usually
+        # be on local addresses
+        nil
       end
     end
   end
 
-  private_constant :ClientLimit, :Socket
+  private_constant :ClientLimit, :Socket, :ProxySocket
 end