@alwaysai/device-agent 1.5.0 → 2.0.1

package/src/cloud-connection/device-agent-message-handler.ts

@@ -0,0 +1,642 @@
+import {
+  keyMirrors,
+  AppStateControlMessage,
+  buildToClientStatusResponseMessage,
+  AppStateControlPayload,
+  AppVersionControlMessage,
+  AppVersionControlInstallPayload,
+  AppVersionControlUninstallPayload,
+  SignedUrlsRequestPayload,
+  buildSignedUrlsRequestMessage,
+  ToCloudMessage,
+  LiveStateUpdatesToggleMessage,
+  buildDeviceStatsMessage,
+  buildAppStateMessage,
+  buildAppLogsMessage,
+  AppInstallResponsePayload,
+  ModelsInstallResponsePayload,
+  ToDeviceStatusResponseMessage,
+  DeviceActionPayload,
+  validateToDeviceAgentMessage,
+  ToDeviceAgentMessage,
+  EnvVars,
+  validateEnvVarSchemaShadowUpdate
+} from '@alwaysai/device-agent-schemas';
+import {
+  AppInstallResponseMessage,
+  ModelsInstallResponseMessage
+} from '@alwaysai/device-agent-schemas/lib/app-action-schema';
+import { DeviceActionMessage } from '@alwaysai/device-agent-schemas/lib/device-action-schema';
+import { logger, stringifyError } from 'alwaysai/lib/util';
+import {
+  startApp,
+  stopApp,
+  restartApp,
+  getAppLogs,
+  installApp,
+  updateModelsWithPresignedUrls
+} from '../application-control';
+import { reboot } from '../device-control/device-control';
+import { MessageDispatcher, MessageHandler } from './message-dispatcher';
+import { getDeviceStatsPayload, getAppStatePayload } from './messages';
+import { BaseHandler, HandlerContext } from './base-message-handler';
+import { ErrorFunction, SuccessFunction } from './transaction-manager';
+import {
+  AppConfig,
+  validateAppConfig
+} from '@alwaysai/app-configuration-schemas';
+
+export type AppContent = {
+  projectId: string;
+  appCfg?: AppConfig;
+  envVars?: EnvVars;
+};
+
+export type AppContentUpdate = {
+  appCfg?: string;
+  envVars?: string;
+};
+export class DeviceAgentMessageHandler
+  extends MessageDispatcher<ToDeviceAgentMessage>
+  implements MessageHandler<any>
+{
+  private readonly handlerContext: HandlerContext;
+  private readonly errorFn?: ErrorFunction;
+  private readonly successFn?: SuccessFunction;
+
+  constructor(
+    handlerContext: HandlerContext,
+    errorFn?: ErrorFunction,
+    successFn?: SuccessFunction
+  ) {
+    super();
+    this.handlerContext = handlerContext;
+    this.errorFn = errorFn;
+    this.successFn = successFn;
+
+    const handlerMapping: {
+      [key: string]: new (
+        context: HandlerContext,
+        errorFn?: ErrorFunction,
+        successFn?: SuccessFunction
+      ) => MessageHandler<ToDeviceAgentMessage>;
+    } = {
+      [keyMirrors.toDeviceAgentMessageType.app_state_control]:
+        AppStateMessageHandler,
+      [keyMirrors.toDeviceAgentMessageType.live_state_updates]:
+        LiveStateUpdatesMessageHandler,
+      [keyMirrors.toDeviceAgentMessageType.app_version_control]:
+        AppVersionControlMessageHandler,
+      [keyMirrors.toDeviceAgentMessageType.app_install_response]:
+        AppInstallResponseMessageHandler,
+      [keyMirrors.toDeviceAgentMessageType.models_install_response]:
+        ModelsInstallResponseMessageHandler,
+      [keyMirrors.toDeviceAgentMessageType.status_response]:
+        StatusResponseMessageHandler,
+      [keyMirrors.toDeviceAgentMessageType.device_action]:
+        DeviceActionMessageHandler
+    };
+
+    Object.entries(handlerMapping).forEach(([keyMirror, HandlerClass]) => {
+      this.registerHandler(
+        keyMirror,
+        new HandlerClass(this.handlerContext, this.errorFn, this.successFn)
+      );
+    });
+  }
+
+  public handle(message: any, topic?: string): void {
+    const valid = validateToDeviceAgentMessage(message);
+    if (!valid) {
+      logger.error(
+        `Error validating message: ${JSON.stringify(
+          { topic, message, errors: validateToDeviceAgentMessage.errors },
+          null,
+          2
+        )}`
+      );
+      if (this.errorFn) {
+        this.errorFn(
+          '',
+          `Error validating message: ${JSON.stringify(
+            { topic, message, errors: validateToDeviceAgentMessage.errors },
+            null,
+            2
+          )}`
+        );
+      }
+      return;
+    }
+    this.dispatch(message.messageType, message);
+  }
+}
+
+class AppStateMessageHandler
+  extends BaseHandler
+  implements MessageHandler<AppStateControlMessage>
+{
+  public async handle(message: AppStateControlMessage): Promise<void> {
+    const payload = message.payload;
+    const projectId = payload.projectId;
+    const txId = message.txId;
+    try {
+      await this.txnMgr.runTransactionStep({
+        func: () => this.handleAppStateControl(message.payload),
+        projectId,
+        txId,
+        start: true,
+        liveUpdatesPublishFn: async () =>
+          this.publisher.publishToClient(
+            buildToClientStatusResponseMessage(
+              this.clientId,
+              { status: keyMirrors.statusResponse.in_progress },
+              txId
+            ),
+            logger.silly
+          ),
+        stepName: payload.baseCommand,
+        errorFn: this.errorFn,
+        successFn: this.successFn
+      });
+    } catch (e) {
+      logger.error(
+        `Error processing application state control request for ${projectId}!\n${stringifyError(
+          e
+        )}`
+      );
+    }
+  }
+
+  private handleAppStateControl = async (
+    payload: AppStateControlPayload
+  ): Promise<boolean> => {
+    const { baseCommand, projectId } = payload;
+    switch (baseCommand) {
+      case keyMirrors.appStateControl.start:
+        await startApp({ projectId });
+        break;
+      case keyMirrors.appStateControl.stop:
+        await stopApp({ projectId });
+        break;
+      case keyMirrors.appStateControl.restart:
+        await restartApp({ projectId });
+        break;
+    }
+    return true;
+  };
+}
+
+class AppVersionControlMessageHandler
+  extends BaseHandler
+  implements MessageHandler<AppVersionControlMessage>
+{
+  public async handle(message: AppVersionControlMessage): Promise<void> {
+    const payload = message.payload;
+    const projectId = payload.projectId;
+    const txId = message.txId;
+    try {
+      await this.txnMgr.runTransactionStep({
+        func: () => this.handleAppVersionControl(payload, txId),
+        projectId,
+        txId,
+        start: true,
+        liveUpdatesPublishFn: async () =>
+          this.publisher.publishToClient(
+            buildToClientStatusResponseMessage(
+              this.clientId,
+              { status: keyMirrors.statusResponse.in_progress },
+              txId
+            ),
+            logger.silly
+          ),
+        stepName: payload.baseCommand,
+        errorFn: this.errorFn,
+        successFn: this.successFn
+      });
+    } catch (e) {
+      logger.error(
+        `Error processing application install request for ${projectId}!\n${stringifyError(
+          e
+        )}`
+      );
+    }
+  }
+
+  private handleAppVersionControl = async (
+    payload:
+      | AppVersionControlInstallPayload
+      | AppVersionControlUninstallPayload,
+    txId: string
+  ): Promise<boolean> => {
+    switch (payload.baseCommand) {
+      case keyMirrors.appVersionControl.install: {
+        const { projectId, appReleaseHash, appCfg, envVars } = payload;
+        const appContentUpdate: AppContentUpdate = {
+          appCfg,
+          envVars
+        };
+        const appContent = await this.processAppContentUpdate(
+          projectId,
+          appContentUpdate
+        );
+        if (appContent) {
+          this.txnMgr.setAppContentToTx(txId, appContent);
+        }
+
+        const signedUrlsRequestPayload: SignedUrlsRequestPayload = {
+          signedUrlsRequest: {
+            projectId,
+            appReleaseHash
+          }
+        };
+        const message = buildSignedUrlsRequestMessage(
+          this.clientId,
+          signedUrlsRequestPayload,
+          txId
+        );
+        await this.publishCloudRequest(message);
+        return false;
+      }
+      case keyMirrors.appVersionControl.uninstall: {
+        const { projectId } = payload;
+        await this.atomicApplicationUninstall(projectId);
+        return true;
+      }
+      default:
+        logger.warn(
+          `Ignore App Version Control packet: ${JSON.stringify(
+            payload,
+            null,
+            2
+          )}`
+        );
+        return true;
+    }
+  };
+
+  private async publishCloudRequest(message: ToCloudMessage) {
+    this.publisher.publishToCloud(message);
+  }
+
+  private async processAppContentUpdate(
+    projectId: string,
+    appContentUpdate: AppContentUpdate
+  ): Promise<AppContent | null> {
+    const appContent: AppContent = { projectId };
+    // appCfg
+    if (!appContentUpdate.appCfg) {
+      logger.info(`No appCfgUpdate for ${projectId}`);
+    } else {
+      // Handle errors and validation
+      try {
+        const appCfgUpdate = JSON.parse(appContentUpdate.appCfg);
+        if (!validateAppConfig(appCfgUpdate)) {
+          logger.error(
+            `Received invalid app config for ${projectId}!\n${JSON.stringify(
+              validateAppConfig.errors,
+              null,
+              2
+            )}`
+          );
+          throw new Error(
+            `Received invalid app config for ${projectId}!\n${JSON.stringify(
+              validateAppConfig.errors,
+              null,
+              2
+            )}`
+          );
+        } else {
+          appContent.appCfg = appCfgUpdate;
+        }
+      } catch (e) {
+        logger.error(
+          `Could not parse the appConfig for transaction!\n${stringifyError(e)}`
+        );
+      }
+    }
+
+    // envVars
+    if (!appContentUpdate.envVars) {
+      logger.info(`No envVars for ${projectId}`);
+    } else {
+      try {
+        const envvarsUpdate = JSON.parse(appContentUpdate.envVars);
+        if (!validateEnvVarSchemaShadowUpdate(envvarsUpdate)) {
+          logger.error(
+            `Received invalid environment variables update for ${projectId}!\n${JSON.stringify(
+              validateEnvVarSchemaShadowUpdate.errors,
+              null,
+              2
+            )}`
+          );
+          throw new Error(
+            `Received invalid environment variables update for${projectId}!\n${JSON.stringify(
+              validateEnvVarSchemaShadowUpdate.errors,
+              null,
+              2
+            )}`
+          );
+        } else {
+          appContent.envVars = envvarsUpdate;
+        }
+      } catch (e) {
+        // throw here
+        logger.error(
+          `Could not parse the environment variables for transaction!\n${stringifyError(
+            e
+          )}`
+        );
+      }
+    }
+
+    return appContent.appCfg || appContent.envVars ? appContent : null;
+  }
+}
+
+class LiveStateUpdatesMessageHandler
+  extends BaseHandler
+  implements MessageHandler<LiveStateUpdatesToggleMessage>
+{
+  public async handle(message: LiveStateUpdatesToggleMessage): Promise<void> {
+    const { deviceStats, appState, appLogs } = message.payload;
+    const txId = message.txId;
+
+    if (deviceStats !== undefined) {
+      if (deviceStats) {
+        await this.liveUpdatesHandler.enable(
+          keyMirrors.toClientMessageType.device_stats,
+          async () =>
+            this.publisher.publishToClient(
+              buildDeviceStatsMessage(
+                this.clientId,
+                await getDeviceStatsPayload(),
+                txId
+              ),
+              logger.silly
+            )
+        );
+      } else {
+        this.liveUpdatesHandler.disable(
+          keyMirrors.toClientMessageType.device_stats
+        );
+      }
+    }
+
+    if (appState !== undefined) {
+      if (appState) {
+        await this.liveUpdatesHandler.enable(
+          keyMirrors.toClientMessageType.app_state,
+          async () =>
+            this.publisher.publishToClient(
+              buildAppStateMessage(
+                this.clientId,
+                await getAppStatePayload(),
+                txId
+              ),
+              logger.silly
+            )
+        );
+      } else {
+        this.liveUpdatesHandler.disable(
+          keyMirrors.toClientMessageType.app_state
+        );
+      }
+    }
+
+    if (appLogs !== undefined) {
+      if (appLogs.toggle) {
+        await this.liveUpdatesHandler.startStream(
+          appLogs.projectId,
+          async () =>
+            await getAppLogs({
+              projectId: appLogs.projectId,
+              args: ['--tail', '100', '--no-log-prefix']
+            }),
+          async (logChunk: string) =>
+            this.publisher.publishToClient(
+              buildAppLogsMessage(
+                this.clientId,
+                {
+                  projectId: appLogs.projectId,
+                  logChunk
+                },
+                txId
+              ),
+              logger.silly
+            )
+        );
+      } else {
+        this.liveUpdatesHandler.stopStream(appLogs.projectId);
+      }
+    }
+  }
+}
+
+class AppInstallResponseMessageHandler
+  extends BaseHandler
+  implements MessageHandler<AppInstallResponseMessage>
+{
+  public async handle(message: AppInstallResponseMessage): Promise<void> {
+    const payload = message.payload;
+    const txId = message.txId;
+    const { projectId } = payload.appInstallResponse;
+    const start = !this.txnMgr.isOngoingTransactionForProjectID(projectId);
+
+    if (!start && txId !== this.txnMgr.getTransactionFromProject(projectId)) {
+      throw new Error(
+        `App install response received a message for a transaction ID ${txId} that is not currently underway (${this.txnMgr.getTransactionFromProject(
+          projectId
+        )})!`
+      );
+    }
+    const appContent = this.txnMgr.getAppContentFromTxId(txId);
+    await this.txnMgr.runTransactionStep({
+      func: () =>
+        this.handleAppInstallCloudResponsePayload(payload, appContent),
+      projectId,
+      txId,
+      start,
+      stepName: message.messageType,
+      errorFn: this.errorFn,
+      successFn: this.successFn
+    });
+  }
+
+  private handleAppInstallCloudResponsePayload = async (
+    payload: AppInstallResponsePayload,
+    appContent?: AppContent
+  ): Promise<boolean> => {
+    const {
+      projectId,
+      appReleaseHash,
+      appInstallPayload,
+      modelsInstallPayload
+    } = payload.appInstallResponse;
+    const signedUrlsPayload = {
+      appInstallPayload,
+      modelsInstallPayload
+    };
+    await this.atomicApplicationUpdate(async () => {
+      this.shadowHandler.clearProjectShadow(projectId);
+      await installApp({
+        projectId,
+        appReleaseHash,
+        signedUrlsPayload,
+        appCfg: appContent?.appCfg,
+        envVars: appContent?.envVars
+      });
+    }, projectId);
+    return true;
+  };
+}
+
+class ModelsInstallResponseMessageHandler
+  extends BaseHandler
+  implements MessageHandler<ModelsInstallResponseMessage>
+{
+  public async handle(message: ModelsInstallResponseMessage): Promise<void> {
+    const payload = message.payload;
+    const txId = message.txId;
+    const { projectId } = payload.modelsInstallResponse;
+    if (txId !== this.txnMgr.getTransactionFromProject(projectId)) {
+      throw new Error(
+        `Model install response received a message for a transaction ID ${txId} that is not currently underway (${this.txnMgr.getTransactionFromProject(
+          projectId
+        )})!`
+      );
+    }
+    await this.txnMgr.runTransactionStep({
+      func: () => this.handleModelsInstallCloudResponsePayload(payload, txId),
+      projectId,
+      txId,
+      start: false,
+      stepName: message.messageType,
+      errorFn: this.errorFn,
+      successFn: this.successFn
+    });
+  }
+
+  private handleModelsInstallCloudResponsePayload = async (
+    payload: ModelsInstallResponsePayload,
+    txId: string
+  ): Promise<boolean> => {
+    const projectId = payload.modelsInstallResponse.projectId;
+
+    const update = this.txnMgr.getAppCfgUpdateFromTxID(txId);
+    if (update === undefined) {
+      throw new Error(
+        'Unknown error while updating models via application config! No config present for model update.'
+      );
+    }
+    const { appCfgUpdate, envVarUpdate } = update;
+    if (appCfgUpdate) {
+      await this.atomicApplicationUpdate(
+        async () =>
+          await updateModelsWithPresignedUrls({
+            projectId,
+            modelInstallPayloads: payload.modelsInstallResponse.newModels,
+            newAppCfg: appCfgUpdate.newAppCfg
+          }),
+        projectId
+      );
+    }
+
+    if (envVarUpdate) {
+      await this.atomicApplicationUpdate(
+        async () =>
+          await this.shadowHandler.updateProjectEnvVars({
+            projectId,
+            envVars: envVarUpdate.envVars
+          }),
+        projectId,
+        true
+      );
+    }
+    return true;
+  };
+}
+
+class StatusResponseMessageHandler
+  extends BaseHandler
+  implements MessageHandler<ToDeviceStatusResponseMessage>
+{
+  public async handle(message: ToDeviceStatusResponseMessage): Promise<void> {
+    const txId = message.txId;
+    const { failure } = keyMirrors.statusResponse;
+    if (message.payload.status === failure) {
+      this.txnMgr.completeTransaction(txId);
+      const msg = buildToClientStatusResponseMessage(
+        this.clientId,
+        {
+          status: keyMirrors.statusResponse.failure,
+          message: message.payload.message
+        },
+        txId
+      );
+      this.publisher.publishToClient(msg);
+    }
+  }
+}
+
+class DeviceActionMessageHandler
+  extends BaseHandler
+  implements MessageHandler<DeviceActionMessage>
+{
+  public async handle(message: DeviceActionMessage): Promise<void> {
+    const txId = message.txId;
+    try {
+      this.publisher.publishToClient(
+        buildToClientStatusResponseMessage(
+          this.clientId,
+          {
+            status: keyMirrors.statusResponse.in_progress
+          },
+          txId
+        )
+      );
+
+      await this.handleDeviceAction(message.payload);
+
+      this.publisher.publishToClient(
+        buildToClientStatusResponseMessage(
+          this.clientId,
+          {
+            status: keyMirrors.statusResponse.success
+          },
+          txId
+        )
+      );
+    } catch (e) {
+      logger.error(
+        `There was a problem performing device action '${
+          message.payload.action
+        }'!\n${stringifyError(e)}`
+      );
+      this.publisher.publishToClient(
+        buildToClientStatusResponseMessage(
+          this.clientId,
+          {
+            status: keyMirrors.statusResponse.failure,
+            message: e.message
+          },
+          txId
+        )
+      );
+    }
+  }
+
+  private async handleDeviceAction(payload: DeviceActionPayload) {
+    const { system_restart } = keyMirrors.deviceAction;
+    switch (payload.action) {
+      case system_restart: {
+        await reboot();
+        break;
+      }
+      default: {
+        logger.info(
+          `Unrecognized device action requested: '${payload.action}'.`
+        );
+      }
+    }
+  }
+}

