Implement temporary last step of Connect My Computer flow in Discover

web/packages/teleport/src/Discover/ConnectMyComputer/SetupConnect/SetupConnect.story.tsx

@@ -37,6 +37,7 @@ const oneDay = 1000 * 60 * 60 * 24;
 const setupConnectProps = {
   prevStep: () => {},
   nextStep: () => {},
+  updateAgentMeta: () => {},
   // Set high default intervals and timeouts so that stories don't poll for no reason.
   pingInterval: oneDay,
   showHintTimeout: oneDay,
@@ -120,7 +121,12 @@ export const HintTimeout = () => (
 
 HintTimeout.parameters = {
   msw: {
-    handlers: [noNodesHandler],
+    handlers: [
+      noNodesHandler,
+      rest.post(cfg.api.webRenewTokenPath, (req, res, ctx) =>
+        res(ctx.json({}))
+      ),
+    ],
   },
 };
 

web/packages/teleport/src/Discover/ConnectMyComputer/SetupConnect/SetupConnect.test.ts

@@ -17,13 +17,12 @@
 import { renderHook } from '@testing-library/react-hooks';
 
 import * as useTeleport from 'teleport/useTeleport';
-import NodeService from 'teleport/services/nodes/nodes';
+import NodeService, { Node } from 'teleport/services/nodes';
+import UserService from 'teleport/services/user';
 import TeleportContext from 'teleport/teleportContext';
 
 import { nodes } from 'teleport/Nodes/fixtures';
 
-import { Node } from 'teleport/services/nodes';
-
 import { usePollForConnectMyComputerNode } from './SetupConnect';
 
 beforeEach(() => {
@@ -66,6 +65,7 @@ describe('usePollForConnectMyComputerNode', () => {
         username: 'alice',
         clusterId: 'foo',
         pingInterval: 1,
+        reloadUser: false,
       })
     );
 
@@ -78,4 +78,64 @@ describe('usePollForConnectMyComputerNode', () => {
     expect(result.current.node).toEqual(expectedNode);
     expect(result.current.isPolling).toBe(false);
   });
+
+  it('reloads user before each poll if reloadUser is true', async () => {
+    const expectedNode = nodes[0];
+    let hasReloadedUser = false;
+
+    const nodeService = {
+      fetchNodes: jest.fn(),
+    } as Partial<NodeService> as NodeService;
+
+    jest.mocked(nodeService).fetchNodes.mockImplementation(async () => {
+      if (hasReloadedUser) {
+        return { agents: [expectedNode] };
+      } else {
+        return { agents: [] };
+      }
+    });
+
+    const userService = {
+      reloadUser: jest.fn(),
+    } as Partial<typeof UserService> as typeof UserService;
+
+    jest.mocked(userService).reloadUser.mockImplementation(async () => {
+      hasReloadedUser = true;
+    });
+
+    jest
+      .spyOn(useTeleport, 'default')
+      .mockReturnValue({ nodeService, userService } as TeleportContext);
+
+    const { result, rerender, waitFor, waitForValueToChange } = renderHook(
+      usePollForConnectMyComputerNode,
+      {
+        initialProps: {
+          reloadUser: false,
+          username: 'alice',
+          clusterId: 'foo',
+          pingInterval: 1,
+        },
+      }
+    );
+    expect(result.error).toBeUndefined();
+    await waitFor(() => {
+      expect(nodeService.fetchNodes).toHaveBeenCalled();
+    });
+    expect(userService.reloadUser).not.toHaveBeenCalled();
+
+    rerender({
+      reloadUser: true,
+      username: 'alice',
+      clusterId: 'foo',
+      pingInterval: 1,
+    });
+    expect(result.error).toBeUndefined();
+
+    await waitForValueToChange(() => result.current.node, { interval: 3 });
+    expect(userService.reloadUser).toHaveBeenCalled();
+
+    expect(result.current.node).toEqual(expectedNode);
+    expect(result.current.isPolling).toBe(false);
+  });
 });

web/packages/teleport/src/Discover/ConnectMyComputer/SetupConnect/SetupConnect.tsx

@@ -53,6 +53,7 @@ export function SetupConnect(
 ) {
   const pingInterval = props.pingInterval || 1000 * 3; // 3 seconds
   const showHintTimeout = props.showHintTimeout || 1000 * 60 * 5; // 5 minutes
+
   const ctx = useTeleport();
   const clusterId = ctx.storeUser.getClusterId();
   const { cluster, username } = ctx.storeUser.state;
@@ -63,14 +64,46 @@ export function SetupConnect(
     username,
     path: Path.ConnectMyComputer,
   });
+  const [showHint, setShowHint] = useState(false);
 
   const { node, isPolling } = usePollForConnectMyComputerNode({
     username,
     clusterId,
-    pingInterval,
+    // If reloadUser is set to true, the polling callback takes longer to finish so let's increase
+    // the polling interval as well.
+    pingInterval: showHint ? pingInterval * 2 : pingInterval,
+    // Completing the Connect My Computer setup in Connect causes the user to gain a new role. That
+    // role grants access to nodes labeled with `teleport.dev/connect-my-computer/owner:
+    // <current-username>`.
+    //
+    // In certain cases, that role might be the only role which grants the user the visibility of
+    // the Connect My Computer node. For example, if the user doesn't have a role like the built-in
+    // access which gives blanket access to all nodes, the user won't be able to see the node until
+    // they have the Connect My Computer role in their cert.
+    //
+    // As such, if we don't reload the cert during polling, it might never see the node. So let's
+    // flip it to true after a timeout.
+    reloadUser: showHint,
   });
 
-  const [showHint, setShowHint] = useState(false);
+  // TODO(ravicious): Take these from the context rather than from props.
+  const { agentMeta, updateAgentMeta, nextStep } = props;
+  const handleNextStep = () => {
+    if (!node) {
+      return;
+    }
+
+    updateAgentMeta({
+      ...agentMeta,
+      // Node is an oddity in that the hostname is the more
+      // user identifiable resource name and what user expects
+      // as the resource name.
+      resourceName: node.hostname,
+      node,
+    });
+    nextStep();
+  };
+
   useEffect(() => {
     if (isPolling) {
       const id = window.setTimeout(() => setShowHint(true), showHintTimeout);
@@ -196,7 +229,7 @@ export function SetupConnect(
       {pollingStatus}
 
       <ActionButtons
-        onProceed={props.nextStep}
+        onProceed={handleNextStep}
         disableProceed={!node}
         onPrev={props.prevStep}
       />
@@ -225,6 +258,7 @@ export function SetupConnect(
 export const usePollForConnectMyComputerNode = (args: {
   username: string;
   clusterId: string;
+  reloadUser: boolean;
   pingInterval: number;
 }): {
   node: Node | undefined;
@@ -237,6 +271,10 @@ export const usePollForConnectMyComputerNode = (args: {
   const node = usePoll(
     useCallback(
       async signal => {
+        if (args.reloadUser) {
+          await ctx.userService.reloadUser(signal);
+        }
+
         const request = {
           query: `labels["${constants.ConnectMyComputerNodeOwnerLabel}"] == "${args.username}"`,
           // An arbitrary limit where we bank on the fact that no one is going to have 50 Connect My
@@ -268,7 +306,13 @@ export const usePollForConnectMyComputerNode = (args: {
           return node;
         }
       },
-      [ctx.nodeService, args.clusterId, args.username]
+      [
+        ctx.nodeService,
+        ctx.userService,
+        args.clusterId,
+        args.username,
+        args.reloadUser,
+      ]
     ),
     isPolling,
     args.pingInterval

web/packages/teleport/src/Discover/ConnectMyComputer/TestConnection/TestConnection.story.tsx

@@ -0,0 +1,137 @@
+/**
+ * Copyright 2023 Gravitational, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import React from 'react';
+import { MemoryRouter } from 'react-router';
+import { initialize, mswLoader } from 'msw-storybook-addon';
+import { rest } from 'msw';
+
+import { nodes } from 'teleport/Nodes/fixtures';
+import { ContextProvider } from 'teleport';
+import cfg from 'teleport/config';
+import { UserContext } from 'teleport/User/UserContext';
+import { createTeleportContext } from 'teleport/mocks/contexts';
+import { makeDefaultUserPreferences } from 'teleport/services/userPreferences/userPreferences';
+
+import { TestConnection } from './TestConnection';
+
+export default {
+  title: 'Teleport/Discover/ConnectMyComputer/TestConnection',
+  loaders: [mswLoader],
+};
+
+initialize();
+
+const node = nodes[0];
+
+const agentStepProps = {
+  prevStep: () => {},
+  nextStep: () => {},
+  agentMeta: { resourceName: node.hostname, node, agentMatcherLabels: [] },
+};
+
+export const Story = () => {
+  return (
+    <Provider>
+      <TestConnection {...agentStepProps} />
+    </Provider>
+  );
+};
+
+Story.parameters = {
+  msw: {
+    handlers: [
+      rest.post(cfg.api.webRenewTokenPath, (req, res, ctx) =>
+        res(ctx.json({}))
+      ),
+      rest.get(cfg.api.nodesPath, (req, res, ctx) =>
+        res(ctx.json({ items: [node] }))
+      ),
+    ],
+  },
+};
+
+export const ReloadUserProcessing = () => {
+  return (
+    <Provider>
+      <TestConnection {...agentStepProps} />
+    </Provider>
+  );
+};
+
+ReloadUserProcessing.parameters = {
+  msw: {
+    handlers: [
+      rest.post(cfg.api.webRenewTokenPath, (req, res, ctx) =>
+        res(ctx.delay('infinite'))
+      ),
+    ],
+  },
+};
+
+export const ReloadUserError = () => {
+  return (
+    <Provider>
+      <TestConnection {...agentStepProps} />
+    </Provider>
+  );
+};
+
+ReloadUserError.parameters = {
+  msw: {
+    handlers: [
+      // The first handler returns an error immediately. Subsequent requests return after a delay so
+      // that we can show a spinner after clicking on "Retry".
+      rest.post(cfg.api.webRenewTokenPath, (req, res, ctx) =>
+        res.once(
+          ctx.status(500),
+          ctx.json({ message: 'Could not renew session' })
+        )
+      ),
+      rest.post(cfg.api.webRenewTokenPath, (req, res, ctx) =>
+        res(
+          ctx.delay(1000),
+          ctx.status(500),
+          ctx.json({ message: 'Could not renew session' })
+        )
+      ),
+    ],
+  },
+};
+
+const Provider = ({ children }) => {
+  const ctx = createTeleportContext();
+
+  const preferences = makeDefaultUserPreferences();
+  const updatePreferences = () => Promise.resolve();
+  const getClusterPinnedResources = () => Promise.resolve([]);
+  const updateClusterPinnedResources = () => Promise.resolve();
+
+  return (
+    <MemoryRouter>
+      <UserContext.Provider
+        value={{
+          preferences,
+          updatePreferences,
+          getClusterPinnedResources,
+          updateClusterPinnedResources,
+        }}
+      >
+        <ContextProvider ctx={ctx}>{children}</ContextProvider>
+      </UserContext.Provider>
+    </MemoryRouter>
+  );
+};