Skip to content

Commit

Permalink
chore: refactor AuthTokenMiddleware logic (#492)
Browse files Browse the repository at this point in the history
  • Loading branch information
@yash30201
yash30201 authored Nov 15, 2023
1 parent 999e9ce commit 3d68b6d
Show file tree
Hide file tree
Showing 3 changed files with 119 additions and 54 deletions.
47 changes: 27 additions & 20 deletions src/Middleware/AuthTokenMiddleware.php
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,11 @@

namespace Google\Auth\Middleware;

use Google\Auth\FetchAuthTokenCache;
use Google\Auth\FetchAuthTokenInterface;
use Google\Auth\GetQuotaProjectInterface;
use Google\Auth\UpdateMetadataInterface;
use GuzzleHttp\Psr7\Utils;
use Psr\Http\Message\RequestInterface;

/**
Expand All @@ -40,6 +43,9 @@ class AuthTokenMiddleware
private $httpHandler;

/**
* It must be an implementation of FetchAuthTokenInterface.
* It may also implement UpdateMetadataInterface allowing direct
* retrieval of auth related headers
* @var FetchAuthTokenInterface
*/
private $fetcher;
Expand Down Expand Up @@ -99,7 +105,7 @@ public function __invoke(callable $handler)
return $handler($request, $options);
}

$request = $request->withHeader('authorization', 'Bearer ' . $this->fetchToken());
$request = $this->addAuthHeaders($request);

