support pass-through :encoding for deflate and gzip

These allow serving pre-compressed data off disk and
on-the-fly uncompressing for the few clients that do
not accept compressed responses.

6 files changed, 162 insertions, 0 deletions

diff --git a/lib/metropolis.rb b/lib/metropolis.rb
index a7cfc47..3afc3ec 100644
--- a/lib/metropolis.rb
+++ b/lib/metropolis.rb
@@ -3,6 +3,9 @@ require 'rack'
 require 'uri'
 
 module Metropolis
+  autoload :InputWrapper, 'metropolis/input_wrapper'
+  autoload :Deflate, 'metropolis/deflate'
+  autoload :Gzip, 'metropolis/gzip'
   autoload :TC, 'metropolis/tc'
   autoload :Hash, 'metropolis/hash'
 
diff --git a/lib/metropolis/common.rb b/lib/metropolis/common.rb
index 273a1b3..971accd 100644
--- a/lib/metropolis/common.rb
+++ b/lib/metropolis/common.rb
@@ -13,6 +13,15 @@ module Metropolis::Common
     if @readonly && @exclusive
       raise ArgumentError, ":readonly and :exclusive may not be used together"
     end
+    case @encoding = opts[:encoding]
+    when nil
+    when :deflate
+      extend(Metropolis::Deflate)
+    when :gzip
+      extend(Metropolis::Gzip)
+    else
+      raise ArgumentError, "unsupported encoding"
+    end
   end
 
   def r(code, body = nil)
diff --git a/lib/metropolis/deflate.rb b/lib/metropolis/deflate.rb
new file mode 100644
index 0000000..d585b43
--- /dev/null
+++ b/lib/metropolis/deflate.rb
@@ -0,0 +1,45 @@
+# -*- encoding: binary -*-
+require "zlib"
+
+# allows storing pre-deflated data on disk and serving it
+# as-is for clients that accept that deflate encoding
+module Metropolis::Deflate
+  def get(key, env)
+    status, headers, body = r = super
+    if 200 == status && /\bdeflate\b/ !~ env['HTTP_ACCEPT_ENCODING']
+      inflater = Zlib::Inflate.new(-Zlib::MAX_WBITS)
+      body[0] = "#{inflater.inflate(body[0])}#{inflater.finish}"
+      inflater.end
+      headers['Content-Length'] = body[0].size.to_s
+      headers.delete('Content-Encoding')
+      headers.delete('Vary')
+    end
+    r
+  end
+
+  def put(key, env)
+    Wrapper.new(env)
+    super(key, env)
+  end
+
+  def self.extended(obj)
+    obj.instance_eval do
+      @headers['Content-Encoding'] = 'deflate'
+      @headers['Vary'] = 'Accept-Encoding'
+    end
+  end
+
+  class Wrapper < Metropolis::InputWrapper
+
+    def read_all
+      deflater = Zlib::Deflate.new(
+        Zlib::DEFAULT_COMPRESSION,
+        # drop the zlib header which causes both Safari and IE to choke
+        -Zlib::MAX_WBITS,
+        Zlib::DEF_MEM_LEVEL,
+        Zlib::DEFAULT_STRATEGY
+      )
+      "#{deflater.deflate(@input.read)}#{deflater.finish}"
+    end
+  end
+end
diff --git a/lib/metropolis/gzip.rb b/lib/metropolis/gzip.rb
new file mode 100644
index 0000000..50d71ef
--- /dev/null
+++ b/lib/metropolis/gzip.rb
@@ -0,0 +1,38 @@
+# -*- encoding: binary -*-
+require "zlib"
+
+# allows storing pre-gzipped data on disk and serving it
+# as-is for clients that accept that gzip encoding
+module Metropolis::Gzip
+  def get(key, env)
+    status, headers, body = r = super
+    if 200 == status && /\bgzip\b/ !~ env['HTTP_ACCEPT_ENCODING']
+      body[0] = Zlib::GzipReader.new(StringIO.new(body[0])).read
+      headers['Content-Length'] = body[0].size.to_s
+      headers.delete('Content-Encoding')
+      headers.delete('Vary')
+    end
+    r
+  end
+
+  def put(key, env)
+    Wrapper.new(env)
+    super(key, env)
+  end
+
+  def self.extended(obj)
+    obj.instance_eval do
+      @headers['Content-Encoding'] = 'gzip'
+      @headers['Vary'] = 'Accept-Encoding'
+    end
+  end
+
+  class Wrapper < Metropolis::InputWrapper
+
+    def read_all
+      zipped = StringIO.new("")
+      Zlib::GzipWriter.wrap(zipped) { |io| io.write(@input.read) }
+      zipped.string
+    end
+  end
+end
diff --git a/lib/metropolis/input_wrapper.rb b/lib/metropolis/input_wrapper.rb
new file mode 100644
index 0000000..12a1bcc
--- /dev/null
+++ b/lib/metropolis/input_wrapper.rb
@@ -0,0 +1,21 @@
+# -*- encoding: binary -*-
+
+class Metropolis::InputWrapper
+  def initialize(env)
+    @input = env["rack.input"]
+    env["rack.input"] = self
+  end
+
+  def read(*args)
+    args.empty? and return read_all
+    ni
+  end
+
+  def ni(*args)
+    raise NotImplementedError
+  end
+
+  alias gets ni
+  alias each ni
+  alias rewind ni
+end
diff --git a/test/rack_read_write.rb b/test/rack_read_write.rb
index b3a8a1f..46c1764 100644
--- a/test/rack_read_write.rb
+++ b/test/rack_read_write.rb
@@ -7,6 +7,52 @@ require 'rack'
 module TestRackReadWrite
   attr_reader :app
 
+  def test_rack_read_write_deflated
+    @app = Metropolis.new(:uri => uri, :encoding => :deflate)
+    basic_rest
+
+    blob = "." * 1024 * 1024
+    o = { :lint => true, :fatal => true }
+    req = Rack::MockRequest.new(app)
+
+    r = req.put("/asdf", o.merge(:input => blob))
+    assert_equal 201, r.status
+    assert_equal "Created\n", r.body
+
+    r = req.get("/asdf", o.merge("HTTP_ACCEPT_ENCODING" => "deflate"))
+    assert_equal 200, r.status
+    assert_equal "deflate", r.headers['Content-Encoding']
+    assert r.body.size < blob.size
+
+    r = req.get("/asdf", o.merge("HTTP_ACCEPT_ENCODING" => "gzip"))
+    assert_equal 200, r.status
+    assert_nil r.headers['Content-Encoding']
+    assert_equal blob, r.body
+  end
+
+  def test_rack_read_write_gzipped
+    @app = Metropolis.new(:uri => uri, :encoding => :gzip)
+    basic_rest
+
+    blob = "." * 1024 * 1024
+    o = { :lint => true, :fatal => true }
+    req = Rack::MockRequest.new(app)
+
+    r = req.put("/asdf", o.merge(:input => blob))
+    assert_equal 201, r.status
+    assert_equal "Created\n", r.body
+
+    r = req.get("/asdf", o.merge("HTTP_ACCEPT_ENCODING" => "gzip"))
+    assert_equal 200, r.status
+    assert_equal "gzip", r.headers['Content-Encoding']
+    assert r.body.size < blob.size
+
+    r = req.get("/asdf", o.merge("HTTP_ACCEPT_ENCODING" => "deflate"))
+    assert_equal 200, r.status
+    assert_nil r.headers['Content-Encoding']
+    assert_equal blob, r.body
+  end
+
   def test_rack_read_write
     @app = Metropolis.new(:uri => uri)
     basic_rest