RSDK-12351 - remove location secret connection code (#659)

Author: stuqdog

src/app/viam-client.spec.ts

@@ -1,7 +1,10 @@
 // @vitest-environment happy-dom
 
 import { beforeEach, describe, expect, it, vi } from 'vitest';
-import { Location, RobotPart, SharedSecret_State } from '../gen/app/v1/app_pb';
+import {
+  GetRobotPartByNameAndLocationResponse,
+  RobotPart,
+} from '../gen/app/v1/app_pb';
 import { createRobotClient } from '../robot/dial';
 import { AppClient } from './app-client';
 import { BillingClient } from './billing-client';
@@ -156,98 +159,64 @@ describe('ViamClient', () => {
       ).rejects.toThrowError('not provided and could not be obtained');
     });
 
-    it('gets location secret if credential is access token -- host', async () => {
+    it('gets robot secret if credential is access token -- host', async () => {
       options = { credentials: testAccessToken };
       const client = await subject();
 
-      const location = new Location({
-        auth: {
-          secrets: [
-            {
-              id: '0',
-              state: SharedSecret_State.DISABLED, // eslint-disable-line camelcase
-              secret: 'disabled secret',
-            },
-            {
-              id: '1',
-              state: SharedSecret_State.UNSPECIFIED, // eslint-disable-line camelcase
-              secret: 'unspecified secret',
-            },
-            {
-              id: '2',
-              state: SharedSecret_State.ENABLED, // eslint-disable-line camelcase
-              secret: 'enabled secret',
-            },
-          ],
-          locationId: 'location',
-          secret: 'secret',
-        },
+      const MAIN_PART = new RobotPart({
+        mainPart: true,
+        name: 'main-part',
+        secret: 'fake-robot-secret',
       });
-      const getLocationMock = vi.fn().mockImplementation(() => location);
-      AppClient.prototype.getLocation = getLocationMock;
+      const partByNameAndLocationResponse =
+        new GetRobotPartByNameAndLocationResponse({
+          part: MAIN_PART,
+        });
+      const getRobotPartByNameAndLocationMock = vi
+        .fn()
+        .mockImplementation(() => partByNameAndLocationResponse);
+      AppClient.prototype.getRobotPartByNameAndLocation =
+        getRobotPartByNameAndLocationMock;
 
       await client.connectToMachine({
         host: 'main-part.location.viam.cloud',
       });
-      expect(getLocationMock).toHaveBeenCalledWith('location');
+      expect(getRobotPartByNameAndLocationMock).toHaveBeenCalledWith(
+        'main-part',
+        'location'
+      );
       expect(createRobotClient).toHaveBeenCalledWith(
         expect.objectContaining({
           credentials: expect.objectContaining({
-            type: 'robot-location-secret',
-            payload: 'enabled secret',
+            type: 'robot-secret',
+            payload: 'fake-robot-secret',
           }),
         })
       );
     });
 
-    it('gets location secret if credential is access token -- id', async () => {
+    it('gets robot secret if credential is access token -- id', async () => {
       options = { credentials: testAccessToken };
       const client = await subject();
 
       const MAIN_PART = new RobotPart({
         mainPart: true,
-        locationId: 'location-id',
         fqdn: 'main-part.fqdn',
+        secret: 'fake-robot-secret',
       });
       const robotParts = [MAIN_PART];
       const getRobotPartsMock = vi.fn().mockImplementation(() => robotParts);
       AppClient.prototype.getRobotParts = getRobotPartsMock;
 
-      const location = new Location({
-        auth: {
-          secrets: [
-            {
-              id: '0',
-              state: SharedSecret_State.DISABLED, // eslint-disable-line camelcase
-              secret: 'disabled secret',
-            },
-            {
-              id: '1',
-              state: SharedSecret_State.UNSPECIFIED, // eslint-disable-line camelcase
-              secret: 'unspecified secret',
-            },
-            {
-              id: '2',
-              state: SharedSecret_State.ENABLED, // eslint-disable-line camelcase
-              secret: 'enabled secret',
-            },
-          ],
-          locationId: 'location',
-          secret: 'secret',
-        },
-      });
-      const getLocationMock = vi.fn().mockImplementation(() => location);
-      AppClient.prototype.getLocation = getLocationMock;
-
       await client.connectToMachine({
         id: 'machine-uuid',
       });
-      expect(getLocationMock).toHaveBeenCalledWith('location-id');
+      expect(getRobotPartsMock).toHaveBeenCalledWith('machine-uuid');
       expect(createRobotClient).toHaveBeenCalledWith(
         expect.objectContaining({
           credentials: expect.objectContaining({
-            type: 'robot-location-secret',
-            payload: 'enabled secret',
+            type: 'robot-secret',
+            payload: 'fake-robot-secret',
           }),
         })
       );

src/app/viam-client.ts

@@ -1,5 +1,4 @@
 import type { Transport } from '@connectrpc/connect';
-import { SharedSecret_State } from '../gen/app/v1/app_pb';
 import { createRobotClient } from '../robot/dial';
 import { AppClient } from './app-client';
 import { BillingClient } from './billing-client';
@@ -54,6 +53,26 @@ export class ViamClient {
     this.billingClient = new BillingClient(this.transport);
   }
 
+  async getRobotSecretFromHost(host: string): Promise<string | undefined> {
+    const firstHalf = host.split('.viam.');
+    const locationSplit = firstHalf[0]?.split('.');
+    if (locationSplit !== undefined) {
+      const locationId = locationSplit.at(-1);
+      if (locationId === undefined) {
+        return undefined;
+      }
+      const name = host.split('.').at(0);
+      if (name !== undefined) {
+        const resp = await this.appClient.getRobotPartByNameAndLocation(
+          name,
+          locationId
+        );
+        return resp.part?.secret;
+      }
+    }
+    return undefined;
+  }
+
   public async connectToMachine({
     host = undefined,
     id = undefined,
@@ -62,7 +81,7 @@ export class ViamClient {
       throw new Error('Either a machine address or ID must be provided');
     }
     let address = host;
-    let locationId: string | undefined = undefined;
+    let robotSecret: string | undefined = undefined;
 
     // Get address if only ID was provided
     if (id !== undefined && host === undefined) {
@@ -74,7 +93,7 @@ export class ViamClient {
         );
       }
       address = mainPart.fqdn;
-      locationId = mainPart.locationId;
+      robotSecret = mainPart.secret;
     }
 
     if (address === undefined || address === '') {
@@ -83,31 +102,20 @@ export class ViamClient {
       );
     }
 
-    // If credentials is AccessToken, then attempt to get the robot location secret
+    // If credentials is AccessToken, then attempt to use the robot part secret
     let creds = this.credentials;
     if (!isCredential(creds)) {
-      if (locationId === undefined) {
-        // If we don't have a location, try to get it from the address
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        const firstHalf = address.split('.viam.');
-        const locationSplit = firstHalf[0]?.split('.');
-        if (locationSplit !== undefined) {
-          locationId = locationSplit.at(-1);
-        }
-      }
-      if (locationId !== undefined) {
-        // If we found the location, then attempt to get its secret
-        const location = await this.appClient.getLocation(locationId);
-        const secret = location?.auth?.secrets.find(
-          // eslint-disable-next-line camelcase
-          (sec) => sec.state === SharedSecret_State.ENABLED
-        );
-        creds = {
-          type: 'robot-location-secret',
-          payload: secret?.secret,
-          authEntity: address,
-        } as Credential;
+      if (robotSecret === undefined) {
+        robotSecret = await this.getRobotSecretFromHost(address);
       }
+      creds =
+        robotSecret === undefined
+          ? creds
+          : ({
+              type: 'robot-secret',
+              payload: robotSecret,
+              authEntity: address,
+            } as Credential);
     }
 
     return createRobotClient({

src/app/viam-transport.ts

@@ -16,10 +16,7 @@ export interface Credential {
   payload: string;
 }
 
-export type CredentialType =
-  | 'robot-location-secret'
-  | 'api-key'
-  | 'robot-secret';
+export type CredentialType = 'api-key' | 'robot-secret';
 
 /** An access token used to access protected resources. */
 export interface AccessToken {