Use the graceful deny mechanism for device state PATCH v3

Change-type: major

src/features/device-state/index.ts

@@ -8,7 +8,11 @@ import { resolveOrDenyDevicesWithStatus } from './middleware';
 import { stateV2 } from './routes/state-get-v2';
 import { stateV3 } from './routes/state-get-v3';
 import { statePatchV2 } from './routes/state-patch-v2';
-import { statePatchV3 } from './routes/state-patch-v3';
+import {
+	StatePatchV3Body,
+	resolveDeviceUuids,
+	statePatchV3,
+} from './routes/state-patch-v3';
 import { fleetStateV3 } from './routes/fleet-state-get-v3';
 import { Device } from '../../balena-model';
 
@@ -49,7 +53,14 @@ export const setup = (app: Application) => {
 		middleware.authenticatedApiKey,
 		statePatchV2,
 	);
-	app.patch('/device/v3/state', middleware.authenticatedApiKey, statePatchV3);
+	app.patch(
+		'/device/v3/state',
+		resolveOrDenyDevicesWithStatus(401, (req) =>
+			resolveDeviceUuids(req.body as StatePatchV3Body),
+		),
+		middleware.authenticatedApiKey,
+		statePatchV3,
+	);
 	app.get(
 		'/device/v3/fleet/:fleetUuid/state',
 		middleware.authenticated,

src/features/device-state/middleware.ts

@@ -5,6 +5,7 @@ import type { Device } from '../../balena-model';
 
 import { sbvrUtils, permissions } from '@balena/pinejs';
 import { DEVICE_EXISTS_CACHE_TIMEOUT } from '../../lib/config';
+import type { Request } from 'express-serve-static-core';
 
 const { api } = sbvrUtils;
 
@@ -39,35 +40,50 @@ export const checkDeviceExistsIsFrozen = multiCacheMemoizee(
 );
 
 export interface ResolveDeviceInfoCustomObject {
-	resolvedDevice: Device['id'];
+	resolvedDeviceIds: Array<Device['id']>;
 }
 
+const requestParamsUuidResolver = (req: Request) => [req.params.uuid];
+
 /**
  * This checks if a device is deleted or frozen and responds according to the passed statusCode(s)
  */
 export const resolveOrDenyDevicesWithStatus = (
 	statusCode: number | { deleted: number; frozen: number },
+	uuidResolver: (req: Request) => string[] = requestParamsUuidResolver,
 ): RequestHandler => {
 	const deletedStatusCode =
 		typeof statusCode === 'number' ? statusCode : statusCode.deleted;
 	const frozenStatusCode =
 		typeof statusCode === 'number' ? statusCode : statusCode.frozen;
 
 	return async (req, res, next) => {
-		const device = await checkDeviceExistsIsFrozen(req.params.uuid);
-		if (device == null) {
-			// Gracefully deny deleted devices
+		const uuids = uuidResolver(req);
+		if (!uuids.length) {
 			res.status(deletedStatusCode).end();
 			return;
 		}
-		if (device.is_frozen) {
-			// Gracefully deny frozen devices
-			res.status(frozenStatusCode).end();
-			return;
+		const deviceIds: number[] = [];
+		for (const uuid of uuids) {
+			const device = await checkDeviceExistsIsFrozen(uuid);
+			// Heads-up: if any of the provided devices is deleted/frozen
+			// then the whole request is rejected! We should revisit this
+			// if we later add again support for handling multiple devices
+			// per request.
+			if (device == null) {
+				// Gracefully deny deleted devices
+				res.status(deletedStatusCode).end();
+				return;
+			}
+			if (device.is_frozen) {
+				// Gracefully deny frozen devices
+				res.status(frozenStatusCode).end();
+				return;
+			}
+			deviceIds.push(device.id);
 		}
-
 		req.custom ??= {};
-		(req.custom as ResolveDeviceInfoCustomObject).resolvedDevice = device.id;
+		(req.custom as ResolveDeviceInfoCustomObject).resolvedDeviceIds = deviceIds;
 		next();
 	};
 };

src/features/device-state/routes/state-get-v2.ts

@@ -239,7 +239,8 @@ const stateQuery = _.once(() =>
 );
 
 const getStateV2 = async (req: Request, uuid: string): Promise<StateV2> => {
-	const deviceId = (req.custom as ResolveDeviceInfoCustomObject).resolvedDevice;
+	const [deviceId] = (req.custom as ResolveDeviceInfoCustomObject)
+		.resolvedDeviceIds;
 
 	const device = await getDevice(req, uuid);
 	const config = getConfig(device);

src/features/device-state/routes/state-get-v3.ts

@@ -277,7 +277,8 @@ const stateQuery = _.once(() =>
 );
 
 const getStateV3 = async (req: Request, uuid: string): Promise<StateV3> => {
-	const deviceId = (req.custom as ResolveDeviceInfoCustomObject).resolvedDevice;
+	const [deviceId] = (req.custom as ResolveDeviceInfoCustomObject)
+		.resolvedDeviceIds;
 
 	const device = await getDevice(req, uuid);
 	const config = getConfig(device);

src/features/device-state/routes/state-patch-v2.ts

@@ -116,8 +116,9 @@ export const statePatchV2: RequestHandler = async (req, res) => {
 		if (!uuid) {
 			throw new BadRequestError();
 		}
-		const { resolvedDevice: deviceId } =
-			req.custom as ResolveDeviceInfoCustomObject;
+		const {
+			resolvedDeviceIds: [deviceId],
+		} = req.custom as ResolveDeviceInfoCustomObject;
 		if (deviceId == null) {
 			// We are supposed to have already checked this.
 			throw new UnauthorizedError();

src/features/device-state/routes/state-patch-v3.ts

@@ -23,6 +23,7 @@ import {
 	shouldUpdateMetrics,
 	truncateShortTextFields,
 } from '../state-patch-utils';
+import { ResolveDeviceInfoCustomObject } from '../middleware';
 
 const { BadRequestError, UnauthorizedError, InternalRequestError } = errors;
 const { api } = sbvrUtils;
@@ -80,7 +81,7 @@ export type StatePatchV3Body = {
 const fetchData = async (
 	req: Express.Request,
 	custom: AnyObject,
-	uuids: string[],
+	deviceIds: number[],
 	appReleaseUuids: {
 		[appUuid: string]: {
 			releaseUuids: Set<string>;
@@ -97,7 +98,7 @@ const fetchData = async (
 			options: {
 				$select: ['id', 'uuid'],
 				$filter: {
-					uuid: { $in: uuids },
+					id: { $in: deviceIds },
 				},
 				$expand: {
 					belongs_to__application: {
@@ -110,7 +111,7 @@ const fetchData = async (
 				belongs_to__application: Array<Pick<Application, 'uuid'>>;
 			}
 		>;
-		if (devices.length !== uuids.length) {
+		if (devices.length !== deviceIds.length) {
 			throw new UnauthorizedError();
 		}
 
@@ -205,6 +206,9 @@ const fetchData = async (
 		return { devicesByUuid, images, releasesByAppUuid };
 	});
 
+export const resolveDeviceUuids = (body: StatePatchV3Body) =>
+	Object.keys(body).filter((uuid) => body[uuid] != null);
+
 export const statePatchV3: RequestHandler = async (req, res) => {
 	try {
 		const body = req.body as StatePatchV3Body;
@@ -215,7 +219,7 @@ export const statePatchV3: RequestHandler = async (req, res) => {
 			custom.ipAddress = getIP(req);
 		}
 
-		const uuids = Object.keys(body).filter((uuid) => body[uuid] != null);
+		const uuids = resolveDeviceUuids(body);
 		if (uuids.length === 0) {
 			throw new BadRequestError();
 		}
@@ -298,8 +302,15 @@ export const statePatchV3: RequestHandler = async (req, res) => {
 			}
 
 			if (apps != null || Object.keys(deviceBody).length > 0) {
+				const { resolvedDeviceIds } =
+					req.custom as ResolveDeviceInfoCustomObject;
 				// We lazily fetch the necessary data only if we absolutely must to avoid unnecessary work if it turns out we don't need it
-				data ??= await fetchData(req, custom, uuids, appReleasesCriteria);
+				data ??= await fetchData(
+					req,
+					custom,
+					resolvedDeviceIds,
+					appReleasesCriteria,
+				);
 				const { images, releasesByAppUuid } = data;
 				const device = data.devicesByUuid[uuid];