Extend Base64() function support

3 files changed, 89 insertions, 8 deletions

diff --git a/README.markdown b/README.markdown
index c2d8388..602eaf7 100644
--- a/README.markdown
+++ b/README.markdown
@@ -146,7 +146,36 @@ Converts any object to an array containing that object. Empty argument lists are
 
 #### `base64`
 
-Converts a string to and from base64 encoding. Requires an action ('encode', 'decode') and either a plain or base64-encoded string. *Type*: rvalue.
+Converts a string to and from base64 encoding.
+Requires an `action` ('encode', 'decode') and either a plain or base64-encoded `string`,
+and an optional `method` ('default', 'strict', 'urlsafe')
+
+for backward compatibility, `metohd` will be set as `default` if not specified.
+
+*Examples:*
+~~~
+base64('encode', 'hello')
+base64('encode', 'hello', 'default')
+# return: "aGVsbG8=\n"
+
+base64('encode', 'hello', 'strict')
+# return: "aGVsbG8="
+
+base64('decode', 'aGVsbG8=')
+base64('decode', 'aGVsbG8=\n')
+base64('decode', 'aGVsbG8=', 'default')
+base64('decode', 'aGVsbG8=\n', 'default')
+base64('decode', 'aGVsbG8=', 'strict')
+# return: "hello"
+
+base64('encode', 'https://puppetlabs.com', 'urlsafe')
+# return: "aHR0cHM6Ly9wdXBwZXRsYWJzLmNvbQ=="
+
+base64('decode', 'aHR0cHM6Ly9wdXBwZXRsYWJzLmNvbQ==', 'urlsafe')
+# return: "https://puppetlabs.com"
+~~~
+
+*Type*: rvalue.
 
 #### `basename`
 
diff --git a/lib/puppet/parser/functions/base64.rb b/lib/puppet/parser/functions/base64.rb
index 617ba31..a8998f2 100644
--- a/lib/puppet/parser/functions/base64.rb
+++ b/lib/puppet/parser/functions/base64.rb
@@ -6,14 +6,19 @@ module Puppet::Parser::Functions
 
     Usage:
 
-      $encodestring = base64('encode','thestring')
-      $decodestring = base64('decode','dGhlc3RyaW5n')
+      $encodestring = base64('encode', 'thestring')
+      $decodestring = base64('decode', 'dGhlc3RyaW5n')
+
+      # explicitly define encode/decode method: default, strict, urlsafe
+      $method = 'default'
+      $encodestring = base64('encode', 'thestring', $method)
+      $decodestring = base64('decode', 'dGhlc3RyaW5n', $method)
 
     ENDHEREDOC
 
     require 'base64'
 
-    raise Puppet::ParseError, ("base64(): Wrong number of arguments (#{args.length}; must be = 2)") unless args.length == 2
+    raise Puppet::ParseError, ("base64(): Wrong number of arguments (#{args.length}; must be >= 2)") unless args.length >= 2
 
     actions = ['encode','decode']
 
@@ -25,11 +30,37 @@ module Puppet::Parser::Functions
       raise Puppet::ParseError, ("base64(): the second argument must be a string to base64")
     end
 
+    method = ['default','strict','urlsafe']
+
+    if args.length <= 2
+      chosenMethod = 'default'
+    else
+      chosenMethod = args[2]
+    end
+
+    unless method.include?(chosenMethod)
+      raise Puppet::ParseError, ("base64(): the third argument must be one of 'default', 'strict', or 'urlsafe'")
+    end
+
     case args[0]
       when 'encode'
-        result = Base64.encode64(args[1])
+        case chosenMethod
+          when 'default'
+            result = Base64.encode64(args[1])
+          when 'strict'
+            result = Base64.strict_encode64(args[1])
+          when 'urlsafe'
+            result = Base64.urlsafe_encode64(args[1])
+        end
       when 'decode'
-        result = Base64.decode64(args[1])
+        case chosenMethod
+          when 'default'
+            result = Base64.decode64(args[1])
+          when 'strict'
+            result = Base64.strict_decode64(args[1])
+          when 'urlsafe'
+            result = Base64.urlsafe_decode64(args[1])
+        end
     end
 
     return result
diff --git a/spec/functions/base64_spec.rb b/spec/functions/base64_spec.rb
index c529e9e..842a37a 100755
--- a/spec/functions/base64_spec.rb
+++ b/spec/functions/base64_spec.rb
@@ -4,13 +4,34 @@ describe 'base64' do
   it { is_expected.not_to eq(nil) }
   it { is_expected.to run.with_params().and_raise_error(Puppet::ParseError) }
   it { is_expected.to run.with_params("one").and_raise_error(Puppet::ParseError) }
