lib/agouti/rack/package_limiter.rb in agouti-0.0.2 vs lib/agouti/rack/package_limiter.rb in agouti-0.0.3

@@ -46,11 +46,11 @@
         status, headers, body = @app.call(env)
 
         set_limit(env)
 
         if enabled?(env)
-          unless headers['Content-Type'] == 'text/html'
+          unless headers['Content-Type'] && headers['Content-Type'].include?('text/html')
             return [204, {}, []]
           end
 
           headers = ::Rack::Utils::HeaderHash.new(headers)
 
@@ -69,30 +69,37 @@
       def get_http_header env, header
         env["HTTP_#{header.upcase.gsub('-', '_')}"]
       end
 
       def enabled? env
-        get_http_header(env, ENABLE_HEADER) && get_http_header(env, ENABLE_HEADER) == 1
+        get_http_header(env, ENABLE_HEADER).to_i == 1
       end
 
       def set_limit env
         @limit = (get_http_header(env, LIMIT_HEADER)) ? get_http_header(env, LIMIT_HEADER).to_i : DEFAULT_LIMIT
       end
 
-      def valid_enable_header? env
-        header = get_http_header(env, ENABLE_HEADER)
+      def valid? env
+        valid_enable_header?(env) && valid_limit_header?(env)
+      end
 
-        (0..1).include?(header) || header.nil?
+      def parseable?(header)
+        (header =~ /^[0-9]+$/) == 0
       end
 
-      def valid_limit_header? env
-        header = get_http_header(env, LIMIT_HEADER)
+      # Defines two methods: valid_enable_header? and valid_limit_header?
+      ['enable', 'limit'].each do |action|
+        define_method("valid_#{action}_header?") do |env|
+          header = get_http_header(env, self.class.const_get("#{action.upcase}_HEADER"))
 
-        (header.kind_of?(Integer) && header > 0) || header.nil?
-      end
+          return true if header.nil?
 
-      def valid? env
-        valid_enable_header?(env) && valid_limit_header?(env)
+          if parseable?(header)
+            header.to_i >= 0
+          else
+            false
+          end
+        end
       end
 
       # Public: class responsible for truncating the gzip stream to a given number of bytes.
       class GzipTruncatedStream < ::Rack::Deflater::GzipStream
         # Public: Constructor.