Don't follow external redirects

lib/webrat/core/session.rb

@@ -110,7 +110,7 @@ For example:
       @http_method  = http_method
       @data         = data
 
-      request_page(response.headers["Location"], :get, data) if redirect?
+      request_page(response_location, :get, data) if internal_redirect?
 
       return response
     end
@@ -118,11 +118,15 @@ For example:
     def success_code? #:nodoc:
       (200..499).include?(response_code)
     end
-    
+
     def redirect? #:nodoc:
       response_code / 100 == 3
     end
 
+    def internal_redirect? #:nodoc:
+      redirect? && current_host == response_location_host
+    end
+
     def exception_caught? #:nodoc:
       response_body =~ /Exception caught/
     end
@@ -222,6 +226,18 @@ For example:
 
   private
 
+    def response_location
+      response.headers["Location"]
+    end
+
+    def current_host
+      URI.parse(current_url).host || "www.example.com"
+    end
+
+    def response_location_host
+      URI.parse(response_location).host || "www.example.com"
+    end
+
     def reset
       @elements     = {}
       @_scopes      = nil

spec/integration/merb/app/controllers/testing.rb

@@ -7,8 +7,12 @@ class Testing < Application
   def submit_form
   end
   
-  def redirect_to_root
+  def internal_redirect
     redirect "/"
   end
   
+  def external_redirect
+    redirect "http://google.com"
+  end
+  
 end

spec/integration/merb/config/router.rb

@@ -28,5 +28,6 @@
 Merb.logger.info("Compiling routes...")
 Merb::Router.prepare do
   match("/").to(:controller => "testing", :action => "show_form")
-  match("/redirect").to(:controller => "testing", :action => "redirect_to_root")
+  match("/internal_redirect").to(:controller => "testing", :action => "internal_redirect")
+  match("/external_redirect").to(:controller => "testing", :action => "external_redirect")
 end

spec/integration/merb/spec/webrat_spec.rb

@@ -14,8 +14,14 @@ describe "Webrat" do
     click_button "Test"
   end
 
-  it "should follow redirects" do
-    response = visit "/redirect"
+  it "should follow internal redirects" do
+    response = visit "/internal_redirect"
+    response.status.should == 200
     response.should contain("Webrat Form")
   end
+
+  it "should not follow external redirects" do
+    response = visit "/external_redirect"
+    response.status.should == 302
+  end
 end

spec/integration/rails/app/controllers/webrat_controller.rb

@@ -7,8 +7,12 @@ class WebratController < ApplicationController
     render :text => "OK"
   end
 
-  def redirect
+  def internal_redirect
     redirect_to :submit
   end
+  
+  def external_redirect
+    redirect_to "http://google.com"
+  end
 
 end

spec/integration/rails/config/routes.rb

@@ -1,7 +1,8 @@
 ActionController::Routing::Routes.draw do |map|
   map.with_options :controller => "webrat" do |webrat|
-    webrat.submit   "/submit",    :action => "submit"
-    webrat.redirect "/redirect",  :action => "redirect"
+    webrat.submit             "/submit",            :action => "submit"
+    webrat.internal_redirect  "/internal_redirect", :action => "internal_redirect"
+    webrat.external_redirect  "/external_redirect", :action => "external_redirect"
 
     webrat.root :action => "form"
   end

spec/integration/rails/test/integration/webrat_test.rb

@@ -15,8 +15,13 @@ class WebratTest < ActionController::IntegrationTest
     click_button "Test"
   end
 
-  test "should follow redirects" do
-    visit redirect_path
+  test "should follow internal redirects" do
+    visit internal_redirect_path
     assert response.body.include?("OK")
   end
+  
+  test "should not follow external redirects" do
+    visit external_redirect_path
+    assert response.redirect?
+  end
 end

spec/integration/sinatra/app.rb

@@ -1,26 +1,30 @@
 require "rubygems"
 require "sinatra"
- 
+
 use_in_file_templates!
- 
+
 get "/" do
   erb :home
 end
- 
+
 get "/go" do
   erb :go
 end
 
-get "/redirect" do
+get "/internal_redirect" do
   redirect "/"
 end
- 
+
+get "/external_redirect" do
+  redirect "http://google.com"
+end
+
 post "/go" do
   @user = params[:name]
   @email = params[:email]
   erb :hello
 end
