Merge pull request #1291 from balena-io/api-key-role-cache

Add a cache for looking up api key roles
balena-io · Apr 6, 2023 · fcbe596 · fcbe596
2 parents 762a4bf + 50a69c9
commit fcbe596
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 28 deletions.
diff --git a/src/features/api-keys/lib.ts b/src/features/api-keys/lib.ts
@@ -4,6 +4,8 @@ import _ from 'lodash';
 
 import { sbvrUtils, permissions, errors } from '@balena/pinejs';
 import { Deferred, Application, Device, Role, User } from '../../balena-model';
+import { multiCacheMemoizee } from '../../infra/cache';
+import { API_KEY_ROLE_CACHE_TIMEOUT } from '../../lib/config';
 
 const { api } = sbvrUtils;
 const { BadRequestError } = errors;
@@ -274,31 +276,29 @@ export const createGenericApiKey = async (
 	);
 };
 
-export const isApiKeyWithRole = async (
-	key: string,
-	roleName: string,
-	tx?: Tx,
-) => {
-	const role = await api.Auth.get({
-		resource: 'role',
-		passthrough: { tx, req: permissions.root },
-		id: {
-			name: roleName,
-		},
-		options: {
-			$select: 'id',
-			$filter: {
-				is_of__api_key: {
-					$any: {
-						$alias: 'khr',
-						$expr: {
-							khr: {
-								api_key: {
-									$any: {
-										$alias: 'k',
-										$expr: {
-											k: {
-												key,
+export const isApiKeyWithRole = (() => {
+	const authQuery = _.once(() =>
+		api.Auth.prepare<{ key: string; roleName: string }>({
+			resource: 'role',
+			passthrough: { req: permissions.root },
+			id: {
+				name: { '@': 'roleName' },
+			},
+			options: {
+				$select: 'id',
+				$filter: {
+					is_of__api_key: {
+						$any: {
+							$alias: 'khr',
+							$expr: {
+								khr: {
+									api_key: {
+										$any: {
+											$alias: 'k',
+											$expr: {
+												k: {
+													key: { '@': 'key' },
+												},
 											},
 										},
 									},
@@ -308,10 +308,24 @@ export const isApiKeyWithRole = async (
 					},
 				},
 			},
+		}),
+	);
+	return multiCacheMemoizee(
+		async (key: string, roleName: string, tx?: Tx): Promise<boolean> => {
+			const role = await authQuery()({ key, roleName }, undefined, {
+				tx,
+			});
+			return role != null;
 		},
-	});
-	return role != null;
-};
+		{
+			cacheKey: 'isApiKeyWithRole',
+			promise: true,
+			primitive: true,
+			maxAge: API_KEY_ROLE_CACHE_TIMEOUT,
+			normalizer: ([key, roleName]) => `${roleName}$${key}`,
+		},
+	);
+})();
 
 /**
  * Temporarily augments the request's api key with the specified permissions.

diff --git a/src/lib/config.ts b/src/lib/config.ts
@@ -333,6 +333,10 @@ export const DEVICE_LOGS_WRITE_AUTH_CACHE_TIMEOUT = intVar(
 	'DEVICE_LOGS_WRITE_AUTH_CACHE_TIMEOUT',
 	5 * MINUTES,
 );
+export const API_KEY_ROLE_CACHE_TIMEOUT = intVar(
+	'API_KEY_ROLE_CACHE_TIMEOUT',
+	5 * MINUTES,
+);
 
 export const GZIP_COMPRESSION_QUALITY = intVar('GZIP_COMPRESSION_QUALITY', -1);
 export const BROTLI_COMPRESSION_QUALITY = intVar(