From b76e76b49d76253dc6bbc2161d88e89df50b06fd Mon Sep 17 00:00:00 2001
From: "A.B" <a.b320012@gmail.com>
Date: Fri, 4 May 2018 13:50:53 -0400
Subject: [PATCH] simplified logic and code quality in decode

---
 lib/jwt.rb        |  1 +
 lib/jwt/decode.rb | 46 ++++++++++++++++++++++++++++++++--------------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/lib/jwt.rb b/lib/jwt.rb
index 75db8379..b95e7f28 100644
--- a/lib/jwt.rb
+++ b/lib/jwt.rb
@@ -42,6 +42,7 @@ def decode(jwt, key = nil, verify = true, options = {}, &keyfinder)
 
     [@payload, @header]
   end
+
   private_class_method
   def verify_signature(&keyfinder)
     @key = find_key(&keyfinder) if keyfinder
diff --git a/lib/jwt/decode.rb b/lib/jwt/decode.rb
index 770a8aa9..ca3bd9d9 100644
--- a/lib/jwt/decode.rb
+++ b/lib/jwt/decode.rb
@@ -15,6 +15,7 @@ def self.base64url_decode(str)
 
     def initialize(jwt, verify)
       @jwt = jwt
+      @segments = jwt.split('.')
       @verify = verify
       @header = ''
       @payload = ''
@@ -22,28 +23,45 @@ def initialize(jwt, verify)
     end
 
     def decode_segments
-      header_segment, payload_segment, crypto_segment = raw_segments
-      @header, @payload = decode_header_and_payload(header_segment, payload_segment)
-      @signature = Decode.base64url_decode(crypto_segment.to_s) if @verify
-      signing_input = [header_segment, payload_segment].join('.')
-      [@header, @payload, @signature, signing_input]
+      validate_segment_count
+      decode_crypto if @verify
+      return_values
     end
 
     private
 
-    def raw_segments
-      segments = @jwt.split('.')
-      required_num_segments = @verify ? [3] : [2, 3]
-      raise(JWT::DecodeError, 'Not enough or too many segments') unless required_num_segments.include? segments.length
-      segments
+    private
+    def validate_segment_count
+      raise(JWT::DecodeError, 'Not enough or too many segments') unless
+        (@verify and @segments.length != 3) ||
+        (@segments.length != 3 or @segments.length != 2)
+    end
+
+    def decode_crypto
+      @signature = Decode.base64url_decode(@segments[2])
+    end
+
+    def return_values
+      [header, payload, @signature, signing_input]
+    end
+
+    def header
+      parse_and_decode @segments[0]
+    end
+
+    def payload
+      parse_and_decode @segments[1]
     end
 
-    def decode_header_and_payload(header_segment, payload_segment)
-      header = JSON.parse(Decode.base64url_decode(header_segment))
-      payload = JSON.parse(Decode.base64url_decode(payload_segment))
-      [header, payload]
+    def signing_input
+      @segments.first(2).join('.')
+    end
+
+    def parse_and_decode segment
+      JSON.parse(Decode.base64url_decode(segment))
     rescue JSON::ParserError
       raise JWT::DecodeError, 'Invalid segment encoding'
     end
+
   end
 end