Fixing 2855 Inadvertent zlib dependency

Disables the b64_zlib_yaml format if zlib cannot be loaded.

I've added a --no-zlib to make it possible to test this on a single
machine, but it might also be useful if someone finds themselves failing
to connect to a server that doesn't have zlib installed.

FactHandler' format is still hard-coded to YAML rather than using
facts.class.default_format

Signed-off-by: Jesse Wolfe <jes5199@gmail.com>

6 files changed, 81 insertions, 9 deletions

diff --git a/lib/puppet/configurer/fact_handler.rb b/lib/puppet/configurer/fact_handler.rb
index cbb627680..40e79b6c6 100644
--- a/lib/puppet/configurer/fact_handler.rb
+++ b/lib/puppet/configurer/fact_handler.rb
@@ -27,12 +27,15 @@ module Puppet::Configurer::FactHandler
         facts = find_facts
         #format = facts.class.default_format
 
-        # Hard-code yaml, because I couldn't get marshal to work.
-        format = :b64_zlib_yaml
+        if facts.support_format?(:b64_zlib_yaml)
+            format = :b64_zlib_yaml
+        else
+            format = :yaml
+        end
 
         text = facts.render(format)
 
-        return {:facts_format => :b64_zlib_yaml, :facts => CGI.escape(text)}
+        return {:facts_format => format, :facts => CGI.escape(text)}
     end
 
     # Retrieve facts from the central server.
diff --git a/lib/puppet/defaults.rb b/lib/puppet/defaults.rb
index 67d4a42a5..0153957ef 100644
--- a/lib/puppet/defaults.rb
+++ b/lib/puppet/defaults.rb
@@ -196,7 +196,10 @@ module Puppet
         :config_version => ["", "How to determine the configuration version.  By default, it will be the
             time that the configuration is parsed, but you can provide a shell script to override how the
             version is determined.  The output of this script will be added to every log message in the
-            reports, allowing you to correlate changes on your hosts to the source version on the server."]
+            reports, allowing you to correlate changes on your hosts to the source version on the server."],
+        :zlib => [true, 
+            "Boolean; whether to use the zlib library",
+        ]
     )
 
     hostname = Facter["hostname"].value
diff --git a/lib/puppet/feature/zlib.rb b/lib/puppet/feature/zlib.rb
new file mode 100644
index 000000000..ddd65b234
--- /dev/null
+++ b/lib/puppet/feature/zlib.rb
@@ -0,0 +1,6 @@
+require 'puppet/util/feature'
+
+# We want this to load if possible, but it's not automatically
+# required.
+Puppet.features.rubygems?
+Puppet.features.add(:zlib, :libs => %{zlib})
diff --git a/lib/puppet/network/formats.rb b/lib/puppet/network/formats.rb
index 010c23521..a98dcbcc5 100644
--- a/lib/puppet/network/formats.rb
+++ b/lib/puppet/network/formats.rb
@@ -44,7 +44,18 @@ end
 # 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 use_zlib?
+        Puppet.features.zlib? && Puppet[:zlib]
+    end
+
+    def requiring_zlib
+        if use_zlib?
+            yield
+        else
+            raise Puppet::Error, "the zlib library is not installed or is disabled."
+        end
+    end
 
     def intern(klass, text)
         decode(text)
