Store the PATH_INFO and response code in env["rack-unbasic.*"]

lib/rack/unbasic.rb

@@ -1,13 +1,16 @@
 module Rack
   class Unbasic
-    def initialize(app, &block) # :yields: middleware
+    def initialize(app, &middleware)
       @app = app
-      block.call(self) if block
+      @locations = {}
+      middleware.call(self) if middleware
     end
 
     def call(env)
       @env = env
 
+      clean_session_data
+
       response = @app.call(@env)
 
       case response[0].to_i
@@ -18,23 +21,38 @@ def call(env)
     end
 
     def unauthorized(location)
-      @unauthorized_location = location
+      @locations["unauthorized"] = location
     end
 
     def bad_request(location)
-      @bad_request_location = location
+      @locations["bad_request"] = location
     end
 
     private
 
     def handle_unauthorized
-      return nil if @unauthorized_location.nil?
-      [302, {"Location" => @unauthorized_location}, []]
+      return nil if @locations["unauthorized"].nil?
+      store_location_and_response_code(401)
+      [302, {"Location" => @locations["unauthorized"]}, []]
     end
 
     def handle_bad_request
-      return nil if @bad_request_location.nil?
-      [302, {"Location" => @bad_request_location}, []]
+      return nil if @locations["bad_request"].nil?
+      store_location_and_response_code(400)
+      [302, {"Location" => @locations["bad_request"]}, []]
+    end
+
+    def store_location_and_response_code(code)
+      @env["rack.session"]["rack-unbasic.return-to"] = @env["PATH_INFO"]
+      @env["rack.session"]["rack-unbasic.code"] = code
+    end
+
+    def clean_session_data
+      unless @env["rack.session"].respond_to?("delete")
+        raise "You need to enable sessions for this middleware to work"
+      end
+      @env["rack-unbasic.return-to"] = @env["rack.session"].delete("rack-unbasic.return-to")
+      @env["rack-unbasic.code"] = @env["rack.session"].delete("rack-unbasic.code")
     end
   end
 end

test/test_unbasic.rb

@@ -6,21 +6,24 @@
 class TestUnbasic < Test::Unit::TestCase
   include Rack::Test::Methods
 
-  def app_failing_with(code)
-    lambda do |env|
-      if env["PATH_INFO"] == "/login"
-        [200, {}, "Please login"]
-      else
-        [code, {}, "Go away!"]
-      end
-    end
+  def unbasic_app_failing_with(code, &unbasic)
+    Rack::Builder.new {
+      use Rack::Session::Cookie
+      use Rack::Unbasic, &unbasic
+      run lambda {|env| 
+        if env["PATH_INFO"] == "/login"
+          [200, {}, "Please login"]
+        else
+          [code, {}, "Go away!"]
+        end 
+      }
+    }
   end
 
   context "when the downstream app returns a 401 status code" do
     context "and Unbasic handles unauthorized requests" do
       def app
-        downstream = app_failing_with(401)
-        Rack::Unbasic.new(downstream) do |on|
+        @app ||= unbasic_app_failing_with 401 do |on|
           on.unauthorized "/login"
         end
       end
@@ -35,11 +38,17 @@ def app
         follow_redirect!
         assert_equal "http://example.org/login", last_request.url
       end
+
+      test "it saves the previous url in env['rack-unbasic.return-to']" do
+        get "/foobar"
+        follow_redirect!
+        assert_equal "/foobar", last_request.env["rack-unbasic.return-to"]
+      end
     end
 
     context "but Unbasic doesn't handle unauthorized requests" do
       def app
-        Rack::Unbasic.new(app_failing_with(401))
+        @app ||= unbasic_app_failing_with 401
       end
 
       test "it just returns the original response to the user" do
@@ -53,8 +62,7 @@ def app
   context "when the downstream app returns a 400 status code" do
     context "and Unbasic handles bad requests" do
       def app
-        downstream = app_failing_with(400)
-        Rack::Unbasic.new(downstream) do |on|
+        @app ||= unbasic_app_failing_with 400 do |on|
           on.bad_request "/login"
         end
       end
@@ -69,11 +77,17 @@ def app
         follow_redirect!
         assert_equal "http://example.org/login", last_request.url
       end
+
+      test "it saves the previous url in env['rack-unbasic.return-to']" do
+        get "/foobar"
+        follow_redirect!
+        assert_equal "/foobar", last_request.env["rack-unbasic.return-to"]
+      end
     end
 
     context "but Unbasic doesn't handle bad requests" do
       def app
-        Rack::Unbasic.new(app_failing_with(400))
+        @app ||= unbasic_app_failing_with 400
       end
 
       test "it just returns the original response to the user" do