Allows response_cb to register a body filter so chunked and contentfi…

…lter works in streamed response as well.
ray1729 · Oct 17, 2009 · 8ab5bb9 · 8ab5bb9
1 parent 1e4026d
commit 8ab5bb9
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 43 deletions.
diff --git a/lib/Plack/Middleware.pm b/lib/Plack/Middleware.pm
@@ -3,6 +3,7 @@ use strict;
 use warnings;
 use parent qw/Class::Accessor::Fast/;
 use Carp ();
+use Plack::Util;
 use overload '&{}' => sub { shift->to_app(@_) }, fallback => 1;
 
 __PACKAGE__->mk_accessors(qw/app/);
@@ -32,16 +33,42 @@ sub to_app {
 sub response_cb {
     my($self, $res, $cb) = @_;
 
+    my $body_filter = sub {
+        my($cb, $res) = @_;
+        my $filter_cb = $cb->($res);
+        # If response_cb returns a callback, treat it as a $body filter
+        if (defined $filter_cb && ref $filter_cb eq 'CODE') {
+            if (defined $res->[2]) {
+                my $body    = $res->[2];
+                my $getline = ref $body eq 'ARRAY' ? sub { shift @$body } : sub { $body->getline };
+                $res->[2] = Plack::Util::inline_object
+                    getline => sub { $filter_cb->($getline->()) },
+                    close => sub { $body->close if ref $body ne 'ARRAY' };
+            } else {
+                return $filter_cb;
+            }
+        }
+    };
+
     if (ref $res eq 'ARRAY') {
-        $cb->($res);
+        $body_filter->($cb, $res);
         return $res;
     } elsif (ref $res eq 'CODE') {
         return sub {
             my $respond = shift;
             $res->(sub {
                 my $res = shift;
-                $cb->($res);
-                $respond->($res);
+                my $filter_cb = $body_filter->($cb, $res);
+                if ($filter_cb) {
+                    my $writer = $respond->($res);
+                    if ($writer) {
+                        return Plack::Util::inline_object
+                            write => sub { $writer->write($filter_cb->(@_)) },
+                            close => sub { $writer->close };
+                    }
+                } else {
+                    return $respond->($res);
+                }
             });
         };
     }

diff --git a/lib/Plack/Middleware/Chunked.pm b/lib/Plack/Middleware/Chunked.pm
@@ -13,25 +13,20 @@ sub call {
         if ($env->{'SERVER_PROTOCOL'} ne 'HTTP/1.0' and
             ! Plack::Util::status_with_no_entity_body($res->[0]) and
             ! $h->exists('Content-Length') and
-            ! $h->exists('Transfer-Encoding') and
-            defined $res->[2]
+            ! $h->exists('Transfer-Encoding')
         ) {
             $h->set('Transfer-Encoding' => 'chunked');
-            my $body    = $res->[2];
-            my $getline = ref $body eq 'ARRAY' ? sub { shift @$body } : sub { $body->getline };
             my $done;
-            $res->[2] = Plack::Util::inline_object
-                getline => sub {
-                    my $chunk = $getline->();
-                    return if $done;
-                    unless (defined $chunk) {
-                        $done = 1;
-                        return "0\015\012\015\012";
-                    }
-                    return '' unless length $chunk;
-                    return sprintf('%x', length $chunk) . "\015\012$chunk\015\012";
-                },
-                close => sub { $body->close if ref $body ne 'ARRAY' };
+            return sub {
+                my $chunk = shift;
+                return if $done;
+                unless (defined $chunk) {
+                    $done = 1;
+                    return "0\015\012\015\012";
+                }
+                return '' unless length $chunk;
+                return sprintf('%x', length $chunk) . "\015\012$chunk\015\012";
+            };
         }
     });
 }

diff --git a/lib/Plack/Middleware/SimpleContentFilter.pm b/lib/Plack/Middleware/SimpleContentFilter.pm
@@ -9,30 +9,20 @@ use Plack::Util;
 sub call {
     my $self = shift;
 
-    my($status, $header, $body) = @{$self->app->(@_)};
-
-    my $h = Plack::Util::headers($header);
-
-    unless ($h->get('Content-Type') =~ m!^text/!) {
-        return [ $status, $header, $body ]
-    }
-
-    my $getline = ref $body eq 'ARRAY' ? sub { shift @$body } : sub { $body->getline };
-
-    my $body_filter = Plack::Util::inline_object(
-        getline => sub {
-            my $line = $getline->();
-            return unless defined $line;
-            local $_ = $line;
-            $self->filter->();
-            return $_;
-        },
-        close => sub {
-            $body->close unless ref $body eq 'ARRAY';
-        },
-    );
-
-    return [ $status, $header, $body_filter ];
+    my $res = $self->app->(@_);
+    $self->response_cb($res, sub {
+        my $res = shift;
+        my $h = Plack::Util::headers($res->[1]);
+        if ($h->get('Content-Type') =~ m!^text/!) {
+            return sub {
+                my $chunk = shift;
+                return unless defined $chunk;
+                local $_ = $chunk;
+                $self->filter->();
+                return $_;
+            };
+        }
+    });
 }
 
 1;