coder · kylecarbs · Dec 18, 2022 · Dec 5, 2022 · Dec 6, 2022 · Dec 6, 2022
diff --git a/cli/root.go b/cli/root.go
@@ -97,6 +97,7 @@ func Core() []*cobra.Command {
 		users(),
 		versionCmd(),
 		workspaceAgent(),
+		vscodeipcCmd(),
 	}
 }
 

diff --git a/cli/speedtest.go b/cli/speedtest.go
@@ -71,7 +71,7 @@ func speedtest() *cobra.Command {
 						return ctx.Err()
 					case <-ticker.C:
 					}
-					dur, err := conn.Ping(ctx)
+					dur, p2p, err := conn.Ping(ctx)
 					if err != nil {
 						continue
 					}
@@ -80,7 +80,7 @@ func speedtest() *cobra.Command {
 						continue
 					}
 					peer := status.Peer[status.Peers()[0]]
-					if peer.CurAddr == "" && direct {
+					if !p2p && direct {
 						cmd.Printf("Waiting for a direct connection... (%dms via %s)\n", dur.Milliseconds(), peer.Relay)
 						continue
 					}

diff --git a/cli/vscodeipc.go b/cli/vscodeipc.go
@@ -0,0 +1,68 @@
+package cli
+
+import (
+	"fmt"
+	"net"
+	"net/http"
+	"net/url"
+	"os"
+
+	"github.com/google/uuid"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+
+	"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
+	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")
+			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()
+			url, err := url.Parse(rawURL)
+			if err != nil {
+				return err
+			}
+			agentID, err := uuid.Parse(args[0])
+			if err != nil {
+				return err
+			}
+			client := codersdk.New(url)
+			client.SetSessionToken(token)
+
+			handler, closer, err := vscodeipc.New(cmd.Context(), client, agentID, nil)
+			if err != nil {
+				return err
+			}
+			defer closer.Close()
+			server := http.Server{
+				Handler: handler,
+			}
+			cmd.Printf("Ready\n")
+			return server.Serve(listener)
+		},
+	}
+	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
@@ -0,0 +1,269 @@
+package vscodeipc
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/google/uuid"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/xerrors"
+	"tailscale.com/tailcfg"
+
+	"github.com/coder/coder/agent"
+	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/codersdk"
+)
+
+// New creates a VS Code IPC client that can be used to communicate with workspaces.
+//
+// Creating this IPC was required instead of using SSH, because we're unable to get
+// connection information to display in the bottom-bar when using SSH. It's possible
+// we could jank around this (maybe by using a temporary SSH host), but that's not
+// ideal.
+//
+// This persists a single workspace connection, and lets you execute commands, check
+// for network information, and forward ports.
+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{}
+	}
+	// We need this to track upload and download!
+	options.EnableTrafficStats = true
+
+	agentConn, err := client.DialWorkspaceAgent(ctx, agentID, options)
+	if err != nil {
+		return nil, nil, err
+	}
+	api := &api{
+		agentConn: agentConn,
+	}
+	r := chi.NewRouter()
+	r.Get("/port/{port}", api.port)
+	r.Get("/network", api.network)
+	r.Post("/execute", api.execute)
+	return r, api, nil
+}
+
+type api struct {
+	agentConn     *codersdk.AgentConn
+	sshClient     *ssh.Client
+	sshClientErr  error
+	sshClientOnce sync.Once
+
+	lastNetwork time.Time
+}
+
+func (api *api) Close() error {
+	if api.sshClient != nil {
+		api.sshClient.Close()
+	}
+	return api.agentConn.Close()
+}
+
+type NetworkResponse struct {
+	P2P              bool               `json:"p2p"`
+	Latency          float64            `json:"latency"`
+	PreferredDERP    string             `json:"preferred_derp"`
+	DERPLatency      map[string]float64 `json:"derp_latency"`
+	UploadBytesSec   int64              `json:"upload_bytes_sec"`
+	DownloadBytesSec int64              `json:"download_bytes_sec"`
+}
+
+// port accepts an HTTP request to dial a port on the workspace agent.
+// It uses an HTTP connection upgrade to transfer the connection to TCP.
+func (api *api) port(w http.ResponseWriter, r *http.Request) {
+	port, err := strconv.Atoi(chi.URLParam(r, "port"))
+	if err != nil {
+		httpapi.Write(r.Context(), w, http.StatusBadRequest, codersdk.Response{
+			Message: "Port must be an integer!",
+		})
+		return
+	}
+	remoteConn, err := api.agentConn.DialContext(r.Context(), "tcp", fmt.Sprintf("127.0.0.1:%d", port))
+	if err != nil {
+		httpapi.InternalServerError(w, err)
+		return
+	}
+	defer remoteConn.Close()
+
+	// Upgrade an switch to TCP!
+	w.Header().Set("Connection", "Upgrade")
+	w.Header().Set("Upgrade", "tcp")
+	w.WriteHeader(http.StatusSwitchingProtocols)
+
+	hijacker, ok := w.(http.Hijacker)
+	if !ok {
+		httpapi.InternalServerError(w, xerrors.Errorf("unable to hijack connection: %T", w))
+		return
+	}
+
+	localConn, brw, err := hijacker.Hijack()
+	if err != nil {
+		httpapi.InternalServerError(w, err)
+		return
+	}
+
+	_ = brw.Flush()
+	defer localConn.Close()
+	agent.Bicopy(r.Context(), localConn, remoteConn)
+}
+
+// network returns network information about the workspace.
+func (api *api) network(w http.ResponseWriter, r *http.Request) {
+	// Ping the workspace agent to get the latency.
+	latency, p2p, err := api.agentConn.Ping(r.Context())
+	if err != nil {
+		httpapi.Write(r.Context(), w, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to ping the workspace agent.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	node := api.agentConn.Node()
+	derpMap := api.agentConn.DERPMap()
+	derpLatency := map[string]float64{}
+
+	// Convert DERP region IDs to friendly names for display in the UI.
+	for rawRegion, latency := range node.DERPLatency {
+		regionParts := strings.SplitN(rawRegion, "-", 2)
+		regionID, err := strconv.Atoi(regionParts[0])
+		if err != nil {
+			continue
+		}
+		region, found := derpMap.Regions[regionID]
+		if !found {
+			// It's possible that a workspace agent is using an old DERPMap
+			// and reports regions that do not exist. If that's the case,
+			// report the region as unknown!
+			region = &tailcfg.DERPRegion{
+				RegionID:   regionID,
+				RegionName: fmt.Sprintf("Unnamed %d", regionID),
+			}
+		}
+		// Convert the microseconds to milliseconds.
+		derpLatency[region.RegionName] = latency * 1000
+	}
+
+	totalRx := uint64(0)
+	totalTx := uint64(0)
+	for _, stat := range api.agentConn.ExtractTrafficStats() {
+		totalRx += stat.RxBytes
+		totalTx += stat.TxBytes
+	}
+	dur := time.Since(api.lastNetwork)
+
+	uploadSecs := float64(totalTx) / dur.Seconds()
+	downloadSecs := float64(totalRx) / dur.Seconds()
+
+	api.lastNetwork = time.Now()
+
+	httpapi.Write(r.Context(), w, http.StatusOK, NetworkResponse{
+		P2P:              p2p,
+		Latency:          float64(latency.Microseconds()) / 1000,
+		PreferredDERP:    derpMap.Regions[node.PreferredDERP].RegionName,
+		DERPLatency:      derpLatency,
+		UploadBytesSec:   int64(uploadSecs),
+		DownloadBytesSec: int64(downloadSecs),
+	})
+}
+
+type ExecuteRequest struct {
+	Command string `json:"command"`
+}
+
+type ExecuteResponse struct {
+	Data     string `json:"data"`
+	ExitCode *int   `json:"exit_code"`
+}
+
+// execute runs the command provided, streams the output back, and returns the exit code.
+func (api *api) execute(w http.ResponseWriter, r *http.Request) {
+	var req ExecuteRequest
+	if !httpapi.Read(r.Context(), w, r, &req) {
+		return
+	}
+	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.
+		api.sshClient, api.sshClientErr = api.agentConn.SSHClient(context.Background())
+	})
+	if api.sshClientErr != nil {
+		fmt.Printf("WE GOT TO BEGIN ERR! %s\n", api.sshClientErr)
+		httpapi.Write(r.Context(), w, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to create SSH client.",
+			Detail:  api.sshClientErr.Error(),
+		})
+		return
+	}
+	session, err := api.sshClient.NewSession()
+	if err != nil {
+		httpapi.Write(r.Context(), w, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to create SSH session.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	defer session.Close()
+	f, ok := w.(http.Flusher)
+	if !ok {
+		panic("http.ResponseWriter is not http.Flusher")
+	}
+
+	execWriter := &execWriter{w, f}
+	session.Stdout = execWriter
+	session.Stderr = execWriter
+	session.Stdin = strings.NewReader(req.Command + "\n")
+	err = session.Start("sh")
+	if err != nil {
+		httpapi.Write(r.Context(), w, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to start SSH session.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	err = session.Wait()
+
+	writeExit := func(exitCode int) {
+		data, _ := json.Marshal(&ExecuteResponse{
+			ExitCode: &exitCode,
+		})
+		_, _ = w.Write(data)
+		f.Flush()
+	}
+
+	if err != nil {
+		var exitError *ssh.ExitError
+		if errors.As(err, &exitError) {
+			writeExit(exitError.ExitStatus())
+			return
+		}
+	}
+	writeExit(0)
+}
+
+type execWriter struct {
+	w http.ResponseWriter
+	f http.Flusher
+}
+
+func (e *execWriter) Write(data []byte) (int, error) {
+	js, err := json.Marshal(&ExecuteResponse{
+		Data: string(data),
+	})
+	if err != nil {
+		return 0, err
+	}
+	_, _ = e.w.Write(js)
+	e.f.Flush()
+	return len(data), nil
+}