chore: Refactor AuthTokenMiddleware's token logic

googleapis · Oct 26, 2023 · c9345e4 · c9345e4
1 parent b271a52
commit c9345e4
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 15 deletions.
diff --git a/src/Middleware/AuthTokenMiddleware.php b/src/Middleware/AuthTokenMiddleware.php
@@ -17,8 +17,10 @@
 
 namespace Google\Auth\Middleware;
 
+use Google\Auth\FetchAuthTokenCache;
 use Google\Auth\FetchAuthTokenInterface;
 use Google\Auth\GetQuotaProjectInterface;
+use Google\Auth\UpdateMetadataInterface;
 use Psr\Http\Message\RequestInterface;
 
 /**
@@ -99,7 +101,13 @@ public function __invoke(callable $handler)
                 return $handler($request, $options);
             }
 
-            $request = $request->withHeader('authorization', 'Bearer ' . $this->fetchToken());
+            foreach ($this->fetchAuthHeaders() as $key => $value) {
+                if ($key == 'authorization') {
+                    $request = $request->withHeader($key, $value);
+                } else {
+                    $request = $request->withAddedHeader($key, $value);
+                }
+            }
 
             if ($quotaProject = $this->getQuotaProject()) {
                 $request = $request->withHeader(
@@ -113,32 +121,54 @@ public function __invoke(callable $handler)
     }
 
     /**
-     * Call fetcher to fetch the token.
+     * Fetch auth headers.
      *
-     * @return string|null
+     * @return array
      */
-    private function fetchToken()
+    private function fetchAuthHeaders()
     {
-        $auth_tokens = (array) $this->fetcher->fetchAuthToken($this->httpHandler);
+        $authHeaders = [];
+        $authTokens = [];
+
+        // We need to find the actual fetcher incase of a cache wrapper
+        // so that we can avoid the exception case where actual fetcher
+        // does not implements UpdateMetadataInterface with cache wrapper's
+        // `updateMetadata` being called.
+        $actualFetcher = $this->fetcher;
+        if ($actualFetcher instanceof FetchAuthTokenCache) {
+            $actualFetcher = $actualFetcher->getFetcher();
+        }
+
+        if ($actualFetcher instanceof UpdateMetadataInterface) {
+            $headers = $this->fetcher->updateMetadata([], null, $this->httpHandler);
+            if (array_key_exists('authorization', $headers)) {
+                $authHeaders = $headers;
+            }
+        } else {
+            $authTokens = (array) $this->fetcher->fetchAuthToken($this->httpHandler);
+            if (array_key_exists('access_token', $authTokens)) {
+                $authHeaders = ['authorization' => 'Bearer ' . $authTokens['access_token']];
+            } elseif (array_key_exists('id_token', $authTokens)) {
+                $authHeaders = ['authorization' => 'Bearer ' . $authTokens['id_token']];
+            }
+        }
+
+        if (!empty($authHeaders)) {
+            if (empty($authTokens)) {
+                $authTokens = $this->fetcher->getLastReceivedToken();
+            }
 
-        if (array_key_exists('access_token', $auth_tokens)) {
             // notify the callback if applicable
-            if ($this->tokenCallback) {
+            if (array_key_exists('access_token', $authTokens) && $this->tokenCallback) {
                 call_user_func(
                     $this->tokenCallback,
                     $this->fetcher->getCacheKey(),
-                    $auth_tokens['access_token']
+                    $authTokens['access_token']
                 );
             }
-
-            return $auth_tokens['access_token'];
         }
 
-        if (array_key_exists('id_token', $auth_tokens)) {
-            return $auth_tokens['id_token'];
-        }
-
-        return null;
+        return $authHeaders;
     }
 
     /**

diff --git a/tests/Middleware/AuthTokenMiddlewareTest.php b/tests/Middleware/AuthTokenMiddlewareTest.php
@@ -20,6 +20,7 @@
 use Google\Auth\FetchAuthTokenCache;
 use Google\Auth\Middleware\AuthTokenMiddleware;
 use Google\Auth\Tests\BaseTest;
+use Google\Auth\UpdateMetadataInterface;
 use GuzzleHttp\Handler\MockHandler;
 use GuzzleHttp\Psr7\Response;
 use Prophecy\Argument;
@@ -281,6 +282,50 @@ public function testShouldNotifyTokenCallback(callable $tokenCallback)
         $this->assertTrue(MiddlewareCallback::$called);
     }
 
+    public function testAddAuthHeadersFromUpdateMetadata()
+    {
+        $authResult = [
+            'authorization' => 'Bearer 1/abcdef1234567890',
+        ];
+
+        $mockFetcher = $this->prophesize('Google\Auth\FetchAuthTokenInterface');
+        $mockFetcher->willImplement(UpdateMetadataInterface::class);
+        $mockFetcher->updateMetadata(Argument::cetera())
+            ->shouldBeCalledTimes(1)
+            ->willReturn($authResult);
+        $mockFetcher->getLastReceivedToken()
+            ->willReturn(['access_token' => '1/abcdef1234567890']);
+        $this->mockRequest->withHeader('authorization', $authResult['authorization'])
+            ->shouldBeCalledTimes(1)
+            ->willReturn($this->mockRequest->reveal());
+
+        $this->runTestCase($mockFetcher->reveal(), true);
+    }
+
+    public function testOverlappingAddAuthHeadersFromUpdateMetadata()
+    {
+        $authHeaders = [
+            'authorization' => 'Bearer 1/abcdef1234567890',
+            'x-goog-api-client' => 'extra-value'
+        ];
+
+        $mockFetcher = $this->prophesize('Google\Auth\FetchAuthTokenInterface');
+        $mockFetcher->willImplement(UpdateMetadataInterface::class);
+        $mockFetcher->updateMetadata(Argument::cetera())
+            ->shouldBeCalledTimes(1)
+            ->willReturn($authHeaders);
+        $mockFetcher->getLastReceivedToken()
+            ->willReturn(['access_token' => '1/abcdef1234567890']);
+        $this->mockRequest->withHeader('authorization', $authHeaders['authorization'])
+            ->shouldBeCalledTimes(1)
+            ->willReturn($this->mockRequest->reveal());
+        $this->mockRequest->withAddedHeader('x-goog-api-client', $authHeaders['x-goog-api-client'])
+            ->shouldBeCalledTimes(1)
+            ->willReturn($this->mockRequest->reveal());
+
+        $this->runTestCase($mockFetcher->reveal(), true);
+    }
+
     private function runTestCase($fetcher, bool $isGoogleAuth)
     {
         $authType = $isGoogleAuth ? 'google_auth' :'not_google_auth';