From 6765476ebeb8db954190c351ab47bb907ff9e2fc Mon Sep 17 00:00:00 2001
From: yash30201 <54198301+yash30201@users.noreply.github.com>
Date: Wed, 27 Dec 2023 18:15:38 +0000
Subject: [PATCH] fix: Allow id token caching incase of gce cred

* Implement logic
* Test logic
---
 src/Credentials/GCECredentials.php       | 13 +++++--
 src/FetchAuthTokenCache.php              |  4 ++
 tests/Credentials/GCECredentialsTest.php | 15 ++++++++
 tests/FetchAuthTokenCacheTest.php        | 49 ++++++++++++++++++++++++
 4 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/src/Credentials/GCECredentials.php b/src/Credentials/GCECredentials.php
index 7849eccfc..87a1b1106 100644
--- a/src/Credentials/GCECredentials.php
+++ b/src/Credentials/GCECredentials.php
@@ -424,7 +424,8 @@ public function fetchAuthToken(callable $httpHandler = null)
         $response = $this->getFromMetadata($httpHandler, $this->tokenUri);
 
         if ($this->targetAudience) {
-            return ['id_token' => $response];
+            $this->lastReceivedToken = ['id_token' => $response];
+            return $this->lastReceivedToken;
         }
 
         if (null === $json = json_decode($response, true)) {
@@ -448,14 +449,20 @@ public function getCacheKey()
     }
 
     /**
-     * @return array{access_token:string,expires_at:int}|null
+     * @return array<mixed>|null
      */
     public function getLastReceivedToken()
     {
         if ($this->lastReceivedToken) {
+            if (array_key_exists('id_token', $this->lastReceivedToken)) {
+                return [
+                    'id_token' => $this->lastReceivedToken['id_token']
+                ];
+            }
+
             return [
                 'access_token' => $this->lastReceivedToken['access_token'],
-                'expires_at' => $this->lastReceivedToken['expires_at'],
+                'expires_at' => $this->lastReceivedToken['expires_at']
             ];
         }
 
diff --git a/src/FetchAuthTokenCache.php b/src/FetchAuthTokenCache.php
index cac1984ab..a76c4f601 100644
--- a/src/FetchAuthTokenCache.php
+++ b/src/FetchAuthTokenCache.php
@@ -237,6 +237,10 @@ public function updateMetadata(
                 $metadata[self::AUTH_METADATA_KEY] = [
                     'Bearer ' . $cached['access_token']
                 ];
+            } elseif (isset($cached['id_token'])) {
+                $metadata[self::AUTH_METADATA_KEY] = [
+                    'Bearer ' . $cached['id_token']
+                ];
             }
         }
 
diff --git a/tests/Credentials/GCECredentialsTest.php b/tests/Credentials/GCECredentialsTest.php
index 695ba0195..e10a005eb 100644
--- a/tests/Credentials/GCECredentialsTest.php
+++ b/tests/Credentials/GCECredentialsTest.php
@@ -263,6 +263,21 @@ public function testGetLastReceivedTokenIsNullByDefault()
         $this->assertNull($creds->getLastReceivedToken());
     }
 
+    public function testGetLastReceivedTokenShouldWorkWithIdToken()
+    {
+        $idToken = '123asdfghjkl';
+        $httpHandler = getHandler([
+            new Response(200, [GCECredentials::FLAVOR_HEADER => 'Google']),
+            new Response(200, [], Utils::streamFor($idToken)),
+        ]);
+        $g = new GCECredentials(null, null, 'https://example.test.com');
+        $g->fetchAuthToken($httpHandler);
+        $this->assertEquals(
+            $idToken,
+            $g->getLastReceivedToken()['id_token']
+        );
+    }
+
     public function testGetClientName()
     {
         $expected = 'foobar';
diff --git a/tests/FetchAuthTokenCacheTest.php b/tests/FetchAuthTokenCacheTest.php
index 21a68e702..09845bb00 100644
--- a/tests/FetchAuthTokenCacheTest.php
+++ b/tests/FetchAuthTokenCacheTest.php
@@ -18,10 +18,13 @@
 namespace Google\Auth\Tests;
 
 use Google\Auth\Cache\MemoryCacheItemPool;
+use Google\Auth\Credentials\GCECredentials;
 use Google\Auth\Credentials\ServiceAccountCredentials;
 use Google\Auth\CredentialsLoader;
 use Google\Auth\FetchAuthTokenCache;
 use Google\Auth\GetUniverseDomainInterface;
+use GuzzleHttp\Psr7\Response;
+use GuzzleHttp\Psr7\Utils;
 use Prophecy\Argument;
 use Prophecy\PhpUnit\ProphecyTrait;
 use RuntimeException;
@@ -235,6 +238,52 @@ public function testUpdateMetadataWithJwtAccess()
         $this->assertNotEquals($metadata, $metadata3);
     }
 
+    public function testUpdateMetadataWithGceCredForIdToken()
+    {
+        $idToken = '123asdfghjkl';
+        $httpHandler = getHandler([
+            new Response(200, [GCECredentials::FLAVOR_HEADER => 'Google']),
+            new Response(200, [], Utils::streamFor($idToken)),
+        ]);
+        $fetcher = new GCECredentials(null, null, 'https://example.test.com');
+        $cache = new MemoryCacheItemPool();
+
+        $cachedFetcher = new FetchAuthTokenCache(
+            $fetcher,
+            null,
+            $cache
+        );
+        $metadata = $cachedFetcher->updateMetadata(
+            [],
+            'http://test-auth-uri',
+            $httpHandler
+        );
+        $this->assertArrayHasKey(
+            CredentialsLoader::AUTH_METADATA_KEY,
+            $metadata
+        );
+
+        $authorization = $metadata[CredentialsLoader::AUTH_METADATA_KEY];
+        $this->assertTrue(is_array($authorization));
+
+        $bearerToken = current($authorization);
+        $this->assertTrue(is_string($bearerToken));
+        $this->assertEquals(0, strpos($bearerToken, 'Bearer '));
+        $token = str_replace('Bearer ', '', $bearerToken);
+
+        $lastReceivedToken = $cachedFetcher->getLastReceivedToken();
+        $this->assertArrayHasKey('id_token', $lastReceivedToken);
+        $this->assertEquals($idToken, $lastReceivedToken['id_token']);
+
+        // Ensure token is cached
+        $metadata2 = $cachedFetcher->updateMetadata([], 'http://test-auth-uri');
+        $this->assertEquals($metadata, $metadata2);
+
+        // Ensure token for different URI is NOT cached
+        $metadata3 = $cachedFetcher->updateMetadata([], 'http://test-auth-uri-2');
+        $this->assertNotEquals($metadata, $metadata3);
+    }
+
     public function testUpdateMetadataWithInvalidFetcher()
     {
         $this->expectException(RuntimeException::class);