coder · kylecarbs · Apr 11, 2022 · Apr 8, 2022 · Apr 8, 2022 · Apr 9, 2022
diff --git a/agent/agent.go b/agent/agent.go
@@ -101,6 +101,11 @@ func (a *agent) run(ctx context.Context) {
 
 func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 	go func() {
+		select {
+		case <-a.closed:
+			_ = conn.Close()
+		case <-conn.Closed():
+		}
 		<-conn.Closed()
 		a.connCloseWait.Done()
 	}()

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2,7 +2,12 @@ package agent_test
 
 import (
 	"context"
+	"fmt"
+	"io"
+	"net"
+	"os/exec"
 	"runtime"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -29,7 +34,8 @@ func TestAgent(t *testing.T) {
 	t.Parallel()
 	t.Run("SessionExec", func(t *testing.T) {
 		t.Parallel()
-		session := setupSSH(t)
+		session := setupSSHSession(t)
+
 		command := "echo test"
 		if runtime.GOOS == "windows" {
 			command = "cmd.exe /c echo test"
@@ -41,7 +47,7 @@ func TestAgent(t *testing.T) {
 
 	t.Run("GitSSH", func(t *testing.T) {
 		t.Parallel()
-		session := setupSSH(t)
+		session := setupSSHSession(t)
 		command := "sh -c 'echo $GIT_SSH_COMMAND'"
 		if runtime.GOOS == "windows" {
 			command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
@@ -53,7 +59,7 @@ func TestAgent(t *testing.T) {
 
 	t.Run("SessionTTY", func(t *testing.T) {
 		t.Parallel()
-		session := setupSSH(t)
+		session := setupSSHSession(t)
 		prompt := "$"
 		command := "bash"
 		if runtime.GOOS == "windows" {
@@ -76,9 +82,77 @@ func TestAgent(t *testing.T) {
 		err = session.Wait()
 		require.NoError(t, err)
 	})
+
+	t.Run("LocalForwarding", func(t *testing.T) {
+		t.Parallel()
+		random, err := net.Listen("tcp", "127.0.0.1:0")
+		require.NoError(t, err)
+		_ = random.Close()
+		tcpAddr, valid := random.Addr().(*net.TCPAddr)
+		require.True(t, valid)
+		randomPort := tcpAddr.Port
+
+		local, err := net.Listen("tcp", "127.0.0.1:0")
+		require.NoError(t, err)
+		tcpAddr, valid = local.Addr().(*net.TCPAddr)
+		require.True(t, valid)
+		localPort := tcpAddr.Port
+		done := make(chan struct{})
+		go func() {
+			conn, err := local.Accept()
+			require.NoError(t, err)
+			_ = conn.Close()
+			close(done)
+		}()
+
+		err = setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, localPort)}, []string{"echo", "test"}).Start()
+		require.NoError(t, err)
+
+		conn, err := net.Dial("tcp", "127.0.0.1:"+strconv.Itoa(localPort))
+		require.NoError(t, err)
+		conn.Close()
+		<-done
+	})
 }
 
-func setupSSH(t *testing.T) *ssh.Session {
+func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
+	agentConn := setupAgent(t)
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			ssh, err := agentConn.SSH()
+			require.NoError(t, err)
+			go io.Copy(conn, ssh)
+			go io.Copy(ssh, conn)
+		}
+	}()
+	t.Cleanup(func() {
+		_ = listener.Close()
+	})
+	tcpAddr, valid := listener.Addr().(*net.TCPAddr)
+	require.True(t, valid)
+	args := append(beforeArgs,
+		"-o", "HostName "+tcpAddr.IP.String(),
+		"-o", "Port "+strconv.Itoa(tcpAddr.Port),
+		"-o", "StrictHostKeyChecking=no", "host")
+	args = append(args, afterArgs...)
+	return exec.Command("ssh", args...)
+}
+
+func setupSSHSession(t *testing.T) *ssh.Session {
+	sshClient, err := setupAgent(t).SSHClient()
+	require.NoError(t, err)
+	session, err := sshClient.NewSession()
+	require.NoError(t, err)
+	return session
+}
+
+func setupAgent(t *testing.T) *agent.Conn {
 	client, server := provisionersdk.TransportPipe()
 	closer := agent.New(func(ctx context.Context, opts *peer.ConnOptions) (*peerbroker.Listener, error) {
 		return peerbroker.Listen(server, nil, opts)
@@ -100,14 +174,9 @@ func setupSSH(t *testing.T) *ssh.Session {
 	t.Cleanup(func() {
 		_ = conn.Close()
 	})
-	agentClient := &agent.Conn{
+
+	return &agent.Conn{
 		Negotiator: api,
 		Conn:       conn,
 	}
-	sshClient, err := agentClient.SSHClient()
-	require.NoError(t, err)
-	session, err := sshClient.NewSession()
-	require.NoError(t, err)
-
-	return session
 }
diff --git a/cli/cliui/cliui.go b/cli/cliui/cliui.go
@@ -26,6 +26,7 @@ var Styles = struct {
 	Checkmark,
 	Code,
 	Crossmark,
+	Error,
 	Field,
 	Keyword,
 	Paragraph,
@@ -41,6 +42,7 @@ var Styles = struct {
 	Checkmark:     defaultStyles.Checkmark,
 	Code:          defaultStyles.Code,
 	Crossmark:     defaultStyles.Error.Copy().SetString("✘"),
+	Error:         defaultStyles.Error,
 	Field:         defaultStyles.Code.Copy().Foreground(lipgloss.AdaptiveColor{Light: "#000000", Dark: "#FFFFFF"}),
 	Keyword:       defaultStyles.Keyword,
 	Paragraph:     defaultStyles.Paragraph,

diff --git a/cli/cliui/prompt.go b/cli/cliui/prompt.go
@@ -5,7 +5,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"io"
 	"os"
 	"os/signal"
 	"runtime"
@@ -45,11 +44,11 @@ func Prompt(cmd *cobra.Command, opts PromptOptions) (string, error) {
 		var line string
 		var err error
 
-		inFile, valid := cmd.InOrStdin().(*os.File)
-		if opts.Secret && valid && isatty.IsTerminal(inFile.Fd()) {
+		inFile, isInputFile := cmd.InOrStdin().(*os.File)
+		if opts.Secret && isInputFile && isatty.IsTerminal(inFile.Fd()) {
 			line, err = speakeasy.Ask("")
 		} else {
-			if !opts.IsConfirm && runtime.GOOS == "darwin" && valid {
+			if !opts.IsConfirm && runtime.GOOS == "darwin" && isInputFile {
 				var restore func()
 				restore, err = removeLineLengthLimit(int(inFile.Fd()))
 				if err != nil {
@@ -66,22 +65,7 @@ func Prompt(cmd *cobra.Command, opts PromptOptions) (string, error) {
 			// This enables multiline JSON to be pasted into an input, and have
 			// it parse properly.
 			if err == nil && (strings.HasPrefix(line, "{") || strings.HasPrefix(line, "[")) {
-				pipeReader, pipeWriter := io.Pipe()
-				defer pipeWriter.Close()
-				defer pipeReader.Close()
-				go func() {
-					_, _ = pipeWriter.Write([]byte(line))
-					_, _ = reader.WriteTo(pipeWriter)
-				}()
-				var rawMessage json.RawMessage
-				err := json.NewDecoder(pipeReader).Decode(&rawMessage)
-				if err == nil {
-					var buf bytes.Buffer
-					err = json.Compact(&buf, rawMessage)
-					if err == nil {
-						line = buf.String()
-					}
-				}
+				line, err = promptJSON(reader, line)
 			}
 		}
 		if err != nil {
@@ -118,3 +102,39 @@ func Prompt(cmd *cobra.Command, opts PromptOptions) (string, error) {
 		return "", Canceled
 	}
 }
+
+func promptJSON(reader *bufio.Reader, line string) (string, error) {
+	var data bytes.Buffer
+	for {
+		_, _ = data.WriteString(line)
+		var rawMessage json.RawMessage
+		err := json.Unmarshal(data.Bytes(), &rawMessage)
+		if err != nil {
+			if err.Error() != "unexpected end of JSON input" {
+				// If a real syntax error occurs in JSON,
+				// we want to return that partial line to the user.
+				err = nil
+				line = data.String()
+				break
+			}
+
+			// Read line-by-line. We can't use a JSON decoder
+			// here because it doesn't work by newline, so
+			// reads will block.
+			line, err = reader.ReadString('\n')
+			if err != nil {
+				break
+			}
+			continue
+		}
+		// Compacting the JSON makes it easier for parsing and testing.
+		rawJSON := data.Bytes()
+		data.Reset()
+		err = json.Compact(&data, rawJSON)
+		if err != nil {
+			return line, xerrors.Errorf("compact json: %w", err)
+		}
+		return data.String(), nil
+	}
+	return line, nil
+}
diff --git a/cli/cliui/resources.go b/cli/cliui/resources.go
@@ -0,0 +1,140 @@
+package cliui
+
+import (
+	"fmt"
+	"io"
+	"sort"
+	"strconv"
+
+	"github.com/jedib0t/go-pretty/v6/table"
+
+	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/codersdk"
+)
+
+type WorkspaceResourcesOptions struct {
+	WorkspaceName  string
+	HideAgentState bool
+	HideAccess     bool
+	Title          string
+}
+
+// WorkspaceResources displays the connection status and tree-view of provided resources.
+// ┌────────────────────────────────────────────────────────────────────────────┐
+// │ RESOURCE                     STATUS               ACCESS                   │
+// ├────────────────────────────────────────────────────────────────────────────┤
+// │ google_compute_disk.root     persistent                                    │
+// ├────────────────────────────────────────────────────────────────────────────┤
+// │ google_compute_instance.dev  ephemeral                                     │
+// │ └─ dev (linux, amd64)        ⦾ connecting [10s]    coder ssh dev.dev       │
+// ├────────────────────────────────────────────────────────────────────────────┤
+// │ kubernetes_pod.dev           ephemeral                                     │
+// │ ├─ go (linux, amd64)         ⦿ connected           coder ssh dev.go        │
+// │ └─ postgres (linux, amd64)   ⦾ disconnected [4s]   coder ssh dev.postgres  │
+// └────────────────────────────────────────────────────────────────────────────┘
+func WorkspaceResources(writer io.Writer, resources []codersdk.WorkspaceResource, options WorkspaceResourcesOptions) error {
+	// Sort resources by type for consistent output.
+	sort.Slice(resources, func(i, j int) bool {
+		return resources[i].Type < resources[j].Type
+	})
+
+	// Address on stop indexes whether a resource still exists when in the stopped transition.
+	addressOnStop := map[string]codersdk.WorkspaceResource{}
+	for _, resource := range resources {
+		if resource.Transition != database.WorkspaceTransitionStop {
+			continue
+		}
+		addressOnStop[resource.Address] = resource
+	}
+	// Displayed stores whether a resource has already been shown.
+	// Resources can be stored with numerous states, which we
+	// process prior to display.
+	displayed := map[string]struct{}{}
+
+	tableWriter := table.NewWriter()
+	if options.Title != "" {
+		tableWriter.SetTitle(options.Title)
+	}
+	tableWriter.SetStyle(table.StyleLight)
+	tableWriter.Style().Options.SeparateColumns = false
+	row := table.Row{"Resource", "Status"}
+	if !options.HideAccess {
+		row = append(row, "Access")
+	}
+	tableWriter.AppendHeader(row)
+
+	totalAgents := 0
+	for _, resource := range resources {
+		totalAgents += len(resource.Agents)
+	}
+
+	for _, resource := range resources {
+		if resource.Type == "random_string" {
+			// Hide resources that aren't substantial to a user!
+			// This is an unfortunate case, and we should allow
+			// callers to hide resources eventually.
+			continue
+		}
+		if _, shown := displayed[resource.Address]; shown {
+			// The same resource can have multiple transitions.
+			continue
+		}
+		displayed[resource.Address] = struct{}{}
+
+		// Sort agents by name for consistent output.
+		sort.Slice(resource.Agents, func(i, j int) bool {
+			return resource.Agents[i].Name < resource.Agents[j].Name
+		})
+		_, existsOnStop := addressOnStop[resource.Address]
+		resourceState := "ephemeral"
+		if existsOnStop {
+			resourceState = "persistent"
+		}
+		// Display a line for the resource.
+		tableWriter.AppendRow(table.Row{
+			Styles.Bold.Render(resource.Type + "." + resource.Name),
+			Styles.Placeholder.Render(resourceState),
+			"",
+		})
+		// Display all agents associated with the resource.
+		for index, agent := range resource.Agents {
+			sshCommand := "coder ssh " + options.WorkspaceName
+			if totalAgents > 1 {
+				sshCommand += "." + agent.Name
+			}
+			sshCommand = Styles.Code.Render(sshCommand)
+			var agentStatus string
+			if !options.HideAgentState {
+				switch agent.Status {
+				case codersdk.WorkspaceAgentConnecting:
+					since := database.Now().Sub(agent.CreatedAt)
+					agentStatus = Styles.Warn.Render("⦾ connecting") + " " +
+						Styles.Placeholder.Render("["+strconv.Itoa(int(since.Seconds()))+"s]")
+				case codersdk.WorkspaceAgentDisconnected:
+					since := database.Now().Sub(*agent.DisconnectedAt)
+					agentStatus = Styles.Error.Render("⦾ disconnected") + " " +
+						Styles.Placeholder.Render("["+strconv.Itoa(int(since.Seconds()))+"s]")
+				case codersdk.WorkspaceAgentConnected:
+					agentStatus = Styles.Keyword.Render("⦿ connected")
+				}
+			}
+
+			pipe := "├"
+			if index == len(resource.Agents)-1 {
+				pipe = "└"
+			}
+			row := table.Row{
+				// These tree from a resource!
+				fmt.Sprintf("%s─ %s (%s, %s)", pipe, agent.Name, agent.OperatingSystem, agent.Architecture),
+				agentStatus,
+			}
+			if !options.HideAccess {
+				row = append(row, sshCommand)
+			}
+			tableWriter.AppendRow(row)
+		}
+		tableWriter.AppendSeparator()
+	}
+	_, err := fmt.Fprintln(writer, tableWriter.Render())
+	return err
+}