feat(perf): add cache for authtoken lookup

Signed-off-by: Benjamin Gaussorgues <benjamin.gaussorgues@nextcloud.com>

4 files changed, 89 insertions, 77 deletions

diff --git a/lib/private/Authentication/Token/PublicKeyTokenMapper.php b/lib/private/Authentication/Token/PublicKeyTokenMapper.php
index 8feb275b3b7..2621e72e2a1 100644
--- a/lib/private/Authentication/Token/PublicKeyTokenMapper.php
+++ b/lib/private/Authentication/Token/PublicKeyTokenMapper.php
@@ -42,8 +42,6 @@ class PublicKeyTokenMapper extends QBMapper {
 
 	/**
 	 * Invalidate (delete) a given token
-	 *
-	 * @param string $token
 	 */
 	public function invalidate(string $token) {
 		/* @var $qb IQueryBuilder */
@@ -141,14 +139,15 @@ class PublicKeyTokenMapper extends QBMapper {
 		return $entities;
 	}
 
-	public function deleteById(string $uid, int $id) {
+	public function getTokenByUserAndId(string $uid, int $id): ?string {
 		/* @var $qb IQueryBuilder */
 		$qb = $this->db->getQueryBuilder();
-		$qb->delete($this->tableName)
+		$qb->select('token')
+			->from($this->tableName)
 			->where($qb->expr()->eq('id', $qb->createNamedParameter($id)))
 			->andWhere($qb->expr()->eq('uid', $qb->createNamedParameter($uid)))
 			->andWhere($qb->expr()->eq('version', $qb->createNamedParameter(PublicKeyToken::VERSION, IQueryBuilder::PARAM_INT)));
-		$qb->execute();
+		return $qb->executeQuery()->fetchOne() ?: null;
 	}
 
 	/**
diff --git a/lib/private/Authentication/Token/PublicKeyTokenProvider.php b/lib/private/Authentication/Token/PublicKeyTokenProvider.php
index f5fcd4dcef2..7a5e7f6fd5d 100644
--- a/lib/private/Authentication/Token/PublicKeyTokenProvider.php
+++ b/lib/private/Authentication/Token/PublicKeyTokenProvider.php
@@ -29,13 +29,14 @@ declare(strict_types=1);
  */
 namespace OC\Authentication\Token;
 
+use OCP\ICache;
+use OCP\ICacheFactory;
 use OC\Authentication\Exceptions\ExpiredTokenException;
 use OC\Authentication\Exceptions\InvalidTokenException;
 use OC\Authentication\Exceptions\TokenPasswordExpiredException;
 use OC\Authentication\Exceptions\PasswordlessTokenException;
 use OC\Authentication\Exceptions\WipeTokenException;
 use OCP\AppFramework\Db\TTransactional;
-use OCP\Cache\CappedMemoryCache;
 use OCP\AppFramework\Db\DoesNotExistException;
 use OCP\AppFramework\Utility\ITimeFactory;
 use OCP\IConfig;
@@ -47,6 +48,8 @@ use Psr\Log\LoggerInterface;
 
 class PublicKeyTokenProvider implements IProvider {
 	public const TOKEN_MIN_LENGTH = 22;
+	/** Token cache TTL in seconds */
+	private const TOKEN_CACHE_TTL = 10;
 
 	use TTransactional;
 
@@ -67,10 +70,11 @@ class PublicKeyTokenProvider implements IProvider {
 	/** @var ITimeFactory */
 	private $time;
 
-	/** @var CappedMemoryCache */
+	/** @var ICache */
 	private $cache;
 
-	private IHasher $hasher;
+	/** @var IHasher */
+	private $hasher;
 
 	public function __construct(PublicKeyTokenMapper $mapper,
 								ICrypto $crypto,
@@ -78,7 +82,8 @@ class PublicKeyTokenProvider implements IProvider {
 								IDBConnection $db,
 								LoggerInterface $logger,
 								ITimeFactory $time,
-								IHasher $hasher) {
+								IHasher $hasher,
+								ICacheFactory $cacheFactory) {
 		$this->mapper = $mapper;
 		$this->crypto = $crypto;
 		$this->config = $config;
@@ -86,7 +91,7 @@ class PublicKeyTokenProvider implements IProvider {
 		$this->logger = $logger;
 		$this->time = $time;
 
-		$this->cache = new CappedMemoryCache();
+		$this->cache = $cacheFactory->createLocal('authtoken_');
 		$this->hasher = $hasher;
 	}
 
@@ -128,7 +133,7 @@ class PublicKeyTokenProvider implements IProvider {
 		}
 
 		// Add the token to the cache
-		$this->cache[$dbToken->getToken()] = $dbToken;
+		$this->cacheToken($dbToken);
 
 		return $dbToken;
 	}
@@ -156,43 +161,56 @@ class PublicKeyTokenProvider implements IProvider {
 		}
 
 		$tokenHash = $this->hashToken($tokenId);
+		if ($token = $this->getTokenFromCache($tokenHash)) {
+			$this->checkToken($token);
+			return $token;
+		}
 
-		if (isset($this->cache[$tokenHash])) {
-			if ($this->cache[$tokenHash] instanceof DoesNotExistException) {
-				$ex = $this->cache[$tokenHash];
-				throw new InvalidTokenException("Token does not exist: " . $ex->getMessage(), 0, $ex);
-			}
-			$token = $this->cache[$tokenHash];
-		} else {
+		try {
+			$token = $this->mapper->getToken($tokenHash);
+			$this->cacheToken($token);
+		} catch (DoesNotExistException $ex) {
 			try {
-				$token = $this->mapper->getToken($tokenHash);
-				$this->cache[$token->getToken()] = $token;
-			} catch (DoesNotExistException $ex) {
-				try {
-					$token = $this->mapper->getToken($this->hashTokenWithEmptySecret($tokenId));
-					$this->cache[$token->getToken()] = $token;
-					$this->rotate($token, $tokenId, $tokenId);
-				} catch (DoesNotExistException $ex2) {
-					$this->cache[$tokenHash] = $ex2;
-					throw new InvalidTokenException("Token does not exist: " . $ex->getMessage(), 0, $ex);
-				}
+				$token = $this->mapper->getToken($this->hashTokenWithEmptySecret($tokenId));
+				$this->rotate($token, $tokenId, $tokenId);
+			} catch (DoesNotExistException) {
+				$this->cacheInvalidHash($tokenHash);
+				throw new InvalidTokenException("Token does not exist: " . $ex->getMessage(), 0, $ex);
 			}
 		}
 
-		if ((int)$token->getExpires() !== 0 && $token->getExpires() < $this->time->getTime()) {
-			throw new ExpiredTokenException($token);
-		}
+		$this->checkToken($token);
 
-		if ($token->getType() === IToken::WIPE_TOKEN) {
-			throw new WipeTokenException($token);
-		}
+		return $token;
+	}
 
-		if ($token->getPasswordInvalid() === true) {
-			//The password is invalid we should throw an TokenPasswordExpiredException
-			throw new TokenPasswordExpiredException($token);
+	/**
+	 * @throws InvalidTokenException when token doesn't exist
+	 */
+	private function getTokenFromCache(string $tokenHash): ?PublicKeyToken {
+		$serializedToken = $this->cache->get($tokenHash);
+		if (null === $serializedToken) {
+			if ($this->cache->hasKey($tokenHash)) {
+				throw new InvalidTokenException('Token does not exist: ' . $tokenHash);
+			}
+
+			return null;
 		}
 
-		return $token;
+		$token = unserialize($serializedToken, [
+			'allowed_classes' => [PublicKeyToken::class],
+		]);
+
+		return $token instanceof PublicKeyToken ? $token : null;
+	}
+
+	private function cacheToken(PublicKeyToken $token): void {
+		$this->cache->set($token->getToken(), serialize($token), self::TOKEN_CACHE_TTL);
+	}
+
+	private function cacheInvalidHash(string $tokenHash) {
+		// Invalid entries can be kept longer in cache since it’s unlikely to reuse them
+		$this->cache->set($tokenHash, null, self::TOKEN_CACHE_TTL * 2);
 	}
 
 	public function getTokenById(int $tokenId): IToken {
@@ -202,6 +220,12 @@ class PublicKeyTokenProvider implements IProvider {
 			throw new InvalidTokenException("Token with ID $tokenId does not exist: " . $ex->getMessage(), 0, $ex);
 		}
 
+		$this->checkToken($token);
+
+		return $token;
+	}
+
+	private function checkToken($token): void {
 		if ((int)$token->getExpires() !== 0 && $token->getExpires() < $this->time->getTime()) {
 			throw new ExpiredTokenException($token);
 		}
@@ -214,13 +238,9 @@ class PublicKeyTokenProvider implements IProvider {
 			//The password is invalid we should throw an TokenPasswordExpiredException
 			throw new TokenPasswordExpiredException($token);
 		}
-
-		return $token;
 	}
 
 	public function renewSessionToken(string $oldSessionId, string $sessionId): IToken {
-		$this->cache->clear();
-
 		return $this->atomic(function () use ($oldSessionId, $sessionId) {
 			$token = $this->getToken($oldSessionId);
 
@@ -242,7 +262,9 @@ class PublicKeyTokenProvider implements IProvider {
 				IToken::TEMPORARY_TOKEN,
 				$token->getRemember()
 			);
+			$this->cacheToken($newToken);
 
+			$this->cacheInvalidHash($token->getToken());
 			$this->mapper->delete($token);
 
 			return $newToken;
@@ -250,21 +272,22 @@ class PublicKeyTokenProvider implements IProvider {
 	}
 
 	public function invalidateToken(string $token) {
-		$this->cache->clear();
-
+		$tokenHash = $this->hashToken($token);
 		$this->mapper->invalidate($this->hashToken($token));
 		$this->mapper->invalidate($this->hashTokenWithEmptySecret($token));
+		$this->cacheInvalidHash($tokenHash);
 	}
 
 	public function invalidateTokenById(string $uid, int $id) {
-		$t