Allow for custom headers on encode and decode

lib/jwt.rb

@@ -42,10 +42,10 @@ def self.base64url_encode(str)
     Base64.encode64(str).gsub("+", "-").gsub("/", "_").gsub("\n", "").gsub('=', '')
   end  
 
-  def self.encode(payload, key, algorithm='HS256')
+  def self.encode(payload, key, algorithm='HS256', header_fields={})
     algorithm ||= "none"
     segments = []
-    header = {"typ" => "JWT", "alg" => algorithm}
+    header = {"typ" => "JWT", "alg" => algorithm}.merge(header_fields)
     segments << base64url_encode(header.to_json)
     segments << base64url_encode(payload.to_json)
     signing_input = segments.join('.')
@@ -58,7 +58,7 @@ def self.encode(payload, key, algorithm='HS256')
     segments.join('.')
   end
 
-  def self.decode(jwt, key=nil, verify=true)
+  def self.decode(jwt, key=nil, verify=true, &keyfinder)
     segments = jwt.split('.')
     raise JWT::DecodeError.new("Not enough or too many segments") unless [2,3].include? segments.length
     header_segment, payload_segment, crypto_segment = segments
@@ -73,6 +73,10 @@ def self.decode(jwt, key=nil, verify=true)
     if verify == true
       algo = header['alg']
 
+      if keyfinder
+        key = keyfinder.call(header)
+      end
+
       if ["HS256", "HS384", "HS512"].include?(algo)
         raise JWT::DecodeError.new("Signature verification failed") unless signature == sign_hmac(algo, signing_input, key)
       elsif ["RS256", "RS384", "RS512"].include?(algo)

spec/jwt.rb

@@ -19,6 +19,16 @@
     decoded_payload.should == @payload
   end
 
+  it "encodes and decodes JWTs with custom header fields" do
+    private_key = OpenSSL::PKey::RSA.generate(512) 
+    jwt = JWT.encode(@payload, private_key, "RS256", {"kid" => 'default'})
+    decoded_payload = JWT.decode(jwt) do |header|
+      header["kid"].should == 'default'
+      private_key.public_key
+    end
+    decoded_payload.should == @payload
+  end
+
   it "decodes valid JWTs" do
     example_payload = {"hello" => "world"}
     example_secret = 'secret'