feat: add endpoint for netstat web socket

code-asher · code-asher · commit d1496edda09d · 2022-05-27T16:22:45.000-05:00
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -285,6 +285,7 @@ func New(options *Options) *API {
 				r.Get("/", api.workspaceAgent)
 				r.Get("/dial", api.workspaceAgentDial)
 				r.Get("/turn", api.workspaceAgentTurn)
+				r.Get("/netstat", api.workspaceAgentNetstat)
 				r.Get("/pty", api.workspaceAgentPTY)
 				r.Get("/iceservers", api.workspaceAgentICEServers)
 			})
diff --git a/coderd/coderd_test.go b/coderd/coderd_test.go
@@ -139,6 +139,7 @@ func TestAuthorizeAllEndpoints(t *testing.T) {
 		"GET:/api/v2/workspaceagents/{workspaceagent}":            {NoAuthorize: true},
 		"GET:/api/v2/workspaceagents/{workspaceagent}/dial":       {NoAuthorize: true},
 		"GET:/api/v2/workspaceagents/{workspaceagent}/iceservers": {NoAuthorize: true},
+		"GET:/api/v2/workspaceagents/{workspaceagent}/netstat":    {NoAuthorize: true},
 		"GET:/api/v2/workspaceagents/{workspaceagent}/pty":        {NoAuthorize: true},
 		"GET:/api/v2/workspaceagents/{workspaceagent}/turn":       {NoAuthorize: true},
 
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -501,3 +501,55 @@ func convertWorkspaceAgent(dbAgent database.WorkspaceAgent, agentUpdateFrequency
 
 	return workspaceAgent, nil
 }
+
+// workspaceAgentNetstat sends listening ports as `agent.NetstatResponse` on an
+// interval.
+func (api *API) workspaceAgentNetstat(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitMutex.Lock()
+	api.websocketWaitGroup.Add(1)
+	api.websocketWaitMutex.Unlock()
+	defer api.websocketWaitGroup.Done()
+
+	workspaceAgent := httpmw.WorkspaceAgentParam(r)
+	apiAgent, err := convertWorkspaceAgent(workspaceAgent, api.AgentConnectionUpdateFrequency)
+	if err != nil {
+		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
+			Message: fmt.Sprintf("convert workspace agent: %s", err),
+		})
+		return
+	}
+	if apiAgent.Status != codersdk.WorkspaceAgentConnected {
+		httpapi.Write(rw, http.StatusPreconditionRequired, httpapi.Response{
+			Message: fmt.Sprintf("agent must be in the connected state: %s", apiAgent.Status),
+		})
+		return
+	}
+
+	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
+			Message: fmt.Sprintf("accept websocket: %s", err),
+		})
+		return
+	}
+	defer func() {
+		_ = conn.Close(websocket.StatusNormalClosure, "ended")
+	}()
+	wsNetConn := websocket.NetConn(r.Context(), conn, websocket.MessageBinary)
+	agentConn, err := api.dialWorkspaceAgent(r, workspaceAgent.ID)
+	if err != nil {
+		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial workspace agent: %s", err))
+		return
+	}
+	defer agentConn.Close()
+	ptNetConn, err := agentConn.Netstat(r.Context())
+	if err != nil {
+		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))
+		return
+	}
+	defer ptNetConn.Close()
+
+	agent.Bicopy(r.Context(), wsNetConn, ptNetConn)
+}
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
+	"fmt"
 	"runtime"
 	"strings"
 	"testing"
@@ -264,3 +265,67 @@ func TestWorkspaceAgentPTY(t *testing.T) {
 	expectLine(matchEchoCommand)
 	expectLine(matchEchoOutput)
 }
+
+func TestWorkspaceAgentNetstat(t *testing.T) {
+	t.Parallel()
+
+	client, coderAPI := coderdtest.NewWithAPI(t, nil)
+	user := coderdtest.CreateFirstUser(t, client)
+	daemonCloser := coderdtest.NewProvisionerDaemon(t, coderAPI)
+	authToken := uuid.NewString()
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:           echo.ParseComplete,
+		ProvisionDryRun: echo.ProvisionComplete,
+		Provision: []*proto.Provision_Response{{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					Resources: []*proto.Resource{{
+						Name: "example",
+						Type: "aws_instance",
+						Agents: []*proto.Agent{{
+							Id: uuid.NewString(),
+							Auth: &proto.Agent_Token{
+								Token: authToken,
+							},
+						}},
+					}},
+				},
+			},
+		}},
+	})
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+	daemonCloser.Close()
+
+	agentClient := codersdk.New(client.URL)
+	agentClient.SessionToken = authToken
+	agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
+		Logger: slogtest.Make(t, nil),
+	})
+	t.Cleanup(func() {
+		_ = agentCloser.Close()
+	})
+	resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
+
+	conn, err := client.WorkspaceAgentNetstat(context.Background(), resources[0].Agents[0].ID)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	decoder := json.NewDecoder(conn)
+
+	expectNetstat := func() {
+		var res agent.NetstatResponse
+		err = decoder.Decode(&res)
+		require.NoError(t, err)
+
+		if runtime.GOOS == "linux" || runtime.GOOS == "windows" {
+			require.NotNil(t, res.Ports)
+		} else {
+			require.Equal(t, fmt.Sprintf("Port scanning is not supported on %s", runtime.GOOS), res.Error)
+		}
+	}
+
+	expectNetstat()
+}
diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -369,6 +369,36 @@ func (c *Client) WorkspaceAgentReconnectingPTY(ctx context.Context, agentID, rec
 	return websocket.NetConn(ctx, conn, websocket.MessageBinary), nil
 }
 
+// WorkspaceAgentNetstat sends listening ports as `agent.NetstatResponse` on an
+// interval.
+func (c *Client) WorkspaceAgentNetstat(ctx context.Context, agentID uuid.UUID) (net.Conn, error) {
+	serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/netstat", agentID))
+	if err != nil {
+		return nil, xerrors.Errorf("parse url: %w", err)
+	}
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, xerrors.Errorf("create cookie jar: %w", err)
+	}
+	jar.SetCookies(serverURL, []*http.Cookie{{
+		Name:  httpmw.SessionTokenKey,
+		Value: c.SessionToken,
+	}})
+	httpClient := &http.Client{
+		Jar: jar,
+	}
+	conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{
+		HTTPClient: httpClient,
+	})
+	if err != nil {
+		if res == nil {
+			return nil, err
+		}
+		return nil, readBodyAsError(res)
+	}
+	return websocket.NetConn(ctx, conn, websocket.MessageBinary), nil
+}
+
 func (c *Client) turnProxyDialer(ctx context.Context, httpClient *http.Client, path string) proxy.Dialer {
 	return turnconn.ProxyDialer(func() (net.Conn, error) {
 		turnURL, err := c.URL.Parse(path)