if ($quotaProject = $this->getQuotaProject()) {
$request = $request->withHeader(
Expand All @@ -113,32 +119,33 @@ public function __invoke(callable $handler)
}

/**
* Call fetcher to fetch the token.
* Adds auth related headers to the request.
*
* @return string|null
* @param RequestInterface $request
* @return RequestInterface
*/
private function fetchToken()
private function addAuthHeaders(RequestInterface $request)
{
$auth_tokens = (array) $this->fetcher->fetchAuthToken($this->httpHandler);

if (array_key_exists('access_token', $auth_tokens)) {
// notify the callback if applicable
if ($this->tokenCallback) {
call_user_func(
$this->tokenCallback,
$this->fetcher->getCacheKey(),
$auth_tokens['access_token']
);
}

return $auth_tokens['access_token'];
if (!$this->fetcher instanceof UpdateMetadataInterface ||
($this->fetcher instanceof FetchAuthTokenCache &&
!$this->fetcher->getFetcher() instanceof UpdateMetadataInterface)
) {
$token = $this->fetcher->fetchAuthToken();
$request = $request->withHeader(
'authorization', 'Bearer ' . ($token['access_token'] ?? $token['id_token'])
);
} else {
$headers = $this->fetcher->updateMetadata($request->getHeaders(), null, $this->httpHandler);
$request = Utils::modifyRequest($request, ['set_headers' => $headers]);
}

if (array_key_exists('id_token', $auth_tokens)) {
return $auth_tokens['id_token'];
if ($this->tokenCallback && ($token = $this->fetcher->getLastReceivedToken())) {
if (array_key_exists('access_token', $token)) {
call_user_func($this->tokenCallback, $this->fetcher->getCacheKey(), $token['access_token']);
}
}

return null;
return $request;
}

/**
Expand Down
17 changes: 14 additions & 3 deletions tests/FetchAuthTokenTest.php
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
use Google\Auth\CredentialsLoader;
use Google\Auth\FetchAuthTokenInterface;
use Google\Auth\OAuth2;
use Google\Auth\UpdateMetadataInterface;
use Prophecy\Argument;
use Prophecy\PhpUnit\ProphecyTrait;

Expand Down Expand Up @@ -53,10 +54,20 @@ class_implements($fetcherClass)
)) {
$mockFetcher->getQuotaProject()->shouldBeCalledTimes(1);
}
$mockFetcher->fetchAuthToken(Argument::any())
->shouldBeCalledTimes(1)
->will($httpHandler);

if (is_a($fetcherClass, UpdateMetadataInterface::class, true)) {
$mockFetcher->updateMetadata(Argument::cetera())
->shouldBeCalledTimes(1)->will(function () use (&$httpHandlerCalled) {
$httpHandlerCalled = true;
return ['authorization' => ['Bearer xyz']];
});
} else {
$mockFetcher->fetchAuthToken(Argument::any())
->shouldBeCalledTimes(1)
->will($httpHandler);
}
$mockFetcher->getCacheKey()->willReturn('');
$mockFetcher->getLastReceivedToken()->willReturn(['access_token' => 'xyz']);

$tokenCallbackCalled = false;
$tokenCallback = function ($cacheKey, $accessToken) use (&$tokenCallbackCalled) {
Expand Down
109 changes: 78 additions & 31 deletions tests/Middleware/AuthTokenMiddlewareTest.php
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
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\Request;
use GuzzleHttp\Psr7\Response;
use Prophecy\Argument;
use Prophecy\PhpUnit\ProphecyTrait;
Expand Down Expand Up @@ -64,11 +66,7 @@ public function testAddsTheTokenAsAnAuthorizationHeader()
->shouldBeCalledTimes(1)
->willReturn($this->mockRequest->reveal());

// Run the test.
$middleware = new AuthTokenMiddleware($this->mockFetcher->reveal());
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
$this->runTestCase($this->mockFetcher->reveal());
}

public function testDoesNotAddAnAuthorizationHeaderOnNoAccessToken()
Expand All @@ -80,11 +78,7 @@ public function testDoesNotAddAnAuthorizationHeaderOnNoAccessToken()
$this->mockRequest->withHeader('authorization', 'Bearer ')
->willReturn($this->mockRequest->reveal());

// Run the test.
$middleware = new AuthTokenMiddleware($this->mockFetcher->reveal());
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
$this->runTestCase($this->mockFetcher->reveal());
}

public function testUsesIdTokenWhenAccessTokenDoesNotExist()
Expand All @@ -96,12 +90,10 @@ public function testUsesIdTokenWhenAccessTokenDoesNotExist()
->willReturn($authResult);
$this->mockRequest->withHeader('authorization', 'Bearer ' . $token)
->shouldBeCalledTimes(1)
->willReturn($this->mockRequest);
->willReturn($this->mockRequest->reveal());

$this->runTestCase($this->mockFetcher->reveal());

$middleware = new AuthTokenMiddleware($this->mockFetcher->reveal());
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
}

public function testUsesCachedAccessToken()
Expand Down Expand Up @@ -133,10 +125,7 @@ public function testUsesCachedAccessToken()
null,
$this->mockCache->reveal()
);
$middleware = new AuthTokenMiddleware($cachedFetcher);
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
$this->runTestCase($cachedFetcher);
}

public function testUsesCachedIdToken()
Expand Down Expand Up @@ -168,10 +157,7 @@ public function testUsesCachedIdToken()
null,
$this->mockCache->reveal()
);
$middleware = new AuthTokenMiddleware($cachedFetcher);
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
$this->runTestCase($cachedFetcher);
}

public function testGetsCachedAuthTokenUsingCacheOptions()
Expand Down Expand Up @@ -204,10 +190,7 @@ public function testGetsCachedAuthTokenUsingCacheOptions()
['prefix' => $prefix],
$this->mockCache->reveal()
);
$middleware = new AuthTokenMiddleware($cachedFetcher);
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
$this->runTestCase($cachedFetcher);
}

public function testShouldSaveValueInCacheWithSpecifiedPrefix()
Expand Down Expand Up @@ -248,10 +231,7 @@ public function testShouldSaveValueInCacheWithSpecifiedPrefix()
['prefix' => $prefix, 'lifetime' => $lifetime],
$this->mockCache->reveal()
);
$middleware = new AuthTokenMiddleware($cachedFetcher);
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
$this->runTestCase($cachedFetcher);
}

/**
Expand Down Expand Up @@ -282,6 +262,8 @@ public function testShouldNotifyTokenCallback(callable $tokenCallback)
$this->mockFetcher->fetchAuthToken(Argument::any())
->shouldBeCalledTimes(1)
->willReturn($cachedValue);
$this->mockFetcher->getLastReceivedToken()
->willReturn($cachedValue);
$this->mockRequest->withHeader(Argument::any(), Argument::any())
->willReturn($this->mockRequest->reveal());

Expand All @@ -306,6 +288,71 @@ public function testShouldNotifyTokenCallback(callable $tokenCallback)
$this->assertTrue(MiddlewareCallback::$called);
}

public function testAddAuthHeadersFromUpdateMetadata()
{
$authResult = [
'authorization' => 'Bearer 1/abcdef1234567890',
];

$this->mockFetcher->willImplement(UpdateMetadataInterface::class);
$this->mockFetcher->updateMetadata(Argument::cetera())
->shouldBeCalledTimes(1)
->willReturn($authResult);
$this->mockFetcher->getLastReceivedToken()
->willReturn(['access_token' => '1/abcdef1234567890']);

$request = new Request('GET', 'http://foo.com');

$middleware = new AuthTokenMiddleware($this->mockFetcher->reveal());
$mockHandlerCalled = false;
$mock = new MockHandler([function ($request, $options) use ($authResult, &$mockHandlerCalled) {
$this->assertEquals($authResult['authorization'], $request->getHeaderLine('authorization'));
$mockHandlerCalled = true;
return new Response(200);
}]);
$callable = $middleware($mock);
$callable($request, ['auth' => 'google_auth']);
$this->assertTrue($mockHandlerCalled);
}

public function testOverlappingAddAuthHeadersFromUpdateMetadata()
{
$authHeaders = [
'authorization' => 'Bearer 1/abcdef1234567890',
'x-goog-api-client' => 'extra-value'
];

$request = new Request('GET', 'http://foo.com');

$this->mockFetcher->willImplement(UpdateMetadataInterface::class);
$this->mockFetcher->updateMetadata(Argument::cetera())
->shouldBeCalledTimes(1)
->willReturn($authHeaders);
$this->mockFetcher->getLastReceivedToken()
->willReturn(['access_token' => '1/abcdef1234567890']);

$middleware = new AuthTokenMiddleware($this->mockFetcher->reveal());

$mockHandlerCalled = false;
$mock = new MockHandler([function ($request, $options) use ($authHeaders, &$mockHandlerCalled) {
$this->assertEquals($authHeaders['authorization'], $request->getHeaderLine('authorization'));
$this->assertArrayHasKey('x-goog-api-client', $request->getHeaders());
$mockHandlerCalled = true;
return new Response(200);
}]);
$callable = $middleware($mock);
$callable($request, ['auth' => 'google_auth']);
$this->assertTrue($mockHandlerCalled);
}

private function runTestCase($fetcher)
{
$middleware = new AuthTokenMiddleware($fetcher);
$mock = new MockHandler([new Response(200)]);
$callable = $middleware($mock);
$callable($this->mockRequest->reveal(), ['auth' => 'google_auth']);
}

public function provideShouldNotifyTokenCallback()
{
MiddlewareCallback::$phpunit = $this;
Expand Down

0 comments on commit 3d68b6d

Please sign in to comment.