package/src/cloud-connection/device-agent.ts

@@ -1,16 +1,20 @@
 // eslint-disable-next-line
 const awsIot = require('aws-iot-device-sdk');
+import { getTargetHardwareUuid } from 'alwaysai/lib/core/app';
 import {
+  DEVICE_CERTIFICATE_FILE_NAME,
+  DEVICE_CERTIFICATE_ID_FILE_NAME,
   DEVICE_PRIVATE_KEY_FILE_NAME,
+  LOCAL_CERT_AND_KEY_DIR
+} from 'alwaysai/lib/infrastructure';
+import { JsSpawner } from 'alwaysai/lib/util';
+import {
   CERTIFICATE_OWNERSHIP_TOKEN_FILE_NAME,
-  DEVICE_ID_FILE_NAME,
-  DEVICE_CERTIFICATE_FILE_NAME
-} from '../util/directories';
+  getBootstrapCertificateDirectoryPath
+} from '../infrastructure/device-certificate';
 import { getDeviceUuid } from '../util/get-device-id';
-import { JsSpawner } from 'alwaysai/lib/util';
 import { logger } from '../util/logger';
-import { getTargetHardwareUuid } from 'alwaysai/lib/core/app';
-import { LOCAL_CERT_AND_KEY_DIR } from 'alwaysai/lib/paths';
+
 // eslint-disable-next-line
 const process = require('process');
 
@@ -85,7 +89,10 @@ export class BootstrapAgent extends DeviceAgent {
 
         await certSpawner.writeFile(DEVICE_PRIVATE_KEY_FILE_NAME, privateKey);
 
-        await certSpawner.writeFile(DEVICE_ID_FILE_NAME, certificateId);
+        await certSpawner.writeFile(
+          DEVICE_CERTIFICATE_ID_FILE_NAME,
+          certificateId
+        );
 
         await certSpawner.writeFile(
           CERTIFICATE_OWNERSHIP_TOKEN_FILE_NAME,
@@ -111,6 +118,8 @@ export class BootstrapAgent extends DeviceAgent {
       }
       case '$aws/provisioning-templates/FleetProvisionTemplate/provision/json/accepted': {
         logger.info('Device agent provisioning: success');
+        const spawner = JsSpawner();
+        await spawner.rimraf(getBootstrapCertificateDirectoryPath());
         process.exit();
       }
     }