coder · kylecarbs · Dec 18, 2022 · Dec 5, 2022 · Dec 6, 2022 · Dec 6, 2022
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -65,6 +65,8 @@ func setupWorkspaceForAgent(t *testing.T, mutate func([]*proto.Agent) []*proto.A
 	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
 	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+	workspace, err := client.Workspace(context.Background(), workspace.ID)
+	require.NoError(t, err)
 
 	return client, workspace, agentToken
 }

diff --git a/cli/vscodeipc.go b/cli/vscodeipc.go
@@ -5,41 +5,46 @@ import (
 	"net"
 	"net/http"
 	"net/url"
-	"os"
 
 	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/cli/vscodeipc"
 	"github.com/coder/coder/codersdk"
 )
 
 // vscodeipcCmd spawns a local HTTP server on the provided port that listens to messages.
 // It's made for use by the Coder VS Code extension. See: https://github.com/coder/vscode-coder
 func vscodeipcCmd() *cobra.Command {
-	var port uint16
+	var (
+		rawURL string
+		token  string
+		port   uint16
+	)
 	cmd := &cobra.Command{
 		Use:    "vscodeipc <workspace-agent>",
 		Args:   cobra.ExactArgs(1),
 		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			rawURL := os.Getenv("CODER_URL")
 			if rawURL == "" {
 				return xerrors.New("CODER_URL must be set!")
 			}
-			token := os.Getenv("CODER_TOKEN")
+			// token is validated in a header on each request to prevent
+			// unauthenticated clients from connecting.
 			if token == "" {
 				return xerrors.New("CODER_TOKEN must be set!")
 			}
-			if port == 0 {
-				return xerrors.Errorf("port must be specified!")
-			}
 			listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
 			if err != nil {
 				return xerrors.Errorf("listen: %w", err)
 			}
 			defer listener.Close()
+			addr, ok := listener.Addr().(*net.TCPAddr)
+			if !ok {
+				return xerrors.Errorf("listener.Addr() is not a *net.TCPAddr: %T", listener.Addr())
+			}
 			url, err := url.Parse(rawURL)
 			if err != nil {
 				return err
@@ -56,13 +61,27 @@ func vscodeipcCmd() *cobra.Command {
 				return err
 			}
 			defer closer.Close()
+			// nolint:gosec
 			server := http.Server{
 				Handler: handler,
 			}
-			cmd.Printf("Ready\n")
-			return server.Serve(listener)
+			defer server.Close()
+			cmd.Printf("%d\n", addr.Port)
+			errChan := make(chan error, 1)
+			go func() {
+				err := server.Serve(listener)
+				errChan <- err
+			}()
+			select {
+			case <-cmd.Context().Done():
+				return cmd.Context().Err()
+			case err := <-errChan:
+				return err
+			}
 		},
 	}
+	cliflag.StringVarP(cmd.Flags(), &rawURL, "url", "u", "CODER_URL", "", "The URL of the Coder instance!")
+	cliflag.StringVarP(cmd.Flags(), &token, "token", "t", "CODER_TOKEN", "", "The session token of the user!")
 	cmd.Flags().Uint16VarP(&port, "port", "p", 0, "The port to listen on!")
 	return cmd
 }
diff --git a/cli/vscodeipc/vscodeipc.go b/cli/vscodeipc/vscodeipc.go
@@ -32,6 +32,11 @@ import (
 //
 // This persists a single workspace connection, and lets you execute commands, check
 // for network information, and forward ports.
+//
+// The VS Code extension is located at https://github.com/coder/vscode-coder. The
+// extension downloads the slim binary from `/bin/*` and executes `coder vscodeipc`
+// which calls this function. This API must maintain backawards compatibility with
+// the extension to support prior versions of Coder.
 func New(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, options *codersdk.DialWorkspaceAgentOptions) (http.Handler, io.Closer, error) {
 	if options == nil {
 		options = &codersdk.DialWorkspaceAgentOptions{}
@@ -47,6 +52,27 @@ func New(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, option
 		agentConn: agentConn,
 	}
 	r := chi.NewRouter()
+	// This is to prevent unauthorized clients on the same machine from executing
+	// requests on behalf of the workspace.
+	r.Use(func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			token := r.Header.Get("Coder-Session-Token")
+			if token == "" {
+				httpapi.Write(r.Context(), w, http.StatusUnauthorized, codersdk.Response{
+					Message: "A session token must be provided in the `Coder-Session-Token` header.",
+				})
+				return
+			}
+			if token != client.SessionToken() {
+				httpapi.Write(r.Context(), w, http.StatusUnauthorized, codersdk.Response{
+					Message: "The session token provided doesn't match the one used to create the client.",
+				})
+				return
+			}
+			w.Header().Set("Access-Control-Allow-Origin", "*")
+			h.ServeHTTP(w, r)
+		})
+	})
 	r.Get("/port/{port}", api.port)
 	r.Get("/network", api.network)
 	r.Post("/execute", api.execute)