@@ -70,7 +81,7 @@ Puppet::Network::FormatHandler.create(:b64_zlib_yaml, :mime => "text/b64_zlib_ya
 
     # Because of yaml issue in ruby 1.8.1...
     def supported?(klass)
-        RUBY_VERSION != '1.8.1'
+        RUBY_VERSION != '1.8.1' and use_zlib?
     end
 
     # fixup invalid yaml as per:
@@ -81,11 +92,15 @@ Puppet::Network::FormatHandler.create(:b64_zlib_yaml, :mime => "text/b64_zlib_ya
     end
 
     def encode(text)
-        Base64.encode64(Zlib::Deflate.deflate(text, Zlib::BEST_COMPRESSION))
+        requiring_zlib do 
+            Base64.encode64(Zlib::Deflate.deflate(text, Zlib::BEST_COMPRESSION))
+        end
     end
 
     def decode(yaml)
-        YAML.load(Zlib::Inflate.inflate(Base64.decode64(yaml)))
+        requiring_zlib do 
+            YAML.load(Zlib::Inflate.inflate(Base64.decode64(yaml)))
+        end
     end
 end
 
diff --git a/spec/unit/configurer/fact_handler.rb b/spec/unit/configurer/fact_handler.rb
index ec60c6dcd..290f1acfc 100755
--- a/spec/unit/configurer/fact_handler.rb
+++ b/spec/unit/configurer/fact_handler.rb
@@ -97,6 +97,7 @@ describe Puppet::Configurer::FactHandler do
     # I couldn't get marshal to work for this, only yaml, so we hard-code yaml.
     it "should serialize and CGI escape the fact values for uploading" do
         facts = stub 'facts'
+        facts.expects(:support_format?).with(:b64_zlib_yaml).returns true
         facts.expects(:render).returns "my text"
         text = CGI.escape("my text")
 
@@ -107,6 +108,7 @@ describe Puppet::Configurer::FactHandler do
 
     it "should properly accept facts containing a '+'" do
         facts = stub 'facts'
+        facts.expects(:support_format?).with(:b64_zlib_yaml).returns true
         facts.expects(:render).returns "my+text"
         text = "my%2Btext"
 
@@ -115,8 +117,9 @@ describe Puppet::Configurer::FactHandler do
         @facthandler.facts_for_uploading.should == {:facts_format => :b64_zlib_yaml, :facts => text}
     end
 
-    it "should hard-code yaml as the serialization" do
+    it "use compressed yaml as the serialization if zlib is supported" do
         facts = stub 'facts'
+        facts.expects(:support_format?).with(:b64_zlib_yaml).returns true
         facts.expects(:render).with(:b64_zlib_yaml).returns "my text"
         text = CGI.escape("my text")
 
@@ -125,6 +128,17 @@ describe Puppet::Configurer::FactHandler do
         @facthandler.facts_for_uploading
     end
 
+    it "should use yaml as the serialization if zlib is not supported" do
+        facts = stub 'facts'
+        facts.expects(:support_format?).with(:b64_zlib_yaml).returns false
+        facts.expects(:render).with(:yaml).returns "my text"
+        text = CGI.escape("my text")
+
+        @facthandler.expects(:find_facts).returns facts
+
+        @facthandler.facts_for_uploading
+    end
+
     describe "when reloading Facter" do
         before do
             Facter.stubs(:clear)
diff --git a/spec/unit/network/formats.rb b/spec/unit/network/formats.rb
index b1ef9ec53..a241306a2 100755
--- a/spec/unit/network/formats.rb
+++ b/spec/unit/network/formats.rb
@@ -91,6 +91,9 @@ describe "Puppet Network Format" do
     end
 
     describe "base64 compressed yaml" do
+        yaml = Puppet::Network::FormatHandler.format(:b64_zlib_yaml)
+        confine "We must have zlib" => Puppet.features.zlib?
+
         before do
             @yaml = Puppet::Network::FormatHandler.format(:b64_zlib_yaml)
         end
@@ -173,6 +176,34 @@ describe "Puppet Network Format" do
         it "should fixup incorrect yaml to correct" do
             @yaml.fixup("&id004 !ruby/object:Puppet::Relationship ?").should == "? &id004 !ruby/object:Puppet::Relationship"
         end
+
+        describe "when zlib is disabled" do
+            before do
+                Puppet[:zlib] = false
+            end
+
+            it "use_zlib? should return false" do
+                @yaml.use_zlib?.should == false  
+            end
+
+            it "should refuse to encode" do
+                lambda{ @yaml.encode("foo") }.should raise_error
+            end
+
+            it "should refuse to decode" do
+                lambda{ @yaml.decode("foo") }.should raise_error
+            end
+        end
+
+        describe "when zlib is not installed" do
+            it "use_zlib? should return false" do
+                Puppet[:zlib] = true
+                Puppet.features.expects(:zlib?).returns(false)
+
+                @yaml.use_zlib?.should == false
+            end
+        end
+
     end
 
     it "should include a marshal format" do