Workaround #2668 - Compress facts so that the request size limit triggers less often

This is not the right fix, but more a hackish workaround.

Since 0.25, the facts are transmitted as GET parameters when a
node asks for a catalog. Most proxies or webserver have a size limit
which is sometimes reached. In this case the request is denied
and the node can't get its catalog.

The idea is to compress facts (some non-scientific studies show a
57% fact size decrease for an average node) when transmitting
those when asking for a catalog.

Signed-off-by: Brice Figureau <brice-puppet@daysofwonder.com>

4 files changed, 140 insertions, 5 deletions

diff --git a/lib/puppet/configurer/fact_handler.rb b/lib/puppet/configurer/fact_handler.rb
index 43e9f35f4..8e0fef71d 100644
--- a/lib/puppet/configurer/fact_handler.rb
+++ b/lib/puppet/configurer/fact_handler.rb
@@ -29,11 +29,11 @@ module Puppet::Configurer::FactHandler
         #format = facts.class.default_format
 
         # Hard-code yaml, because I couldn't get marshal to work.
-        format = :yaml
+        format = :b64_zlib_yaml
 
         text = facts.render(format)
 
-        return {:facts_format => format, :facts => CGI.escape(text)}
+        return {:facts_format => :b64_zlib_yaml, :facts => CGI.escape(text)}
     end
 
     # Retrieve facts from the central server.
diff --git a/lib/puppet/network/formats.rb b/lib/puppet/network/formats.rb
index df6ef399c..010c23521 100644
--- a/lib/puppet/network/formats.rb
+++ b/lib/puppet/network/formats.rb
@@ -39,6 +39,56 @@ Puppet::Network::FormatHandler.create(:yaml, :mime => "text/yaml") do
     end
 end
 
+# This is a "special" format which is used for the moment only when sending facts
+# as REST GET parameters (see Puppet::Configurer::FactHandler).
+# This format combines a yaml serialization, then zlib compression and base64 encoding.
+Puppet::Network::FormatHandler.create(:b64_zlib_yaml, :mime => "text/b64_zlib_yaml") do
+    require 'base64'
+    require 'zlib'
+
+    def intern(klass, text)
+        decode(text)
+    end
+
+    def intern_multiple(klass, text)
+        decode(text)
+    end
+
+    def render(instance)
+        yaml = instance.to_yaml
+
+        yaml = encode(fixup(yaml)) unless yaml.nil?
+        yaml
+    end
+
+    def render_multiple(instances)
+        yaml = instances.to_yaml
+
+        yaml = encode(fixup(yaml)) unless yaml.nil?
+        yaml
+    end
+
+    # Because of yaml issue in ruby 1.8.1...
+    def supported?(klass)
+        RUBY_VERSION != '1.8.1'
+    end
+
+    # fixup invalid yaml as per:
+    # http://redmine.ruby-lang.org/issues/show/1331
+    def fixup(yaml)
+        yaml.gsub!(/((?:&id\d+\s+)?!ruby\/object:.*?)\s*\?/) { "? #{$1}" }
+        yaml
+    end
+
+    def encode(text)
+        Base64.encode64(Zlib::Deflate.deflate(text, Zlib::BEST_COMPRESSION))
+    end
+
+    def decode(yaml)
+        YAML.load(Zlib::Inflate.inflate(Base64.decode64(yaml)))
+    end
+end
+
 
 Puppet::Network::FormatHandler.create(:marshal, :mime => "text/marshal") do
     # Marshal doesn't need the class name; it's serialized.
diff --git a/spec/unit/configurer/fact_handler.rb b/spec/unit/configurer/fact_handler.rb
index 0c4af9554..ec60c6dcd 100755
--- a/spec/unit/configurer/fact_handler.rb
+++ b/spec/unit/configurer/fact_handler.rb
@@ -102,7 +102,7 @@ describe Puppet::Configurer::FactHandler do
 
         @facthandler.expects(:find_facts).returns facts
 
-        @facthandler.facts_for_uploading.should == {:facts_format => :yaml, :facts => text}
+        @facthandler.facts_for_uploading.should == {:facts_format => :b64_zlib_yaml, :facts => text}
     end
 
     it "should properly accept facts containing a '+'" do
