Support sub uri and lambda uri in Shield::Middleware

lib/shield.rb

@@ -3,8 +3,6 @@
 
 module Shield
   class Middleware
-    attr :url
-
     def initialize(app, url = "/login")
       @app = app
       @url = url
@@ -14,15 +12,23 @@ def call(env)
       tuple = @app.call(env)
 
       if tuple[0] == 401
-        [302, headers(env["SCRIPT_NAME"] + env["PATH_INFO"]), []]
+        [302, headers(url(env), env["SCRIPT_NAME"] + env["PATH_INFO"]), []]
       else
         tuple
       end
     end
 
+    def url(env)
+      if @url.respond_to? :call
+        @url.call(env)
+      else
+        env["SCRIPT_NAME"] + @url
+      end
+    end
+
   private
-    def headers(path)
-      { "Location" => "%s?return=%s" % [url, encode(path)],
+    def headers(redirect_url, return_path)
+      { "Location" => "%s?return=%s" % [redirect_url, encode(return_path)],
         "Content-Type" => "text/html",
         "Content-Length" => "0"
       }

test/middleware.rb

@@ -28,3 +28,11 @@
   assert_equal 302, status
   assert_equal "/login?return=%2Fsecured", headers["Location"]
 end
+
+test do
+  env = { "PATH_INFO" => "/secured", "SCRIPT_NAME" => "/suburi" }
+  status, headers, body = Cuba.call(env)
+
+  assert_equal 302, status
+  assert_equal "/suburi/login?return=%2Fsuburi%2Fsecured", headers["Location"]
+end

test/middleware_lambda.rb

@@ -0,0 +1,30 @@
+require_relative "helper"
+require_relative "user"
+require "cuba"
+
+Cuba.use Rack::Session::Cookie, secret: "foo"
+Cuba.use Shield::Middleware, lambda { |env| env["SCRIPT_NAME"] + "/login" } 
+
+Cuba.plugin Shield::Helpers
+
+Cuba.define do
+  on "secured" do
+    if not authenticated(User)
+      halt [401, { "Content-Type" => "text/html" }, []]
+    end
+
+    res.write "You're in"
+  end
+
+  on "foo" do
+    puts env.inspect
+  end
+end
+
+test do
+  env = { "PATH_INFO" => "/secured", "SCRIPT_NAME" => "/lambda" }
+  status, headers, body = Cuba.call(env)
+
+  assert_equal 302, status
+  assert_equal "/lambda/login?return=%2Flambda%2Fsecured", headers["Location"]
+end

test/nested.rb

@@ -6,7 +6,7 @@
 Cuba.plugin Shield::Helpers
 
 class Admin < Cuba
-  use Shield::Middleware, "/admin/login"
+  use Shield::Middleware, "/login"
 
   define do
     on "login" do