Fix requested changes

coder · kylecarbs · Dec 18, 2022 · Dec 5, 2022 · Dec 6, 2022 · Dec 6, 2022
commit ab45e4443c6cb8289ead0bcfe629bf7be2bbb0dd
diff --git a/cli/vscodeipc.go b/cli/vscodeipc.go
@@ -66,7 +66,7 @@ func vscodeipcCmd() *cobra.Command {
 				Handler: handler,
 			}
 			defer server.Close()
-			cmd.Printf("%d\n", addr.Port)
+			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "%s\n", addr.String())
 			errChan := make(chan error, 1)
 			go func() {
 				err := server.Serve(listener)

diff --git a/cli/vscodeipc/vscodeipc.go b/cli/vscodeipc/vscodeipc.go
@@ -35,7 +35,7 @@ import (
 //
 // 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
+// which calls this function. This API must maintain backward 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 {
@@ -54,28 +54,12 @@ func New(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, option
 	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.Use(sessionTokenMiddleware(client.SessionToken()))
+	r.Route("/v1", func(r chi.Router) {
+		r.Get("/port/{port}", api.port)
+		r.Get("/network", api.network)
+		r.Post("/execute", api.execute)
 	})
-	r.Get("/port/{port}", api.port)
-	r.Get("/network", api.network)
-	r.Post("/execute", api.execute)
 	return r, api, nil
 }
 
@@ -137,9 +121,9 @@ func (api *api) port(w http.ResponseWriter, r *http.Request) {
 		httpapi.InternalServerError(w, err)
 		return
 	}
+	defer localConn.Close()
 
 	_ = brw.Flush()
-	defer localConn.Close()
 	agent.Bicopy(r.Context(), localConn, remoteConn)
 }
 
@@ -222,6 +206,9 @@ func (api *api) execute(w http.ResponseWriter, r *http.Request) {
 	api.sshClientOnce.Do(func() {
 		// The SSH client is lazily created because it's not needed for
 		// all requests. It's only needed for the execute endpoint.
+		//
+		// It's alright if this fails on the first execution, because
+		// a new instance of this API is created for each remote SSH request.
 		api.sshClient, api.sshClientErr = api.agentConn.SSHClient(context.Background())
 	})
 	if api.sshClientErr != nil {
@@ -292,7 +279,32 @@ func (e *execWriter) Write(data []byte) (int, error) {
 	if err != nil {
 		return 0, err
 	}
-	_, _ = e.w.Write(js)
+	_, err = e.w.Write(js)
+	if err != nil {
+		return 0, err
+	}
 	e.f.Flush()
 	return len(data), nil
 }
+
+func sessionTokenMiddleware(sessionToken string) func(h http.Handler) http.Handler {
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			token := r.Header.Get("Coder-IPC-Token")
+			if token == "" {
+				httpapi.Write(r.Context(), w, http.StatusUnauthorized, codersdk.Response{
+					Message: "A session token must be provided in the `Coder-IPC-Token` header.",
+				})
+				return
+			}
+			if token != 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)
+		})
+	}
+}
diff --git a/cli/vscodeipc/vscodeipc_test.go b/cli/vscodeipc/vscodeipc_test.go
@@ -115,7 +115,7 @@ func TestVSCodeIPC(t *testing.T) {
 		handler.ServeHTTP(res, req)
 		network := &vscodeipc.NetworkResponse{}
 		err = json.NewDecoder(res.Body).Decode(&network)
-		require.NoError(t, err)
+		assert.NoError(t, err)
 		return network.Latency != 0
 	}, testutil.WaitLong, testutil.IntervalFast)
 

diff --git a/cli/vscodeipc_test.go b/cli/vscodeipc_test.go
@@ -1,8 +1,7 @@
 package cli_test
 
 import (
-	"bytes"
-	"fmt"
+	"io"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -20,8 +19,8 @@ func TestVSCodeIPC(t *testing.T) {
 		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)
+		rdr, wtr := io.Pipe()
+		cmd.SetOut(wtr)
 		ctx, cancelFunc := testutil.Context(t)
 		defer cancelFunc()
 		done := make(chan error, 1)
@@ -30,14 +29,14 @@ func TestVSCodeIPC(t *testing.T) {
 			done <- err
 		}()
 
-		var line string
+		buf := make([]byte, 64)
 		require.Eventually(t, func() bool {
-			t.Log("Looking for port!")
+			t.Log("Looking for address!")
 			var err error
-			line, err = buf.ReadString('\n')
+			_, err = rdr.Read(buf)
 			return err == nil
 		}, testutil.WaitMedium, testutil.IntervalFast)
-		t.Logf("Port: %s\n", line)
+		t.Logf("Address: %s\n", buf)
 
 		cancelFunc()
 		<-done