Initial commit

pegasus/plugins/compress.lua

@@ -0,0 +1,145 @@
+local zlib = require "zlib"
+
+local function zlib_name(lib)
+  if lib._VERSION and string.find(lib._VERSION, 'lua-zlib', nil, true) then
+    return 'lua-zlib'
+  end
+
+  if lib._VERSION and string.find(lib._VERSION, 'lzlib', nil, true) then
+    return 'lzlib'
+  end
+end
+
+local z_lib_name = assert(zlib_name(zlib), 'Unsupported zlib Lua binding')
+
+local ZlibStream = {} do
+  ZlibStream.__index = ZlibStream
+
+  ZlibStream.NO_COMPRESSION      = zlib.NO_COMPRESSION       or  0
+  ZlibStream.BEST_SPEED          = zlib.BEST_SPEED           or  1
+  ZlibStream.BEST_COMPRESSION    = zlib.BEST_COMPRESSION     or  9
+  ZlibStream.DEFAULT_COMPRESSION = zlib.DEFAULT_COMPRESSION  or -1
+  ZlibStream.STORE               = 0
+  ZlibStream.DEFLATE             = 8
+
+  if z_lib_name == 'lzlib' then
+    function ZlibStream:new(writer, level, method, windowBits)
+      level  = level or ZlibStream.DEFAULT_COMPRESSION
+      method = method or ZlibStream.DEFLATE
+
+      local o = setmetatable({
+        zd = assert(zlib.deflate(writer, level, method, windowBits));
+      }, self)
+
+      return o
+    end
+
+    function ZlibStream:write(chunk)
+      assert(self.zd:write(chunk))
+    end
+
+    function ZlibStream:close()
+      self.zd:close()
+    end
+
+  elseif z_lib_name == 'lua-zlib' then
+    function ZlibStream:new(writer, level, method, windowBits)
+      level  = level or ZlibStream.DEFAULT_COMPRESSION
+      method = method or ZlibStream.DEFLATE
+
+      assert(method == ZlibStream.DEFLATE, 'lua-zlib support only deflated method')
+
+      local o = setmetatable({
+        zd = assert(zlib.deflate(level, windowBits));
+        writer = writer;
+      }, self)
+
+      return o
+    end
+
+    function ZlibStream:write(chunk)
+      chunk = assert(self.zd(chunk))
+      self.writer(chunk)
+    end
+
+    function ZlibStream:close()
+      local chunk = self.zd('', 'finish')
+      if chunk and #chunk > 0 then self.writer(chunk) end
+    end
+  end
+end
+
+local Compress = {} do
+  Compress.__index = Compress
+
+  Compress.NO_COMPRESSION      = ZlibStream.NO_COMPRESSION
+  Compress.BEST_SPEED          = ZlibStream.BEST_SPEED
+  Compress.BEST_COMPRESSION    = ZlibStream.BEST_COMPRESSION
+  Compress.DEFAULT_COMPRESSION = ZlibStream.DEFAULT_COMPRESSION
+
+  function Compress:new(options)
+    local compress = {}
+    compress.options = options or {}
+
+    return setmetatable(compress, self)
+  end
+
+  function Compress:processBodyData(data, stayOpen, request, response)
+    local accept_encoding
+
+    if response.headersSended then
+      accept_encoding = response.headers['Content-Encoding'] or ''
+    else
+      local headers = request:headers()
+      accept_encoding = headers and headers['Accept-Encoding'] or ''
+    end
+
+    local accept_gzip = not not accept_encoding:find('gzip', nil, true)
+
+    if accept_gzip and self.options.level ~= ZlibStream.NO_COMPRESSION then
+      local stream = response.compress_stream
+      local buffer = response.compress_buffer
+
+      if not stream then
+        local writer = function (zdata) buffer[#buffer + 1] = zdata end
+        stream, buffer = ZlibStream:new(writer, self.options.level, nil, 31), {}
+      end
+
+      if stayOpen then
+        if data == nil then
+          stream:close()
+          response.compress_stream = nil
+          response.compress_buffer = nil
+        else
+          stream:write(data)
+          response.compress_stream = stream
+          response.compress_buffer = buffer
+        end
+
+        local compressed = table.concat(buffer)
+        for i = 1, #buffer do buffer[i] = nil end
+        if not response.headersSended then
+          response:addHeader('Content-Encoding', 'gzip')
+        end
+
+        return compressed
+      end
+
+      stream:write(data)
+      stream:close()
+      local compressed = table.concat(buffer)
+
+      if #compressed < #data then
+        if not response.headersSended then
+          response:addHeader('Content-Encoding', 'gzip')
+        end
+
+        return compressed
+      end
+    end
+
+    return data
+  end
+end
+
+return Compress