@@ -112,12 +112,12 @@ describe Puppet::Configurer::FactHandler do
 
         @facthandler.expects(:find_facts).returns facts
 
-        @facthandler.facts_for_uploading.should == {:facts_format => :yaml, :facts => text}
+        @facthandler.facts_for_uploading.should == {:facts_format => :b64_zlib_yaml, :facts => text}
     end
 
     it "should hard-code yaml as the serialization" do
         facts = stub 'facts'
-        facts.expects(:render).with(:yaml).returns "my text"
+        facts.expects(:render).with(:b64_zlib_yaml).returns "my text"
         text = CGI.escape("my text")
 
         @facthandler.expects(:find_facts).returns facts
diff --git a/spec/unit/network/formats.rb b/spec/unit/network/formats.rb
index de2e0afe3..b1ef9ec53 100755
--- a/spec/unit/network/formats.rb
+++ b/spec/unit/network/formats.rb
@@ -90,6 +90,91 @@ describe "Puppet Network Format" do
         end
     end
 
+    describe "base64 compressed yaml" do
+        before do
+            @yaml = Puppet::Network::FormatHandler.format(:b64_zlib_yaml)
+        end
+
+        it "should have its mime type set to text/b64_zlib_yaml" do
+            @yaml.mime.should == "text/b64_zlib_yaml"
+        end
+
+        it "should render by calling 'to_yaml' on the instance" do
+            instance = mock 'instance'
+            instance.expects(:to_yaml).returns "foo"
+            @yaml.render(instance)
+        end
+
+        it "should fixup generated yaml on render" do
+            instance = mock 'instance', :to_yaml => "foo"
+
+            @yaml.expects(:fixup).with("foo").returns "bar"
+
+            @yaml.render(instance)
+        end
+
+        it "should encode generated yaml on render" do
+            instance = mock 'instance', :to_yaml => "foo"
+
+            @yaml.expects(:encode).with("foo").returns "bar"
+
+            @yaml.render(instance).should == "bar"
+        end
+
+        it "should render multiple instances by calling 'to_yaml' on the array" do
+            instances = [mock('instance')]
+            instances.expects(:to_yaml).returns "foo"
+            @yaml.render_multiple(instances)
+        end
+
+        it "should fixup generated yaml on render" do
+            instances = [mock('instance')]
+            instances.stubs(:to_yaml).returns "foo"
+
+            @yaml.expects(:fixup).with("foo").returns "bar"
+
+            @yaml.render(instances)
+        end
+
+        it "should encode generated yaml on render" do
+            instances = [mock('instance')]
+            instances.stubs(:to_yaml).returns "foo"
+
+            @yaml.expects(:encode).with("foo").returns "bar"
+
+            @yaml.render(instances).should == "bar"
+        end
+
+        it "should intern by calling decode" do
+            text = "foo"
+            @yaml.expects(:decode).with("foo").returns "bar"
+            @yaml.intern(String, text).should == "bar"
+        end
+
+        it "should intern multiples by calling 'decode'" do
+            text = "foo"
+            @yaml.expects(:decode).with("foo").returns "bar"
+            @yaml.intern_multiple(String, text).should == "bar"
+        end
+
+        it "should decode by base64 decoding, uncompressing and Yaml loading" do
+            Base64.expects(:decode64).with("zorg").returns "foo"
+            Zlib::Inflate.expects(:inflate).with("foo").returns "baz"
+            YAML.expects(:load).with("baz").returns "bar"
+            @yaml.decode("zorg").should == "bar"
+        end
+
+        it "should encode by compressing and base64 encoding" do
+            Zlib::Deflate.expects(:deflate).with("foo", Zlib::BEST_COMPRESSION).returns "bar"
+            Base64.expects(:encode64).with("bar").returns "baz"
+            @yaml.encode("foo").should == "baz"
+        end
+
+        it "should fixup incorrect yaml to correct" do
+            @yaml.fixup("&id004 !ruby/object:Puppet::Relationship ?").should == "? &id004 !ruby/object:Puppet::Relationship"
+        end
+    end
+
     it "should include a marshal format" do
         Puppet::Network::FormatHandler.format(:marshal).should_not be_nil
     end