From 70def9aebc6506e69ba00b02e1a07c663315aeb3 Mon Sep 17 00:00:00 2001
From: kylecarbs <kyle@carberry.com>
Date: Fri, 27 May 2022 17:31:37 +0000
Subject: [PATCH] fix: Add command to reconnecting PTY

This fixes #1708 and opens the door for PTYs to execute
non-shell commands!
---
 agent/agent.go            | 6 +++---
 agent/agent_test.go       | 4 ++--
 agent/conn.go             | 6 ++++--
 coderd/workspaceagents.go | 2 +-
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 09a36e2e6d6a8..4859e35f08395 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -456,8 +456,8 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn ne
 
 	// The ID format is referenced in conn.go.
 	// <uuid>:<height>:<width>
-	idParts := strings.Split(rawID, ":")
-	if len(idParts) != 3 {
+	idParts := strings.SplitN(rawID, ":", 4)
+	if len(idParts) != 4 {
 		a.logger.Warn(ctx, "client sent invalid id format", slog.F("raw-id", rawID))
 		return
 	}
@@ -489,7 +489,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn ne
 		}
 	} else {
 		// Empty command will default to the users shell!
-		cmd, err := a.createCommand(ctx, "", nil)
+		cmd, err := a.createCommand(ctx, idParts[3], nil)
 		if err != nil {
 			a.logger.Warn(ctx, "create reconnecting pty command", slog.Error(err))
 			return
diff --git a/agent/agent_test.go b/agent/agent_test.go
index 6d2206097e8d2..923ce46290b5d 100644
--- a/agent/agent_test.go
+++ b/agent/agent_test.go
@@ -221,7 +221,7 @@ func TestAgent(t *testing.T) {
 
 		conn := setupAgent(t, agent.Metadata{}, 0)
 		id := uuid.NewString()
-		netConn, err := conn.ReconnectingPTY(id, 100, 100)
+		netConn, err := conn.ReconnectingPTY(id, 100, 100, "/bin/bash")
 		require.NoError(t, err)
 		bufRead := bufio.NewReader(netConn)
 
@@ -259,7 +259,7 @@ func TestAgent(t *testing.T) {
 		expectLine(matchEchoOutput)
 
 		_ = netConn.Close()
-		netConn, err = conn.ReconnectingPTY(id, 100, 100)
+		netConn, err = conn.ReconnectingPTY(id, 100, 100, "/bin/bash")
 		require.NoError(t, err)
 		bufRead = bufio.NewReader(netConn)
 
diff --git a/agent/conn.go b/agent/conn.go
index b63c0d0b0da35..d44d6d0c0b0d8 100644
--- a/agent/conn.go
+++ b/agent/conn.go
@@ -34,8 +34,10 @@ type Conn struct {
 
 // ReconnectingPTY returns a connection serving a TTY that can
 // be reconnected to via ID.
-func (c *Conn) ReconnectingPTY(id string, height, width uint16) (net.Conn, error) {
-	channel, err := c.CreateChannel(context.Background(), fmt.Sprintf("%s:%d:%d", id, height, width), &peer.ChannelOptions{
+//
+// The command is optional and defaults to start a shell.
+func (c *Conn) ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error) {
+	channel, err := c.CreateChannel(context.Background(), fmt.Sprintf("%s:%d:%d:%s", id, height, width, command), &peer.ChannelOptions{
 		Protocol: ProtocolReconnectingPTY,
 	})
 	if err != nil {
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
index bbda257eb01d5..cfcdea0404683 100644
--- a/coderd/workspaceagents.go
+++ b/coderd/workspaceagents.go
@@ -381,7 +381,7 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer agentConn.Close()
-	ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width))
+	ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width), "")
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))
 		return