feat: hook up port dropdown to workspace page

code-asher · code-asher · commit 91a194dde979 · 2022-05-27T16:22:48.000-05:00
diff --git a/site/src/components/Resources/Resources.tsx b/site/src/components/Resources/Resources.tsx
@@ -1,18 +1,21 @@
+import Button from "@material-ui/core/Button"
 import { makeStyles, Theme } from "@material-ui/core/styles"
 import Table from "@material-ui/core/Table"
 import TableBody from "@material-ui/core/TableBody"
 import TableCell from "@material-ui/core/TableCell"
 import TableHead from "@material-ui/core/TableHead"
 import TableRow from "@material-ui/core/TableRow"
+import CompareArrowsIcon from "@material-ui/icons/CompareArrows"
 import useTheme from "@material-ui/styles/useTheme"
 import React from "react"
-import { Workspace, WorkspaceResource } from "../../api/typesGenerated"
+import { Workspace, WorkspaceAgent, WorkspaceResource } from "../../api/typesGenerated"
 import { getDisplayAgentStatus } from "../../util/workspace"
 import { TableHeaderRow } from "../TableHeaders/TableHeaders"
 import { TerminalLink } from "../TerminalLink/TerminalLink"
 import { WorkspaceSection } from "../WorkspaceSection/WorkspaceSection"
 
 const Language = {
+  portForwardLabel: "Port forward",
   resources: "Resources",
   resourceLabel: "Resource",
   agentsLabel: "Agents",
@@ -22,12 +25,18 @@ const Language = {
 }
 
 interface ResourcesProps {
+  handleOpenPortForward: (agent: WorkspaceAgent, anchorEl: HTMLElement) => void
   resources?: WorkspaceResource[]
   getResourcesError?: Error
   workspace: Workspace
 }
 
-export const Resources: React.FC<ResourcesProps> = ({ resources, getResourcesError, workspace }) => {
+export const Resources: React.FC<ResourcesProps> = ({
+  handleOpenPortForward,
+  resources,
+  getResourcesError,
+  workspace,
+}) => {
   const styles = useStyles()
   const theme: Theme = useTheme()
 
@@ -89,12 +98,22 @@ export const Resources: React.FC<ResourcesProps> = ({ resources, getResourcesErr
                     </TableCell>
                     <TableCell>
                       {agent.status === "connected" && (
-                        <TerminalLink
-                          className={styles.accessLink}
-                          workspaceName={workspace.name}
-                          agentName={agent.name}
-                          userName={workspace.owner_name}
-                        />
+                        <>
+                          <TerminalLink
+                            className={styles.accessLink}
+                            workspaceName={workspace.name}
+                            agentName={agent.name}
+                            userName={workspace.owner_name}
+                          />
+                          <Button
+                            variant="text"
+                            className={styles.accessLink}
+                            onClick={(event) => handleOpenPortForward(agent, event.currentTarget)}
+                          >
+                            <CompareArrowsIcon />
+                            {Language.portForwardLabel}
+                          </Button>
+                        </>
                       )}
                     </TableCell>
                   </TableRow>
@@ -134,9 +153,16 @@ const useStyles = makeStyles((theme) => ({
   },
 
   accessLink: {
+    alignItems: "center",
     color: theme.palette.text.secondary,
     display: "flex",
-    alignItems: "center",
+    border: 0,
+    padding: 0,
+
+    "&:hover": {
+      backgroundColor: "unset",
+      textDecoration: "underline",
+    },
 
     "& svg": {
       width: 16,
diff --git a/site/src/components/Workspace/Workspace.tsx b/site/src/components/Workspace/Workspace.tsx
@@ -17,6 +17,7 @@ export interface WorkspaceProps {
   handleStop: () => void
   handleUpdate: () => void
   handleCancel: () => void
+  handleOpenPortForward: (agent: TypesGen.WorkspaceAgent, anchorEl: HTMLElement) => void
   workspace: TypesGen.Workspace
   resources?: TypesGen.WorkspaceResource[]
   getResourcesError?: Error
@@ -31,6 +32,7 @@ export const Workspace: React.FC<WorkspaceProps> = ({
   handleStop,
   handleUpdate,
   handleCancel,
+  handleOpenPortForward,
   workspace,
   resources,
   getResourcesError,
@@ -68,7 +70,12 @@ export const Workspace: React.FC<WorkspaceProps> = ({
 
           <WorkspaceStats workspace={workspace} />
 
-          <Resources resources={resources} getResourcesError={getResourcesError} workspace={workspace} />
+          <Resources
+            handleOpenPortForward={handleOpenPortForward}
+            resources={resources}
+            getResourcesError={getResourcesError}
+            workspace={workspace}
+          />
 
           <WorkspaceSection title="Timeline" contentsProps={{ className: styles.timelineContents }}>
             <BuildsTable builds={builds} className={styles.timelineTable} />
diff --git a/site/src/pages/WorkspacePage/WorkspacePage.tsx b/site/src/pages/WorkspacePage/WorkspacePage.tsx
@@ -1,12 +1,14 @@
 import { useMachine } from "@xstate/react"
-import React, { useEffect } from "react"
+import React, { useEffect, useState } from "react"
 import { useParams } from "react-router-dom"
 import { ErrorSummary } from "../../components/ErrorSummary/ErrorSummary"
 import { FullScreenLoader } from "../../components/Loader/FullScreenLoader"
 import { Margins } from "../../components/Margins/Margins"
+import { PortForwardDropdown } from "../../components/PortForwardDropdown/PortForwardDropdown"
 import { Stack } from "../../components/Stack/Stack"
 import { Workspace } from "../../components/Workspace/Workspace"
 import { firstOrItem } from "../../util/array"
+import { agentMachine } from "../../xServices/agent/agentXService"
 import { workspaceMachine } from "../../xServices/workspace/workspaceXService"
 
 export const WorkspacePage: React.FC = () => {
@@ -16,6 +18,10 @@ export const WorkspacePage: React.FC = () => {
   const [workspaceState, workspaceSend] = useMachine(workspaceMachine)
   const { workspace, resources, getWorkspaceError, getResourcesError, builds } = workspaceState.context
 
+  const [agentState, agentSend] = useMachine(agentMachine)
+  const { netstat } = agentState.context
+  const [portForwardAnchorEl, setPortForwardAnchorEl] = useState<HTMLElement>()
+
   /**
    * Get workspace, template, and organization on mount and whenever workspaceId changes.
    * workspaceSend should not change.
@@ -38,10 +44,24 @@ export const WorkspacePage: React.FC = () => {
             handleStop={() => workspaceSend("STOP")}
             handleUpdate={() => workspaceSend("UPDATE")}
             handleCancel={() => workspaceSend("CANCEL")}
+            handleOpenPortForward={(agent, anchorEl) => {
+              agentSend("CONNECT", { agentId: agent.id })
+              setPortForwardAnchorEl(anchorEl)
+            }}
             resources={resources}
             getResourcesError={getResourcesError instanceof Error ? getResourcesError : undefined}
             builds={builds}
           />
+          <PortForwardDropdown
+            open={!!portForwardAnchorEl}
+            anchorEl={portForwardAnchorEl}
+            netstat={netstat}
+            onClose={() => {
+              agentSend("DISCONNECT")
+              setPortForwardAnchorEl(undefined)
+            }}
+            urlFormatter={(port) => `${location.protocol}//${port}--${workspace.owner_name}--${location.host}`}
+          />
         </Stack>
       </Margins>
     )
diff --git a/site/src/xServices/agent/agentXService.ts b/site/src/xServices/agent/agentXService.ts
@@ -0,0 +1,133 @@
+import { assign, createMachine } from "xstate"
+import * as Types from "../../api/types"
+import { errorString } from "../../util/error"
+
+export interface AgentContext {
+  agentId?: string
+  netstat?: Types.NetstatResponse
+  websocket?: WebSocket
+}
+
+export type AgentEvent =
+  | { type: "CONNECT"; agentId: string }
+  | { type: "STAT"; data: Types.NetstatResponse }
+  | { type: "DISCONNECT" }
+
+export const agentMachine = createMachine(
+  {
+    tsTypes: {} as import("./agentXService.typegen").Typegen0,
+    schema: {
+      context: {} as AgentContext,
+      events: {} as AgentEvent,
+      services: {} as {
+        connect: {
+          data: WebSocket
+        }
+      },
+    },
+    id: "agentState",
+    initial: "disconnected",
+    states: {
+      connecting: {
+        invoke: {
+          src: "connect",
+          id: "connect",
+          onDone: [
+            {
+              actions: ["assignWebsocket", "clearNetstat"],
+              target: "connected",
+            },
+          ],
+          onError: [
+            {
+              actions: "assignWebsocketError",
+              target: "disconnected",
+            },
+          ],
+        },
+      },
+      connected: {
+        on: {
+          STAT: {
+            actions: "assignNetstat",
+          },
+          DISCONNECT: {
+            actions: ["disconnect", "clearNetstat"],
+            target: "disconnected",
+          },
+        },
+      },
+      disconnected: {
+        on: {
+          CONNECT: {
+            actions: "assignConnection",
+            target: "connecting",
+          },
+        },
+      },
+    },
+  },
+  {
+    services: {
+      connect: (context) => (send) => {
+        return new Promise<WebSocket>((resolve, reject) => {
+          if (!context.agentId) {
+            return reject("agent ID is not set")
+          }
+          const proto = location.protocol === "https:" ? "wss:" : "ws:"
+          const socket = new WebSocket(`${proto}//${location.host}/api/v2/workspaceagents/${context.agentId}/netstat`)
+          socket.binaryType = "arraybuffer"
+          socket.addEventListener("open", () => {
+            resolve(socket)
+          })
+          socket.addEventListener("error", (error) => {
+            reject(error)
+          })
+          socket.addEventListener("close", () => {
+            send({
+              type: "DISCONNECT",
+            })
+          })
+          socket.addEventListener("message", (event) => {
+            try {
+              send({
+                type: "STAT",
+                data: JSON.parse(new TextDecoder().decode(event.data)),
+              })
+            } catch (error) {
+              send({
+                type: "STAT",
+                data: {
+                  error: errorString(error),
+                },
+              })
+            }
+          })
+        })
+      },
+    },
+    actions: {
+      assignConnection: assign((context, event) => ({
+        ...context,
+        agentId: event.agentId,
+      })),
+      assignWebsocket: assign({
+        websocket: (_, event) => event.data,
+      }),
+      assignWebsocketError: assign({
+        netstat: (_, event) => ({ error: errorString(event.data) }),
+      }),
+      clearNetstat: assign((context: AgentContext) => ({
+        ...context,
+        netstat: undefined,
+      })),
+      assignNetstat: assign({
+        netstat: (_, event) => event.data,
+      }),
+      disconnect: (context: AgentContext) => {
+        // Code 1000 is a successful exit!
+        context.websocket?.close(1000)
+      },
+    },
+  },
+)
