fix: Add command to reconnecting PTY (#1860)

kylecarbs · web-flow · commit d0ed107b08f8 · 2022-05-27T14:51:20.000-05:00
This fixes #1708 and opens the door for PTYs to execute non-shell commands!
diff --git 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
@@ -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
@@ -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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -381,7 +381,7 @@ func (api API) workspaceAgentPTY(rw http.ResponseWriter, r http.Request) {`
`381`	`381`	`return`
`382`	`382`	`}`
`383`	`383`	`defer agentConn.Close()`
`384`		`- ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width))`
	`384`	`+ ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width), "")`
`385`	`385`	`if err != nil {`
`386`	`386`	`_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))`
`387`	`387`	`return`