Merge pull request #2033 from basecamp/limit-webhook-response-size

Stream response in webhook request manually to check size

app/models/webhook/delivery.rb

@@ -1,7 +1,10 @@
 class Webhook::Delivery < ApplicationRecord
+  class ResponseTooLarge < StandardError; end
+
   STALE_TRESHOLD = 7.days
   USER_AGENT = "fizzy/1.0.0 Webhook"
   ENDPOINT_TIMEOUT = 7.seconds
+  MAX_RESPONSE_SIZE = 100.kilobytes
 
   belongs_to :account, default: -> { webhook.account }
   belongs_to :webhook
@@ -52,12 +55,16 @@ class Webhook::Delivery < ApplicationRecord
       if resolved_ip.nil?
         { error: :private_uri }
       else
-        response = http.request(
-          Net::HTTP::Post.new(uri, headers).tap { |request| request.body = payload }
-        )
+        request = Net::HTTP::Post.new(uri, headers).tap { |request| request.body = payload }
+
+        response = http.request(request) do |net_http_response|
+          stream_body_with_limit(net_http_response)
+        end
 
         { code: response.code.to_i }
       end
+    rescue ResponseTooLarge
+      { error: :response_too_large }
     rescue Resolv::ResolvTimeout, Resolv::ResolvError, SocketError
       { error: :dns_lookup_failed }
     rescue Net::OpenTimeout, Net::ReadTimeout, Errno::ETIMEDOUT
@@ -68,6 +75,14 @@ class Webhook::Delivery < ApplicationRecord
       { error: :failed_tls }
     end
 
+    def stream_body_with_limit(response)
+      bytes_read = 0
+      response.read_body do |chunk|
+        bytes_read += chunk.bytesize
+        raise ResponseTooLarge if bytes_read > MAX_RESPONSE_SIZE
+      end
+    end
+
     def resolved_ip
       return @resolved_ip if defined?(@resolved_ip)
       @resolved_ip = SsrfProtection.resolve_public_ip(uri.host)

test/models/webhook/delivery_test.rb

@@ -296,12 +296,15 @@ class Webhook::DeliveryTest < ActiveSupport::TestCase
     stub_dns_resolution(PUBLIC_TEST_IP)
 
     # Verify Net::HTTP.new is called with the pinned IP
+    response_mock = stub(code: "200")
+    response_mock.stubs(:read_body)
+
     http_mock = mock("http")
     http_mock.stubs(:use_ssl=)
     http_mock.stubs(:ipaddr=)
     http_mock.stubs(:open_timeout=)
     http_mock.stubs(:read_timeout=)
-    http_mock.stubs(:request).returns(stub(code: "200"))
+    http_mock.stubs(:request).yields(response_mock).returns(response_mock)
 
     Net::HTTP.expects(:new).with("example.com", 443).returns(http_mock)
 
@@ -310,6 +313,32 @@ class Webhook::DeliveryTest < ActiveSupport::TestCase
     assert delivery.succeeded?
   end
 
+  test "handles response too large error" do
+    delivery = webhook_deliveries(:pending)
+
+    large_body = "x" * 200.kilobytes
+    stub_request(:post, delivery.webhook.url).to_return(status: 200, body: large_body)
+
+    delivery.deliver
+
+    assert_equal "completed", delivery.state
+    assert_equal "response_too_large", delivery.response[:error]
+    assert_not delivery.succeeded?
+  end
+
+  test "allows responses within size limit" do
+    delivery = webhook_deliveries(:pending)
+
+    small_body = "x" * 50.kilobytes
+    stub_request(:post, delivery.webhook.url).to_return(status: 200, body: small_body)
+
+    delivery.deliver
+
+    assert_equal "completed", delivery.state
+    assert_equal 200, delivery.response[:code]
+    assert delivery.succeeded?
+  end
+
   private
     def stub_dns_resolution(*ips)
       dns_mock = mock("dns")