-  it { is_expected.to run.with_params("one", "two", "three").and_raise_error(Puppet::ParseError) }
   it { is_expected.to run.with_params("one", "two").and_raise_error(Puppet::ParseError, /first argument must be one of/) }
   it { is_expected.to run.with_params("encode", ["two"]).and_raise_error(Puppet::ParseError, /second argument must be a string/) }
   it { is_expected.to run.with_params("encode", 2).and_raise_error(Puppet::ParseError, /second argument must be a string/) }
+  it { is_expected.to run.with_params("encode", "thestring", "three").and_raise_error(Puppet::ParseError, /third argument must be one of/) }
+  it { is_expected.to run.with_params("decode", "dGhlc3RyaW5n\n", "strict").and_raise_error(ArgumentError) }
 
   it { is_expected.to run.with_params("encode", "thestring").and_return("dGhlc3RyaW5n\n") }
-  it { is_expected.to run.with_params("encode", "a very long string that will cause the base64 encoder to produce output with multiple lines").and_return("YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0\nIGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5l\ncw==\n") }
   it { is_expected.to run.with_params("decode", "dGhlc3RyaW5n").and_return("thestring") }
   it { is_expected.to run.with_params("decode", "dGhlc3RyaW5n\n").and_return("thestring") }
+
+  it { is_expected.to run.with_params("encode", "thestring", "default").and_return("dGhlc3RyaW5n\n") }
+  it { is_expected.to run.with_params("decode", "dGhlc3RyaW5n", "default").and_return("thestring") }
+  it { is_expected.to run.with_params("decode", "dGhlc3RyaW5n\n", "default").and_return("thestring") }
+
+  it { is_expected.to run.with_params("encode", "thestring", "strict").and_return("dGhlc3RyaW5n") }
+  it { is_expected.to run.with_params("decode", "dGhlc3RyaW5n", "strict").and_return("thestring") }
+
+  it { is_expected.to run.with_params("encode", "a very long string that will cause the base64 encoder to produce output with multiple lines").and_return("YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0\nIGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5l\ncw==\n") }
+  it { is_expected.to run.with_params("decode", "YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0\nIGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5l\ncw==\n").and_return("a very long string that will cause the base64 encoder to produce output with multiple lines") }
+  it { is_expected.to run.with_params("decode", "YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0IGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5lcw==").and_return("a very long string that will cause the base64 encoder to produce output with multiple lines") }
+
+  it { is_expected.to run.with_params("encode", "a very long string that will cause the base64 encoder to produce output with multiple lines", "strict").and_return("YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0IGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5lcw==") }
+  it { is_expected.to run.with_params("decode", "YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0IGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5lcw==").and_return("a very long string that will cause the base64 encoder to produce output with multiple lines") }
+  it { is_expected.to run.with_params("decode", "YSB2ZXJ5IGxvbmcgc3RyaW5nIHRoYXQgd2lsbCBjYXVzZSB0aGUgYmFzZTY0IGVuY29kZXIgdG8gcHJvZHVjZSBvdXRwdXQgd2l0aCBtdWx0aXBsZSBsaW5lcw==", "strict").and_return("a very long string that will cause the base64 encoder to produce output with multiple lines") }
+
+  it { is_expected.to run.with_params("encode", "https://www.google.com.tw/?gws_rd=ssl#q=hello+world", "urlsafe").and_return("aHR0cHM6Ly93d3cuZ29vZ2xlLmNvbS50dy8_Z3dzX3JkPXNzbCNxPWhlbGxvK3dvcmxk") }
+  it { is_expected.to run.with_params("decode", "aHR0cHM6Ly93d3cuZ29vZ2xlLmNvbS50dy8_Z3dzX3JkPXNzbCNxPWhlbGxvK3dvcmxk", "urlsafe").and_return("https://www.google.com.tw/?gws_rd=ssl#q=hello+world") }
+
+  it { is_expected.to run.with_params("encode", "https://github.com/puppetlabs/puppetlabs-stdlib/pulls?utf8=%E2%9C%93&q=is%3Apr+is%3Aopen+Add", "urlsafe").and_return("aHR0cHM6Ly9naXRodWIuY29tL3B1cHBldGxhYnMvcHVwcGV0bGFicy1zdGRsaWIvcHVsbHM_dXRmOD0lRTIlOUMlOTMmcT1pcyUzQXByK2lzJTNBb3BlbitBZGQ=") }
+  it { is_expected.to run.with_params("decode", "aHR0cHM6Ly9naXRodWIuY29tL3B1cHBldGxhYnMvcHVwcGV0bGFicy1zdGRsaWIvcHVsbHM_dXRmOD0lRTIlOUMlOTMmcT1pcyUzQXByK2lzJTNBb3BlbitBZGQ=", "urlsafe").and_return("https://github.com/puppetlabs/puppetlabs-stdlib/pulls?utf8=%E2%9C%93&q=is%3Apr+is%3Aopen+Add") }
 end