fix: Fixed edge cases in token expiration extraction (#73)

googleapis · Oct 4, 2024 · c6ed7bb · c6ed7bb
1 parent fa27352
commit c6ed7bb
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 6 deletions.
diff --git a/lib/google/cloud/env/compute_metadata.rb b/lib/google/cloud/env/compute_metadata.rb
@@ -806,7 +806,7 @@ def determine_data_lifetime path, data
         #
         def access_token_lifetime data
           json = JSON.parse data rescue nil
-          return 0 unless json&.key? "expires_in"
+          return 0 unless json.respond_to?(:key?) && json.key?("expires_in")
           lifetime = json["expires_in"].to_i - TOKEN_EXPIRY_BUFFER
           lifetime = 0 if lifetime.negative?
           lifetime
@@ -818,9 +818,9 @@ def access_token_lifetime data
         #
         def identity_token_lifetime data
           return 0 unless data =~ /^[\w=-]+\.([\w=-]+)\.[\w=-]+$/
-          base64 = Base64.decode64 Regexp.last_match[1]
+          base64 = Base64.urlsafe_decode64 Regexp.last_match[1]
           json = JSON.parse base64 rescue nil
-          return 0 unless json&.key? "exp"
+          return 0 unless json.respond_to?(:key?) && json&.key?("exp")
           lifetime = json["exp"].to_i - Time.now.to_i - TOKEN_EXPIRY_BUFFER
           lifetime = 0 if lifetime.negative?
           lifetime

diff --git a/test/google/cloud/env/compute_metadata_test.rb b/test/google/cloud/env/compute_metadata_test.rb
@@ -254,9 +254,12 @@
     end
 
     it "gets lifetime for an identity token" do
-      token = {data: "abcdef", exp: Time.now.to_i + 1000}
+      # Note: the encoded token below includes one of the extra Base64
+      # characters (specifically underscore for the urlsafe variant) so this
+      # test also checks that urlsafe is properly used to decode.
+      token = {data: "???", exp: Time.now.to_i + 1000}
       token_json = JSON.generate token
-      token_encoded = Base64.strict_encode64 token_json
+      token_encoded = Base64.urlsafe_encode64 token_json
       full_token = "12345.#{token_encoded}.67890"
       compute_metadata.connection.adapter :test do |stub|
         stub.get "/computeMetadata/v1/instance/service-accounts/12345/identity" do |_env|
@@ -267,6 +270,28 @@
       expected_expiry = Process.clock_gettime(Process::CLOCK_MONOTONIC) + 790
       assert_in_delta expected_expiry, compute_metadata.expiration_time_of("instance/service-accounts/12345/identity"), 0.1
     end
+
+    it "defaults to zero lifetime when failing to parse access token response" do
+      compute_metadata.connection.adapter :test do |stub|
+        stub.get "/computeMetadata/v1/instance/service-accounts/12345/token" do |_env|
+          [404, flavor_header, '"Not found"']
+        end
+      end
+      assert_nil compute_metadata.lookup "instance/service-accounts/12345/token"
+      expected_expiry = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      assert_in_delta expected_expiry, compute_metadata.expiration_time_of("instance/service-accounts/12345/token")
+    end
+
+    it "defaults to zero lifetime when failing to parse identity token response" do
+      compute_metadata.connection.adapter :test do |stub|
+        stub.get "/computeMetadata/v1/instance/service-accounts/12345/identity" do |_env|
+          [404, flavor_header, '"Not found"']
+        end
+      end
+      assert_nil compute_metadata.lookup "instance/service-accounts/12345/identity"
+      expected_expiry = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      assert_in_delta expected_expiry, compute_metadata.expiration_time_of("instance/service-accounts/12345/identity")
+    end
   end
 
   describe "check_existence" do

diff --git a/test/google/cloud/env/lazy_value_test.rb b/test/google/cloud/env/lazy_value_test.rb
@@ -522,7 +522,7 @@ class MyError < StandardError
 
     it "reflects computing state" do
       cache = Google::Cloud::Env::LazyValue.new do
-        sleep 0.2
+        sleep 0.3
         2
       end
       start_time = nil