- 
+
 __END__
 
 @@ layout
@@ -31,10 +35,10 @@ __END__
     <%= yield %>
   </body>
 </html>
- 
+
 @@ home
 <p> visit <a href="/go">there</a></p>
- 
+
 @@ go
 <form method="post" action="/go">
   <div>
@@ -47,7 +51,7 @@ __END__
   </div>
   <input type="submit" value="Submit" />
 </form>
- 
+
 @@ hello
 <p>Hello, <%= @user %></p>
 <p>Your email is: <%= @email %></p>

spec/integration/sinatra/test/webrat_test.rb

@@ -8,19 +8,24 @@ class WebratTest < Test::Unit::TestCase
     click_link "there"
     assert response_body.include?('<form method="post" action="/go">')
   end
-  
+
   def test_submits_form
     visit "/go"
     fill_in "Name", :with => "World"
     fill_in "Email", :with => "world@example.org"
     click_button "Submit"
-    
+
     assert response_body.include?("Hello, World")
     assert response_body.include?("Your email is: world@example.org")
   end
-  
-  def test_follows_redirects
-    visit "/redirect"
+
+  def test_follows_internal_redirects
+    visit "/internal_redirect"
     assert response_body.include?("visit")
   end
+
+  def test_does_not_follow_external_redirects
+    visit "/external_redirect"
+    assert response_code == 302
+  end
 end

spec/private/core/session_spec.rb

@@ -113,14 +113,22 @@ describe Webrat::Session do
       lambda { webrat_session.request_page('some url', :get, {}) }.should raise_error(Webrat::PageLoadError)
     end
 
-    it "should follow redirects" do
-      webrat_session.should_receive(:redirect?).twice.and_return(true, false)
+    it "should follow internal redirects" do
+      webrat_session.should_receive(:internal_redirect?).twice.and_return(true, false)
       webrat_session.response.should_receive(:headers).once.and_return({ "Location" => "/newurl" })
 
       webrat_session.request_page("/oldurl", :get, {})
 
       webrat_session.current_url.should == "/newurl"
     end
+
+    it "should now follow external redirects" do
+      webrat_session.should_receive(:internal_redirect?).and_return(false)
+
+      webrat_session.request_page("/oldurl", :get, {})
+
+      webrat_session.current_url.should == "/oldurl"
+    end
   end
 
   describe "#redirect?" do
@@ -138,4 +146,29 @@ describe Webrat::Session do
       webrat_session.redirect?.should be_false
     end
   end
+  
+  describe "#internal_redirect?" do
+    before(:each) do
+      webrat_session = Webrat::Session.new
+    end
+
+    it "should return true if the last response was a redirect and the host of the current_url matches that of the response location" do
+      webrat_session.stub!(:redirect?         => true)
+      webrat_session.stub!(:current_url       => "http://example.com")
+      webrat_session.stub!(:response_location => "http://example.com")
+      webrat_session.internal_redirect?.should be_true
+    end
+    
+    it "should return false if the last response was not a redirect" do
+      webrat_session.stub!(:redirect? => false)
+      webrat_session.internal_redirect?.should be_false
+    end
+    
+    it "should return false if the last response was a redirect but the host of the current_url doesn't matches that of the response location" do
+      webrat_session.stub!(:redirect?         => true)
+      webrat_session.stub!(:current_url       => "http://example.com")
+      webrat_session.stub!(:response_location => "http://google.com")
+      webrat_session.internal_redirect?.should be_false
+    end
+  end
 end

spec/public/visit_spec.rb

@@ -31,14 +31,22 @@ describe "visit" do
     lambda { fill_in "foo", :with => "blah" }.should raise_error(Webrat::WebratError)
   end
 
-  it "should follow redirects" do
-    webrat_session.should_receive(:redirect?).twice.and_return(true, false)
+  it "should follow internal redirects" do
+    webrat_session.should_receive(:internal_redirect?).twice.and_return(true, false)
     webrat_session.response.should_receive(:headers).once.and_return({ "Location" => "/newurl" })
 
     visit("/oldurl")
 
     current_url.should == "/newurl"
   end
+
+  it "should not follow external redirects" do
+    webrat_session.should_receive(:internal_redirect?).and_return(false)
+
+    visit("/oldurl")
+
+    current_url.should == "/oldurl"
+  end
 end
 
 describe "visit with referer" do