@@ -160,8 +186,9 @@ func (api *api) network(w http.ResponseWriter, r *http.Request) {
 		totalRx += stat.RxBytes
 		totalTx += stat.TxBytes
 	}
+	// Tracking the time since last request is required because
+	// ExtractTrafficStats() resets its counters after each call.
 	dur := time.Since(api.lastNetwork)
-
 	uploadSecs := float64(totalTx) / dur.Seconds()
 	downloadSecs := float64(totalRx) / dur.Seconds()
 
@@ -215,7 +242,10 @@ func (api *api) execute(w http.ResponseWriter, r *http.Request) {
 	defer session.Close()
 	f, ok := w.(http.Flusher)
 	if !ok {
-		panic("http.ResponseWriter is not http.Flusher")
+		httpapi.Write(r.Context(), w, http.StatusInternalServerError, codersdk.Response{
+			Message: fmt.Sprintf("http.ResponseWriter is not http.Flusher: %T", w),
+		})
+		return
 	}
 
 	execWriter := &execWriter{w, f}

diff --git a/cli/vscodeipc/vscodeipc_test.go b/cli/vscodeipc/vscodeipc_test.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 	"nhooyr.io/websocket"
@@ -39,34 +40,43 @@ func TestVSCodeIPC(t *testing.T) {
 	id := uuid.New()
 	derpMap := tailnettest.RunDERPAndSTUN(t)
 	coordinator := tailnet.NewCoordinator()
+	t.Cleanup(func() {
+		_ = coordinator.Close()
+	})
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case fmt.Sprintf("/api/v2/workspaceagents/%s/connection", id):
+			assert.Equal(t, r.Method, http.MethodGet)
 			httpapi.Write(ctx, w, http.StatusOK, codersdk.WorkspaceAgentConnectionInfo{
 				DERPMap: derpMap,
 			})
 			return
 		case fmt.Sprintf("/api/v2/workspaceagents/%s/coordinate", id):
+			assert.Equal(t, r.Method, http.MethodGet)
 			ws, err := websocket.Accept(w, r, nil)
 			require.NoError(t, err)
 			conn := websocket.NetConn(ctx, ws, websocket.MessageBinary)
 			_ = coordinator.ServeClient(conn, uuid.New(), id)
 			return
 		case "/api/v2/workspaceagents/me/version":
+			assert.Equal(t, r.Method, http.MethodPost)
 			w.WriteHeader(http.StatusOK)
 			return
 		case "/api/v2/workspaceagents/me/metadata":
+			assert.Equal(t, r.Method, http.MethodGet)
 			httpapi.Write(ctx, w, http.StatusOK, codersdk.WorkspaceAgentMetadata{
 				DERPMap: derpMap,
 			})
 			return
 		case "/api/v2/workspaceagents/me/coordinate":
+			assert.Equal(t, r.Method, http.MethodGet)
 			ws, err := websocket.Accept(w, r, nil)
 			require.NoError(t, err)
 			conn := websocket.NetConn(ctx, ws, websocket.MessageBinary)
 			_ = coordinator.ServeAgent(conn, id)
 			return
 		case "/api/v2/workspaceagents/me/report-stats":
+			assert.Equal(t, r.Method, http.MethodPost)
 			w.WriteHeader(http.StatusOK)
 			return
 		case "/":
@@ -80,6 +90,8 @@ func TestVSCodeIPC(t *testing.T) {
 	srvURL, _ := url.Parse(srv.URL)
 
 	client := codersdk.New(srvURL)
+	token := uuid.New().String()
+	client.SetSessionToken(token)
 	agentConn := agent.New(agent.Options{
 		Client:     client,
 		Filesystem: afero.NewMemMapFs(),
@@ -99,6 +111,7 @@ func TestVSCodeIPC(t *testing.T) {
 	require.Eventually(t, func() bool {
 		res := httptest.NewRecorder()
 		req := httptest.NewRequest(http.MethodGet, "/network", nil)
+		req.Header.Set("Coder-Session-Token", token)
 		handler.ServeHTTP(res, req)
 		network := &vscodeipc.NetworkResponse{}
 		err = json.NewDecoder(res.Body).Decode(&network)
@@ -109,6 +122,23 @@ func TestVSCodeIPC(t *testing.T) {
 	_, port, err := net.SplitHostPort(srvURL.Host)
 	require.NoError(t, err)
 
+	t.Run("NoSessionToken", func(t *testing.T) {
+		t.Parallel()
+		res := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/port/%s", port), nil)
+		handler.ServeHTTP(res, req)
+		require.Equal(t, http.StatusUnauthorized, res.Code)
+	})
+
+	t.Run("MismatchedSessionToken", func(t *testing.T) {
+		t.Parallel()
+		res := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/port/%s", port), nil)
+		req.Header.Set("Coder-Session-Token", uuid.NewString())
+		handler.ServeHTTP(res, req)
+		require.Equal(t, http.StatusUnauthorized, res.Code)
+	})
+
 	t.Run("Port", func(t *testing.T) {
 		// Tests that the port endpoint can be used for forward traffic.
 		// For this test, we simply use the already listening httptest server.
@@ -118,6 +148,7 @@ func TestVSCodeIPC(t *testing.T) {
 		defer output.Close()
 		res := &hijackable{httptest.NewRecorder(), output}
 		req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/port/%s", port), nil)
+		req.Header.Set("Coder-Session-Token", token)
 		go handler.ServeHTTP(res, req)
 
 		req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1/", nil)
@@ -147,6 +178,7 @@ func TestVSCodeIPC(t *testing.T) {
 			Command: "echo test",
 		})
 		req := httptest.NewRequest(http.MethodPost, "/execute", bytes.NewReader(data))
+		req.Header.Set("Coder-Session-Token", token)
 		handler.ServeHTTP(res, req)
 
 		decoder := json.NewDecoder(res.Body)

diff --git a/cli/vscodeipc_test.go b/cli/vscodeipc_test.go
@@ -0,0 +1,45 @@
+package cli_test
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/testutil"
+)
+
+func TestVSCodeIPC(t *testing.T) {
+	t.Parallel()
+	// Ensures the vscodeipc command outputs it's running port!
+	// This signifies to the caller that it's ready to accept requests.
+	t.Run("PortOutputs", func(t *testing.T) {
+		t.Parallel()
+		client, workspace, _ := setupWorkspaceForAgent(t, nil)
+		cmd, _ := clitest.New(t, "vscodeipc", workspace.LatestBuild.Resources[0].Agents[0].ID.String(),
+			"--token", client.SessionToken(), "--url", client.URL.String())
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		ctx, cancelFunc := testutil.Context(t)
+		defer cancelFunc()
+		done := make(chan error, 1)
+		go func() {
+			err := cmd.ExecuteContext(ctx)
+			done <- err
+		}()
+
+		var line string
+		require.Eventually(t, func() bool {
+			t.Log("Looking for port!")
+			var err error
+			line, err = buf.ReadString('\n')
+			return err == nil
+		}, testutil.WaitMedium, testutil.IntervalFast)
+		t.Logf("Port: %s\n", line)
+
+		cancelFunc()
+		<-done
+	})
+}
diff --git a/codersdk/agentconn.go b/codersdk/agentconn.go
@@ -139,6 +139,8 @@ func (c *AgentConn) AwaitReachable(ctx context.Context) bool {
 	return c.Conn.AwaitReachable(ctx, TailnetIP)
 }
 
+// Ping pings the agent and returns the round-trip time.
+// The bool returns true if the ping was made P2P.
 func (c *AgentConn) Ping(ctx context.Context) (time.Duration, bool, error) {
 	ctx, span := tracing.StartSpan(ctx)
 	defer span.End()

diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -408,6 +408,7 @@ func (c *Conn) Status() *ipnstate.Status {
 }
 
 // Ping sends a Disco ping to the Wireguard engine.
+// The bool returned is true if the ping was performed P2P.
 func (c *Conn) Ping(ctx context.Context, ip netip.Addr) (time.Duration, bool, error) {
 	errCh := make(chan error, 1)
 	prChan := make(chan *ipnstate.PingResult, 1)