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