support code option

yappo · Aug 17, 2011 · 0c9980d · 0c9980d
1 parent 08c0e0f
commit 0c9980d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 deletions.
diff --git a/lib/Plack/Middleware/ErrorpageRedirectloopDefence.pm b/lib/Plack/Middleware/ErrorpageRedirectloopDefence.pm
@@ -5,17 +5,24 @@ use parent qw(Plack::Middleware);
 our $VERSION = '0.09';
 
 use Plack::Request;
+use Plack::Util::Accessor qw/code/;
+
+my $res_notcontent = [ 204, [ 'Content-Type', 'text/html', 'Content-Length', '0' ], [ '' ] ];
 
 sub call {
     my $self = shift;
     my $env = shift;
     my $req = Plack::Request->new($env);
+    my $res = $self->app->($env);
+    return $res unless $req->uri eq ($req->header('referer') || '');
 
-    if ($req->uri eq $req->header('referer')) {
-        return [ 204, [ 'Content-Type', 'text/html', 'Content-Length', '0' ], [ '' ] ];
-    }
+    my $code = $self->code;
+    return $res unless $code;
 
-    $self->app->($env);
+    if (((ref $code || '') eq 'Regexp' && $res->[0] =~ $code) || $res->[0].'' eq "$code") {
+        return $res_notcontent;
+    }
+    return $res;
 }
 
 

diff --git a/sample.psgi b/sample.psgi
@@ -19,11 +19,18 @@ my $html = <<'HTML';
 HTML
 
 builder {
-    enable 'ErrorpageRedirectloopDefence';
+    enable 'ErrorpageRedirectloopDefence',
+        code => qr/[45]\d\d/;
+#        code => '502';
+
     sub {
         my $req = Plack::Request->new(shift);
         my $body = sprintf $html, $req->uri, time(), $req->uri;
-        [ 502, [ 'Content-Type', 'text/html', 'Content-Length', length($body) ], [ $body ] ];
+        if ($req->path eq '/ok') {
+            return [ 200, [ 'Content-Type', 'text/html', 'Content-Length', length($body) ], [ $body ] ];
+        } else {
+            return [ 502, [ 'Content-Type', 'text/html', 'Content-Length', length($body) ], [ $body ] ];
+        }
     };
 };