From d89d6ef718a28e298c5326bb455f102e24019210 Mon Sep 17 00:00:00 2001
From: Pisut Sritrakulchai <pudkrong@hotmail.com>
Date: Tue, 29 Aug 2023 10:34:49 +0200
Subject: [PATCH] feat: ignore health check device

---
 cdk/resources/ConvertDeviceMessages.ts | 20 ++++++++++++++++++++
 lambda/onDeviceMessage.ts              | 26 +++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/cdk/resources/ConvertDeviceMessages.ts b/cdk/resources/ConvertDeviceMessages.ts
index 7b9d436e0..164b7be49 100644
--- a/cdk/resources/ConvertDeviceMessages.ts
+++ b/cdk/resources/ConvertDeviceMessages.ts
@@ -4,6 +4,7 @@ import {
 	aws_iot as IoT,
 	aws_lambda as Lambda,
 	aws_logs as Logs,
+	Stack,
 } from 'aws-cdk-lib'
 import { Construct } from 'constructs'
 import type { PackedLambda } from '../helpers/lambdas/packLambda.js'
@@ -11,6 +12,7 @@ import type { DeviceStorage } from './DeviceStorage.js'
 import { IoTActionRole } from './IoTActionRole.js'
 import { LambdaSource } from './LambdaSource.js'
 import type { WebsocketEventBus } from './WebsocketEventBus.js'
+import { Scope } from '../../util/settings.js'
 
 /**
  * Resources needed to convert messages sent by nRF Cloud to the format that hello.nrfcloud.com expects
@@ -50,9 +52,27 @@ export class ConvertDeviceMessages extends Construct {
 				DEVICES_INDEX_NAME: deviceStorage.devicesTableFingerprintIndexName,
 				NODE_NO_WARNINGS: '1',
 				DISABLE_METRICS: this.node.tryGetContext('isTest') === true ? '1' : '0',
+				STACK_NAME: Stack.of(this).stackName,
 			},
 			layers,
 			logRetention: Logs.RetentionDays.ONE_WEEK,
+			initialPolicy: [
+				new IAM.PolicyStatement({
+					actions: ['ssm:GetParametersByPath', 'ssm:GetParameter'],
+					resources: [
+						`arn:aws:ssm:${Stack.of(this).region}:${
+							Stack.of(this).account
+						}:parameter/${Stack.of(this).stackName}/${
+							Scope.NRFCLOUD_ACCOUNT_PREFIX
+						}`,
+						`arn:aws:ssm:${Stack.of(this).region}:${
+							Stack.of(this).account
+						}:parameter/${Stack.of(this).stackName}/${
+							Scope.NRFCLOUD_ACCOUNT_PREFIX
+						}/*`,
+					],
+				}),
+			],
 		})
 		websocketEventBus.eventBus.grantPutEventsTo(onDeviceMessage)
 		deviceStorage.devicesTable.grantReadData(onDeviceMessage)
diff --git a/lambda/onDeviceMessage.ts b/lambda/onDeviceMessage.ts
index 662500901..abbd4d1a7 100644
--- a/lambda/onDeviceMessage.ts
+++ b/lambda/onDeviceMessage.ts
@@ -8,20 +8,38 @@ import { metricsForComponent } from './metrics/metrics.js'
 import type { WebsocketPayload } from './publishToWebsocketClients.js'
 import { logger } from './util/logger.js'
 import { getDeviceAttributesById } from './getDeviceAttributes.js'
+import { getAllAccountsSettings } from '../nrfcloud/allAccounts.js'
+import { SSMClient } from '@aws-sdk/client-ssm'
+import { once } from 'lodash-es'
 
-const { EventBusName, DevicesTableName } = fromEnv({
+const { EventBusName, DevicesTableName, stackName } = fromEnv({
 	EventBusName: 'EVENTBUS_NAME',
 	DevicesTableName: 'DEVICES_TABLE_NAME',
+	stackName: 'STACK_NAME',
 })(process.env)
 
 const log = logger('deviceMessage')
 const db = new DynamoDBClient({})
+const ssm = new SSMClient({})
 const eventBus = new EventBridge({})
 
 const deviceFetcher = getDeviceAttributesById({ db, DevicesTableName })
 
 const { track, metrics } = metricsForComponent('onDeviceMessage')
 
+const getAllNRFCloudAccountSettings = once(
+	getAllAccountsSettings({
+		ssm,
+		stackName,
+	}),
+)
+const getAllHealthCheckClientIds = once(async () => {
+	const settings = await getAllNRFCloudAccountSettings()
+	return Object.values(settings)
+		.map((settings) => settings?.healthCheckSettings?.healthCheckClientId)
+		.filter((x) => x !== undefined)
+})
+
 const h = async (event: {
 	message: unknown
 	deviceId: string
@@ -31,6 +49,12 @@ const h = async (event: {
 	track('deviceMessage', MetricUnits.Count, 1)
 	const { deviceId, message } = event
 
+	const healthCheckClientIds = await getAllHealthCheckClientIds()
+	if (healthCheckClientIds.includes(deviceId)) {
+		log.debug(`Ignoring health check device`, { deviceId })
+		return
+	}
+
 	// Fetch model for device
 	const { model } = await deviceFetcher(deviceId)