From 0014e9372b09e05bf4cb1575f51b9bada5c4e5f7 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        |  2 ++
 lib/jwt/decode.rb | 48 +++++++++++++++++++++++++++++++++--------------
 2 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/lib/jwt.rb b/lib/jwt.rb
index 75db8379..a240cde8 100644
--- a/lib/jwt.rb
+++ b/lib/jwt.rb
@@ -42,7 +42,9 @@ 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..204e8a4a 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,26 +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
+    def validate_segment_count
+      raise(JWT::DecodeError, 'Not enough or too many segments') unless
+        (@verify && segment_length != 3) ||
+          (segment_length != 3 || segment_length != 2)
     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 segment_length
+      @segments.count
+    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 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