From fb70f6061d2ce611526a5e21166e3266e9c8f357 Mon Sep 17 00:00:00 2001
From: Brent Shaffer <betterbrent@google.com>
Date: Fri, 8 Sep 2023 13:55:57 -0700
Subject: [PATCH] feat: respect cache control for access token certs

---
 src/AccessToken.php       | 22 ++++++++++----
 tests/AccessTokenTest.php | 63 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 76 insertions(+), 9 deletions(-)

diff --git a/src/AccessToken.php b/src/AccessToken.php
index e1f92ee7e..e8f78c630 100644
--- a/src/AccessToken.php
+++ b/src/AccessToken.php
@@ -33,6 +33,7 @@
 use phpseclib3\Crypt\PublicKeyLoader;
 use phpseclib3\Math\BigInteger as BigInteger3;
 use Psr\Cache\CacheItemPoolInterface;
+use Psr\Http\Message\ResponseInterface;
 use RuntimeException;
 use SimpleJWT\InvalidTokenException;
 use SimpleJWT\JWT as SimpleJWT;
@@ -312,8 +313,19 @@ private function getCerts($location, $cacheKey, array $options = [])
         $certs = $cacheItem ? $cacheItem->get() : null;
 
         $gotNewCerts = false;
+        $expireTime = '+1 hour';
         if (!$certs) {
-            $certs = $this->retrieveCertsFromLocation($location, $options);
+            $response = $this->retrieveCertsFromLocation($location, $options);
+            $certs = json_decode((string) $response->getBody(), true);
+
+            if ($cacheControl = $response->getHeaderLine('Cache-Control')) {
+                array_map(function($value) use (&$expireTime) {
+                    list($key, $value) = explode('=', $value) + [null, null];
+                    if ($key == 'max-age') {
+                        $expireTime = '+' . $value . ' seconds';
+                    }
+                }, explode(', ', $cacheControl));
+            }
 
             $gotNewCerts = true;
         }
@@ -332,7 +344,7 @@ private function getCerts($location, $cacheKey, array $options = [])
         // Push caching off until after verifying certs are in a valid format.
         // Don't want to cache bad data.
         if ($gotNewCerts) {
-            $cacheItem->expiresAt(new DateTime('+1 hour'));
+            $cacheItem->expiresAt(new DateTime($expireTime));
             $cacheItem->set($certs);
             $this->cache->save($cacheItem);
         }
@@ -345,11 +357,11 @@ private function getCerts($location, $cacheKey, array $options = [])
      *
      * @param string $url location
      * @param array<mixed> $options [optional] Configuration options.
-     * @return array<mixed> certificates
+     * @return ResponseInterface
      * @throws InvalidArgumentException If certs could not be retrieved from a local file.
      * @throws RuntimeException If certs could not be retrieved from a remote location.
      */
-    private function retrieveCertsFromLocation($url, array $options = [])
+    private function retrieveCertsFromLocation($url, array $options = []): ResponseInterface
     {
         // If we're retrieving a local file, just grab it.
         if (strpos($url, 'http') !== 0) {
@@ -367,7 +379,7 @@ private function retrieveCertsFromLocation($url, array $options = [])
         $response = $httpHandler(new Request('GET', $url), $options);
 
         if ($response->getStatusCode() == 200) {
-            return json_decode((string) $response->getBody(), true);
+            return $response;
         }
 
         throw new RuntimeException(sprintf(
diff --git a/tests/AccessTokenTest.php b/tests/AccessTokenTest.php
index f2a6f3c81..3b0a283df 100644
--- a/tests/AccessTokenTest.php
+++ b/tests/AccessTokenTest.php
@@ -264,7 +264,10 @@ public function testEsVerifyEndToEnd()
         $this->assertEquals('https://cloud.google.com/iap', $payload['iss']);
     }
 
-    public function testGetCertsForIap()
+    /**
+     * @dataProvider provideCertsFromUrl
+     */
+    public function testGetCertsFromUrl($certUrl)
     {
         $token = new AccessToken();
         $reflector = new \ReflectionObject($token);
@@ -272,14 +275,22 @@ public function testGetCertsForIap()
         $cacheKeyMethod->setAccessible(true);
         $getCertsMethod = $reflector->getMethod('getCerts');
         $getCertsMethod->setAccessible(true);
-        $cacheKey = $cacheKeyMethod->invoke($token, AccessToken::IAP_CERT_URL);
+        $cacheKey = $cacheKeyMethod->invoke($token, $certUrl);
         $certs = $getCertsMethod->invoke(
             $token,
-            AccessToken::IAP_CERT_URL,
+            $certUrl,
             $cacheKey
         );
         $this->assertTrue(is_array($certs));
-        $this->assertEquals(5, count($certs));
+        $this->assertGreaterThanOrEqual(2, count($certs));
+    }
+
+    public function provideCertsFromUrl()
+    {
+        return [
+            [AccessToken::IAP_CERT_URL],
+            [AccessToken::FEDERATED_SIGNON_CERT_URL],
+        ];
     }
 
     public function testRetrieveCertsFromLocationLocalFile()
@@ -398,6 +409,50 @@ public function testRetrieveCertsFromLocationLocalFileInvalidFileData()
         ]);
     }
 
+    public function testRetrieveCertsFromLocationRespectsCacheControl()
+    {
+        $certsLocation = __DIR__ . '/fixtures/federated-certs.json';
+        $certsJson = file_get_contents($certsLocation);
+        $certsData = json_decode($certsJson, true);
+
+        $httpHandler = function (RequestInterface $request) use ($certsJson) {
+            return new Response(200, [
+                'cache-control' => 'public, max-age=1000',
+            ], $certsJson);
+        };
+
+        $phpunit = $this;
+
+        $item = $this->prophesize('Psr\Cache\CacheItemInterface');
+        $item->get()
+            ->shouldBeCalledTimes(1)
+            ->willReturn(null);
+        $item->set($certsData)
+            ->shouldBeCalledTimes(1)
+            ->willReturn($item->reveal());
+
+        // Assert date-time is set with difference of 1000 (the max-age in the Cache-Control header)
+        $item->expiresAt(Argument::type('\DateTime'))
+            ->shouldBeCalledTimes(1)
+            ->will(function ($value) use ($phpunit) {
+                $phpunit->assertEqualsWithDelta(1000, $value[0]->getTimestamp() - time(), 1);
+            });
+
+        $this->cache->getItem('google_auth_certs_cache|federated_signon_certs_v3')
+            ->shouldBeCalledTimes(1)
+            ->willReturn($item->reveal());
+
+        $this->cache->save(Argument::type('Psr\Cache\CacheItemInterface'))
+            ->shouldBeCalledTimes(1);
+
+        $token = new AccessTokenStub(
+            $httpHandler,
+            $this->cache->reveal()
+        );
+
+        $token->verify($this->token);
+    }
+
     public function testRetrieveCertsFromLocationRemote()
     {
         $certsLocation = __DIR__ . '/fixtures/federated-certs.json';