From f67e0ccaac07266fde4c01029ee3b42c40dd2fcd Mon Sep 17 00:00:00 2001
From: Phil Stewart <p.stewart@imperial.ac.uk>
Date: Fri, 23 Jan 2015 12:45:50 +0000
Subject: [PATCH] Support sub uri and lambda uri in Shield::Middleware

---
 lib/shield.rb             | 16 +++++++++++-----
 test/middleware.rb        |  8 ++++++++
 test/middleware_lambda.rb | 30 ++++++++++++++++++++++++++++++
 test/nested.rb            |  2 +-
 4 files changed, 50 insertions(+), 6 deletions(-)
 create mode 100644 test/middleware_lambda.rb

diff --git a/lib/shield.rb b/lib/shield.rb
index 47985e3..898785e 100644
--- a/lib/shield.rb
+++ b/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"
       }
diff --git a/test/middleware.rb b/test/middleware.rb
index 277a29d..57bad1f 100644
--- a/test/middleware.rb
+++ b/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
diff --git a/test/middleware_lambda.rb b/test/middleware_lambda.rb
new file mode 100644
index 0000000..c4ee8d8
--- /dev/null
+++ b/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
diff --git a/test/nested.rb b/test/nested.rb
index baf4916..80629d5 100644
--- a/test/nested.rb
+++ b/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