Improve logic

googleapis · Oct 31, 2023 · 664b4e2 · 664b4e2
1 parent 5392f7b
commit 664b4e2
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 58 deletions.
diff --git a/src/Middleware/AuthTokenMiddleware.php b/src/Middleware/AuthTokenMiddleware.php
@@ -21,6 +21,7 @@
 use Google\Auth\FetchAuthTokenInterface;
 use Google\Auth\GetQuotaProjectInterface;
 use Google\Auth\UpdateMetadataInterface;
+use GuzzleHttp\Psr7\Utils;
 use Psr\Http\Message\RequestInterface;
 
 /**
@@ -104,11 +105,18 @@ public function __invoke(callable $handler)
                 return $handler($request, $options);
             }
 
-            foreach ($this->fetchAuthHeaders() as $key => $value) {
-                if ($key == 'authorization') {
-                    $request = $request->withHeader($key, $value);
-                } else {
-                    $request = $request->withAddedHeader($key, $value);
+            if (!$this->fetcher instanceof UpdateMetadataInterface ||
+                ($this->fetcher instanceof FetchAuthTokenCache && !$this->fetcher->getFetcher() instanceof UpdateMetadataInterface)
+            ) {
+                $request =$request->withHeader('authorization', 'Bearer ' . $this->fetchToken());
+            } else {
+                $headers = $this->fetcher->updateMetadata($request->getHeaders(), null, $this->httpHandler);
+                $request = Utils::modifyRequest($request, ['set_headers' => $headers]);
+                if ($this->tokenCallback) {
+                    $auth_tokens = $this->fetcher->fetchAuthToken($this->httpHandler);
+                    if (array_key_exists('access_token', $auth_tokens)) {
+                        call_user_func($this->tokenCallback, $this->fetcher->getCacheKey(), $auth_tokens['access_token']);
+                    }
                 }
             }
 
@@ -124,54 +132,32 @@ public function __invoke(callable $handler)
     }
 
     /**
-     * Fetch auth headers.
+     * Call fetcher to fetch the token.
      *
-     * @return array<string, string[]|string>
+     * @return string|null
      */
-    private function fetchAuthHeaders()
+    private function fetchToken()
     {
-        $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();
-            }
+        $auth_tokens = (array) $this->fetcher->fetchAuthToken($this->httpHandler);
 
+        if (array_key_exists('access_token', $auth_tokens)) {
             // notify the callback if applicable
-            if (array_key_exists('access_token', $authTokens) && $this->tokenCallback) {
+            if ($this->tokenCallback) {
                 call_user_func(
                     $this->tokenCallback,
                     $this->fetcher->getCacheKey(),
-                    $authTokens['access_token']
+                    $auth_tokens['access_token']
                 );
             }
+
+            return $auth_tokens['access_token'];
         }
 
-        return $authHeaders;
+        if (array_key_exists('id_token', $auth_tokens)) {
+            return $auth_tokens['id_token'];
+        }
+
+        return null;
     }
 
     /**

diff --git a/tests/FetchAuthTokenTest.php b/tests/FetchAuthTokenTest.php
@@ -59,14 +59,11 @@ class_implements($fetcherClass)
             $mockFetcher->updateMetadata(Argument::cetera())
                 ->shouldBeCalledTimes(1)
                 ->willReturn(['authorization' => 'Bearer xyz']);
-            $mockFetcher->getLastReceivedToken()
-                ->shouldBeCalledTimes(1)
-                ->will($httpHandler);
-        } else {
-            $mockFetcher->fetchAuthToken(Argument::any())
+        }
+
+        $mockFetcher->fetchAuthToken(Argument::any())
             ->shouldBeCalledTimes(1)
             ->will($httpHandler);
-        }
         $mockFetcher->getCacheKey()->willReturn('');
 
         $tokenCallbackCalled = false;

diff --git a/tests/Middleware/AuthTokenMiddlewareTest.php b/tests/Middleware/AuthTokenMiddlewareTest.php
@@ -22,6 +22,7 @@
 use Google\Auth\Tests\BaseTest;
 use Google\Auth\UpdateMetadataInterface;
 use GuzzleHttp\Handler\MockHandler;
+use GuzzleHttp\Psr7\Request;
 use GuzzleHttp\Psr7\Response;
 use Prophecy\Argument;
 use Prophecy\PhpUnit\ProphecyTrait;
@@ -295,11 +296,13 @@ public function testAddAuthHeadersFromUpdateMetadata()
             ->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);
+        $request = new Request('GET', 'http://foo.com');
+
+        $middleware = new AuthTokenMiddleware($mockFetcher->reveal());
+        $mock = new MockHandler([new Response(200)]);
+        $callable = $middleware($mock);
+        $callable($request, ['auth' => 'google_auth']);
     }
 
     public function testOverlappingAddAuthHeadersFromUpdateMetadata()
@@ -309,21 +312,25 @@ public function testOverlappingAddAuthHeadersFromUpdateMetadata()
             'x-goog-api-client' => 'extra-value'
         ];
 
+        $request = new Request('GET', 'http://foo.com');
+        $request->withHeader('x-goog-api-client', 'default-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);
+        $middleware = new AuthTokenMiddleware($mockFetcher->reveal());
+        $mock = new MockHandler([function ($request, $options) use ($authHeaders) {
+            $this->assertEquals($authHeaders['authorization'], $request->getHeaderLine('authorization'));
+            $this->assertArrayHasKey('x-goog-api-client', $request->getHeaders());
+            return new Response(200);
+        }]);
+        $callable = $middleware($mock);
+        $callable($request, ['auth' => 'google_auth']);
     }
 
     private function runTestCase($fetcher, bool $isGoogleAuth)