add max_response_size attribute to Mojo::UserAgent

mojolicious · Jan 28, 2017 · fa35d3b · fa35d3b
1 parent 9438ede
commit fa35d3b
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 15 deletions.
diff --git a/Changes b/Changes
@@ -1,5 +1,6 @@
 
-7.23  2017-01-27
+7.23  2017-01-28
+  - Added max_response_size attribute to Mojo::UserAgent.
   - Added to_unsafe_string method to Mojo::URL.
   - Fixed a data corruption problem in Mojo::IOLoop::Stream, caused by a
     dependency of IO::Socket::SSL on the internal representation of strings,

diff --git a/lib/Mojo/UserAgent.pm b/lib/Mojo/UserAgent.pm
@@ -23,8 +23,9 @@ has ioloop             => sub { Mojo::IOLoop->new };
 has key                => sub { $ENV{MOJO_KEY_FILE} };
 has max_connections    => 5;
 has max_redirects => sub { $ENV{MOJO_MAX_REDIRECTS} || 0 };
-has proxy => sub { Mojo::UserAgent::Proxy->new };
-has request_timeout => sub { $ENV{MOJO_REQUEST_TIMEOUT} // 0 };
+has max_response_size => sub { $ENV{MOJO_MAX_RESPONSE_SIZE} // 2147483648 };
+has proxy             => sub { Mojo::UserAgent::Proxy->new };
+has request_timeout   => sub { $ENV{MOJO_REQUEST_TIMEOUT}   // 0 };
 has server => sub { Mojo::UserAgent::Server->new(ioloop => shift->ioloop) };
 has transactor => sub { Mojo::UserAgent::Transactor->new };
 
@@ -298,6 +299,7 @@ sub _start {
   }
 
   $_->prepare($tx) for $self->proxy, $self->cookie_jar;
+  $tx->res->max_message_size($self->max_response_size);
 
   # Connect and add request timeout if necessary
   my $id = $self->emit(start => $tx)->_connection($loop, $tx, $cb);
@@ -573,6 +575,18 @@ Maximum number of redirects the user agent will follow before it fails,
 defaults to the value of the C<MOJO_MAX_REDIRECTS> environment variable or
 C<0>.
 
+=head2 max_response_size
+
+  my $max = $ua->max_response_size;
+  $ua     = $ua->max_response_size(16777216);
+
+Maximum response size in bytes, defaults to the value of the
+C<MOJO_MAX_RESPONSE_SIZE> environment variable or C<2147483648> (2GB). Setting
+the value to C<0> will allow responses of indefinite size. Note that increasing
+this value can also drastically increase memory usage, should you for example,
+attempt to parse an excessively large response body with the methods
+L<Mojo::Message/"dom"> or L<Mojo::Message/"json">.
+
 =head2 proxy
 
   my $proxy = $ua->proxy;

diff --git a/lib/Mojolicious/Command/get.pm b/lib/Mojolicious/Command/get.pm
@@ -15,14 +15,16 @@ has usage => sub { shift->extract_usage };
 sub run {
   my ($self, @args) = @_;
 
+  my $ua = Mojo::UserAgent->new(ioloop => Mojo::IOLoop->singleton);
   getopt \@args,
     'C|charset=s'            => \my $charset,
     'c|content=s'            => \(my $content = ''),
     'H|header=s'             => \my @headers,
-    'i|inactivity-timeout=i' => \my $inactivity,
+    'i|inactivity-timeout=i' => sub { $ua->inactivity_timeout($_[1]) },
     'M|method=s'             => \(my $method = 'GET'),
-    'o|connect-timeout=i'    => \my $connect,
+    'o|connect-timeout=i'    => sub { $ua->connect_timeout($_[1]) },
     'r|redirect'             => \my $redirect,
+    'S|response-size=i'      => sub { $ua->max_response_size($_[1]) },
     'v|verbose'              => \my $verbose;
 
   @args = map { decode 'UTF-8', $_ } @args;
@@ -33,13 +35,9 @@ sub run {
   my %headers = map { /^\s*([^:]+)\s*:\s*(.*+)$/ ? ($1, $2) : () } @headers;
 
   # Detect proxy for absolute URLs
-  my $ua = Mojo::UserAgent->new(ioloop => Mojo::IOLoop->singleton);
   $url !~ m!^/! ? $ua->proxy->detect : $ua->server->app($self->app);
   $ua->max_redirects(10) if $redirect;
 
-  $ua->inactivity_timeout($inactivity) if defined $inactivity;
-  $ua->connect_timeout($connect) if $connect;
-
   my $buffer = '';
   $ua->on(
     start => sub {
@@ -163,6 +161,10 @@ Mojolicious::Command::get - Get command
     -o, --connect-timeout <seconds>      Connect timeout, defaults to the value
                                          of MOJO_CONNECT_TIMEOUT or 10
     -r, --redirect                       Follow up to 10 redirects
+    -S, --response-size <size>           Maximum response size in bytes,
+                                         defaults to the value of
+                                         MOJO_MAX_RESPONSE_SIZE or
+                                         2147483648 (2GB)
     -v, --verbose                        Print request and response headers to
                                          STDERR
 

diff --git a/lib/Mojolicious/Guides/FAQ.pod b/lib/Mojolicious/Guides/FAQ.pod
@@ -189,9 +189,11 @@ Hypnotoad.
 
 To protect your applications from excessively large requests and responses, our
 HTTP parser has a cap after which it will automatically stop accepting new
-data, and in most cases force the connection to be closed. The limit is 16MB by
-default, you can use the attribute L<Mojo::Message/"max_message_size"> or
-C<MOJO_MAX_MESSAGE_SIZE> environment variable to change this value.
+data, and in most cases force the connection to be closed. The limit is 16MB on
+the server-side and 2GB on the client-side by default, you can use the
+attributes L<Mojo::Message/"max_message_size"> and
+L<Mojo::UserAgent/"max_response_size"> or C<MOJO_MAX_MESSAGE_SIZE> and
+C<MOJO_MAX_RESPONSE_SIZE> environment variables to change these values.
 
 =head2 What does the error "Maximum start-line size exceeded" mean?
 

diff --git a/t/mojo/user_agent.t b/t/mojo/user_agent.t
@@ -68,6 +68,15 @@ get '/one' => sub {
   is(Mojo::UserAgent->new->max_redirects, 0, 'right value');
 }
 
+# Max response size
+{
+  is(Mojo::UserAgent->new->max_response_size, 2147483648, 'right value');
+  local $ENV{MOJO_MAX_RESPONSE_SIZE} = 25;
+  is(Mojo::UserAgent->new->max_response_size, 25, 'right value');
+  local $ENV{MOJO_MAX_RESPONSE_SIZE} = 0;
+  is(Mojo::UserAgent->new->max_response_size, 0, 'right value');
+}
+
 # Timeouts
 {
   is(Mojo::UserAgent->new->connect_timeout, 10, 'right value');
@@ -156,7 +165,8 @@ ok !$tx->kept_alive, 'kept connection not alive';
 is $tx->res->version, '1.1', 'right version';
 is $tx->res->code,    200,   'right status';
 ok !$tx->res->headers->connection, 'no "Connection" value';
-is $tx->res->body, 'works!', 'right content';
+is $tx->res->max_message_size, 2147483648, 'right value';
+is $tx->res->body,             'works!',   'right content';
 
 # Again
 $tx = $ua->get('/');
@@ -166,12 +176,13 @@ is $tx->res->version, '1.1', 'right version';
 is $tx->res->code,    200,   'right status';
 ok !$tx->res->headers->connection, 'no "Connection" value';
 is $tx->res->body, 'works!', 'right content';
-$tx = $ua->get('/');
+$tx = $ua->max_response_size(33554432)->get('/');
 ok $tx->success, 'successful';
 is $tx->res->version, '1.1', 'right version';
 is $tx->res->code,    200,   'right status';
 ok !$tx->res->headers->connection, 'no "Connection" value';
-is $tx->res->body, 'works!', 'right content';
+is $tx->res->max_message_size, 33554432, 'right value';
+is $tx->res->body,             'works!', 'right content';
 
 # Shortcuts for common request methods
 is $ua->delete('/method')->res->body,  'DELETE',  'right content';