feat: Add web terminal with reconnecting TTYs (#1186)

* feat: Add web terminal with reconnecting TTYs

This adds a web terminal that can reconnect to resume sessions!
No more disconnects, and no more bad bufferring!

* Add xstate service

* Add the webpage for accessing a web terminal

* Add terminal page tests

* Use Ticker instead of Timer

* Active Windows mode on Windows

agent/agent.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -12,10 +13,14 @@ import (
 	"os/exec"
 	"os/user"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/armon/circbuf"
+	"github.com/google/uuid"
+
 	gsyslog "github.com/hashicorp/go-syslog"
 	"go.uber.org/atomic"
 
@@ -33,6 +38,11 @@ import (
 	"golang.org/x/xerrors"
 )
 
+type Options struct {
+	ReconnectingPTYTimeout time.Duration
+	Logger                 slog.Logger
+}
+
 type Metadata struct {
 	OwnerEmail           string            `json:"owner_email"`
 	OwnerUsername        string            `json:"owner_username"`
@@ -42,13 +52,20 @@ type Metadata struct {
 
 type Dialer func(ctx context.Context, logger slog.Logger) (Metadata, *peerbroker.Listener, error)
 
-func New(dialer Dialer, logger slog.Logger) io.Closer {
+func New(dialer Dialer, options *Options) io.Closer {
+	if options == nil {
+		options = &Options{}
+	}
+	if options.ReconnectingPTYTimeout == 0 {
+		options.ReconnectingPTYTimeout = 5 * time.Minute
+	}
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	server := &agent{
-		dialer:      dialer,
-		logger:      logger,
-		closeCancel: cancelFunc,
-		closed:      make(chan struct{}),
+		dialer:                 dialer,
+		reconnectingPTYTimeout: options.ReconnectingPTYTimeout,
+		logger:                 options.Logger,
+		closeCancel:            cancelFunc,
+		closed:                 make(chan struct{}),
 	}
 	server.init(ctx)
 	return server
@@ -58,6 +75,9 @@ type agent struct {
 	dialer Dialer
 	logger slog.Logger
 
+	reconnectingPTYs       sync.Map
+	reconnectingPTYTimeout time.Duration
+
 	connCloseWait sync.WaitGroup
 	closeCancel   context.CancelFunc
 	closeMutex    sync.Mutex
@@ -196,6 +216,8 @@ func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 		switch channel.Protocol() {
 		case "ssh":
 			go a.sshServer.HandleConn(channel.NetConn())
+		case "reconnecting-pty":
+			go a.handleReconnectingPTY(ctx, channel.Label(), channel.NetConn())
 		default:
 			a.logger.Warn(ctx, "unhandled protocol from channel",
 				slog.F("protocol", channel.Protocol()),
@@ -282,22 +304,25 @@ func (a *agent) init(ctx context.Context) {
 	go a.run(ctx)
 }
 
-func (a *agent) handleSSHSession(session ssh.Session) error {
+// createCommand processes raw command input with OpenSSH-like behavior.
+// If the rawCommand provided is empty, it will default to the users shell.
+// This injects environment variables specified by the user at launch too.
+func (a *agent) createCommand(ctx context.Context, rawCommand string, env []string) (*exec.Cmd, error) {
 	currentUser, err := user.Current()
 	if err != nil {
-		return xerrors.Errorf("get current user: %w", err)
+		return nil, xerrors.Errorf("get current user: %w", err)
 	}
 	username := currentUser.Username
 
 	shell, err := usershell.Get(username)
 	if err != nil {
-		return xerrors.Errorf("get user shell: %w", err)
+		return nil, xerrors.Errorf("get user shell: %w", err)
 	}
 
 	// gliderlabs/ssh returns a command slice of zero
 	// when a shell is requested.
-	command := session.RawCommand()
-	if len(session.Command()) == 0 {
+	command := rawCommand
+	if len(command) == 0 {
 		command = shell
 	}
 
@@ -307,11 +332,11 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
 	if runtime.GOOS == "windows" {
 		caller = "/c"
 	}
-	cmd := exec.CommandContext(session.Context(), shell, caller, command)
-	cmd.Env = append(os.Environ(), session.Environ()...)
+	cmd := exec.CommandContext(ctx, shell, caller, command)
+	cmd.Env = append(os.Environ(), env...)
 	executablePath, err := os.Executable()
 	if err != nil {
-		return xerrors.Errorf("getting os executable: %w", err)
+		return nil, xerrors.Errorf("getting os executable: %w", err)
 	}
 	// Git on Windows resolves with UNIX-style paths.
 	// If using backslashes, it's unable to find the executable.
@@ -332,6 +357,14 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
 			}
 		}
 	}
+	return cmd, nil
+}
+
+func (a *agent) handleSSHSession(session ssh.Session) error {
+	cmd, err := a.createCommand(session.Context(), session.RawCommand(), session.Environ())
+	if err != nil {
+		return err
+	}
 
 	sshPty, windowSize, isPty := session.Pty()
 	if isPty {
@@ -381,6 +414,194 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
 	return cmd.Wait()
 }
 
+func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn net.Conn) {
+	defer conn.Close()
+
+	// The ID format is referenced in conn.go.
+	// <uuid>:<height>:<width>
+	idParts := strings.Split(rawID, ":")
+	if len(idParts) != 3 {
+		a.logger.Warn(ctx, "client sent invalid id format", slog.F("raw-id", rawID))
+		return
+	}
+	id := idParts[0]
+	// Enforce a consistent format for IDs.
+	_, err := uuid.Parse(id)
+	if err != nil {
+		a.logger.Warn(ctx, "client sent reconnection token that isn't a uuid", slog.F("id", id), slog.Error(err))
+		return
+	}
+	// Parse the initial terminal dimensions.
+	height, err := strconv.Atoi(idParts[1])
+	if err != nil {
+		a.logger.Warn(ctx, "client sent invalid height", slog.F("id", id), slog.F("height", idParts[1]))
+		return
+	}
+	width, err := strconv.Atoi(idParts[2])
+	if err != nil {
+		a.logger.Warn(ctx, "client sent invalid width", slog.F("id", id), slog.F("width", idParts[2]))
+		return
+	}
+
+	var rpty *reconnectingPTY
+	rawRPTY, ok := a.reconnectingPTYs.Load(id)
+	if ok {
+		rpty, ok = rawRPTY.(*reconnectingPTY)
+		if !ok {
+			a.logger.Warn(ctx, "found invalid type in reconnecting pty map", slog.F("id", id))
+		}
+	} else {
+		// Empty command will default to the users shell!
+		cmd, err := a.createCommand(ctx, "", nil)
+		if err != nil {
+			a.logger.Warn(ctx, "create reconnecting pty command", slog.Error(err))
+			return
+		}
+		cmd.Env = append(cmd.Env, "TERM=xterm-256color")
+
+		ptty, process, err := pty.Start(cmd)
+		if err != nil {
+			a.logger.Warn(ctx, "start reconnecting pty command", slog.F("id", id))
+		}
+
+		// Default to buffer 64KB.
+		circularBuffer, err := circbuf.NewBuffer(64 * 1024)
+		if err != nil {
+			a.logger.Warn(ctx, "create circular buffer", slog.Error(err))
+			return
+		}
+
+		a.closeMutex.Lock()
+		a.connCloseWait.Add(1)
+		a.closeMutex.Unlock()
+		ctx, cancelFunc := context.WithCancel(ctx)
+		rpty = &reconnectingPTY{
+			activeConns: make(map[string]net.Conn),
+			ptty:        ptty,
+			// Timeouts created with an after func can be reset!
+			timeout:        time.AfterFunc(a.reconnectingPTYTimeout, cancelFunc),
+			circularBuffer: circularBuffer,
+		}
+		a.reconnectingPTYs.Store(id, rpty)
+		go func() {
+			// CommandContext isn't respected for Windows PTYs right now,
+			// so we need to manually track the lifecycle.
+			// When the context has been completed either:
+			// 1. The timeout completed.
+			// 2. The parent context was canceled.
+			<-ctx.Done()
+			_ = process.Kill()
+		}()
+		go func() {
+			// If the process dies randomly, we should
+			// close the pty.
+			_, _ = process.Wait()
+			rpty.Close()
+		}()
+		go func() {
+			buffer := make([]byte, 1024)
+			for {
+				read, err := rpty.ptty.Output().Read(buffer)
+				if err != nil {
+					// When the PTY is closed, this is triggered.
+					break
+				}
+				part := buffer[:read]
+				_, err = rpty.circularBuffer.Write(part)
+				if err != nil {
+					a.logger.Error(ctx, "reconnecting pty write buffer", slog.Error(err), slog.F("id", id))
+					break
+				}
+				rpty.activeConnsMutex.Lock()
+				for _, conn := range rpty.activeConns {
+					_, _ = conn.Write(part)
+				}
+				rpty.activeConnsMutex.Unlock()
+			}
+
+			// Cleanup the process, PTY, and delete it's
+			// ID from memory.
+			_ = process.Kill()
+			rpty.Close()
+			a.reconnectingPTYs.Delete(id)
+			a.connCloseWait.Done()
+		}()
+	}
+	// Resize the PTY to initial height + width.
+	err = rpty.ptty.Resize(uint16(height), uint16(width))
+	if err != nil {
+		// We can continue after this, it's not fatal!
+		a.logger.Error(ctx, "resize reconnecting pty", slog.F("id", id), slog.Error(err))
+	}
+	// Write any previously stored data for the TTY.
+	_, err = conn.Write(rpty.circularBuffer.Bytes())
+	if err != nil {
+		a.logger.Warn(ctx, "write reconnecting pty buffer", slog.F("id", id), slog.Error(err))
+		return
+	}
+	connectionID := uuid.NewString()
+	// Multiple connections to the same TTY are permitted.
+	// This could easily be used for terminal sharing, but
+	// we do it because it's a nice user experience to
+	// copy/paste a terminal URL and have it _just work_.
+	rpty.activeConnsMutex.Lock()
+	rpty.activeConns[connectionID] = conn
+	rpty.activeConnsMutex.Unlock()
+	// Resetting this timeout prevents the PTY from exiting.
+	rpty.timeout.Reset(a.reconnectingPTYTimeout)
+
+	ctx, cancelFunc := context.WithCancel(ctx)
+	defer cancelFunc()
+	heartbeat := time.NewTicker(a.reconnectingPTYTimeout / 2)
+	defer heartbeat.Stop()
+	go func() {
+		// Keep updating the activity while this
+		// connection is alive!
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-heartbeat.C:
+			}
+			rpty.timeout.Reset(a.reconnectingPTYTimeout)
+		}
+	}()
+	defer func() {
+		// After this connection ends, remove it from
+		// the PTYs active connections. If it isn't
+		// removed, all PTY data will be sent to it.
+		rpty.activeConnsMutex.Lock()
+		delete(rpty.activeConns, connectionID)
+		rpty.activeConnsMutex.Unlock()
+	}()
+	decoder := json.NewDecoder(conn)
+	var req ReconnectingPTYRequest
+	for {
+		err = decoder.Decode(&req)
+		if xerrors.Is(err, io.EOF) {
+			return
+		}
+		if err != nil {
+			a.logger.Warn(ctx, "reconnecting pty buffer read error", slog.F("id", id), slog.Error(err))
+			return
+		}
+		_, err = rpty.ptty.Input().Write([]byte(req.Data))
+		if err != nil {
+			a.logger.Warn(ctx, "write to reconnecting pty", slog.F("id", id), slog.Error(err))
+			return
+		}
+		// Check if a resize needs to happen!
+		if req.Height == 0 || req.Width == 0 {
+			continue
+		}
+		err = rpty.ptty.Resize(req.Height, req.Width)
+		if err != nil {
+			// We can continue after this, it's not fatal!
+			a.logger.Error(ctx, "resize reconnecting pty", slog.F("id", id), slog.Error(err))
+		}
+	}
+}
+
 // isClosed returns whether the API is closed or not.
 func (a *agent) isClosed() bool {
 	select {
@@ -403,3 +624,25 @@ func (a *agent) Close() error {
 	a.connCloseWait.Wait()
 	return nil
 }
+
+type reconnectingPTY struct {
+	activeConnsMutex sync.Mutex
+	activeConns      map[string]net.Conn
+
+	circularBuffer *circbuf.Buffer
+	timeout        *time.Timer
+	ptty           pty.PTY
+}
+
+// Close ends all connections to the reconnecting
+// PTY and clear the circular buffer.
+func (r *reconnectingPTY) Close() {
+	r.activeConnsMutex.Lock()
+	defer r.activeConnsMutex.Unlock()
+	for _, conn := range r.activeConns {
+		_ = conn.Close()
+	}
+	_ = r.ptty.Close()
+	r.circularBuffer.Reset()
+	r.timeout.Stop()
+}

agent/agent_test.go

@@ -2,6 +2,7 @@ package agent_test
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net"
@@ -14,6 +15,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/pion/webrtc/v3"
 	"github.com/pkg/sftp"
 	"github.com/stretchr/testify/require"
@@ -131,7 +133,7 @@ func TestAgent(t *testing.T) {
 
 	t.Run("SFTP", func(t *testing.T) {
 		t.Parallel()
-		sshClient, err := setupAgent(t, agent.Metadata{}).SSHClient()
+		sshClient, err := setupAgent(t, agent.Metadata{}, 0).SSHClient()
 		require.NoError(t, err)
 		client, err := sftp.NewClient(sshClient)
 		require.NoError(t, err)
@@ -168,7 +170,7 @@ func TestAgent(t *testing.T) {
 		content := "somethingnice"
 		setupAgent(t, agent.Metadata{
 			StartupScript: "echo " + content + " > " + tempPath,
-		})
+		}, 0)
 		var gotContent string
 		require.Eventually(t, func() bool {
 			content, err := os.ReadFile(tempPath)
@@ -188,10 +190,54 @@ func TestAgent(t *testing.T) {
 		}, 15*time.Second, 100*time.Millisecond)
 		require.Equal(t, content, strings.TrimSpace(gotContent))
 	})
+
+	t.Run("ReconnectingPTY", func(t *testing.T) {
+		t.Parallel()
+		if runtime.GOOS == "windows" {
+			// This might be our implementation, or ConPTY itself.
+			// It's difficult to find extensive tests for it, so
+			// it seems like it could be either.
+			t.Skip("ConPTY appears to be inconsistent on Windows.")
+		}
+		conn := setupAgent(t, agent.Metadata{}, 0)
+		id := uuid.NewString()
+		netConn, err := conn.ReconnectingPTY(id, 100, 100)
+		require.NoError(t, err)
+
+		data, err := json.Marshal(agent.ReconnectingPTYRequest{
+			Data: "echo test\r\n",
+		})
+		require.NoError(t, err)
+		_, err = netConn.Write(data)
+		require.NoError(t, err)
+
+		findEcho := func() {
+			for {
+				read, err := netConn.Read(data)
+				require.NoError(t, err)
+				if strings.Contains(string(data[:read]), "test") {
+					break
+				}
+			}
+		}
+
+		// Once for typing the command...
+		findEcho()
+		// And another time for the actual output.
+		findEcho()
+
+		_ = netConn.Close()
+		netConn, err = conn.ReconnectingPTY(id, 100, 100)
+		require.NoError(t, err)
+
+		// Same output again!
+		findEcho()
+		findEcho()
+	})
 }
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
-	agentConn := setupAgent(t, agent.Metadata{})
+	agentConn := setupAgent(t, agent.Metadata{}, 0)
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 	go func() {
@@ -220,19 +266,22 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
 }
 
 func setupSSHSession(t *testing.T, options agent.Metadata) *ssh.Session {
-	sshClient, err := setupAgent(t, options).SSHClient()
+	sshClient, err := setupAgent(t, options, 0).SSHClient()
 	require.NoError(t, err)
 	session, err := sshClient.NewSession()
 	require.NoError(t, err)
 	return session
 }
 
-func setupAgent(t *testing.T, options agent.Metadata) *agent.Conn {
+func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) *agent.Conn {
 	client, server := provisionersdk.TransportPipe()
 	closer := agent.New(func(ctx context.Context, logger slog.Logger) (agent.Metadata, *peerbroker.Listener, error) {
 		listener, err := peerbroker.Listen(server, nil)
-		return options, listener, err
-	}, slogtest.Make(t, nil).Leveled(slog.LevelDebug))
+		return metadata, listener, err
+	}, &agent.Options{
+		Logger:                 slogtest.Make(t, nil).Leveled(slog.LevelDebug),
+		ReconnectingPTYTimeout: ptyTimeout,
+	})
 	t.Cleanup(func() {
 		_ = client.Close()
 		_ = server.Close()

agent/conn.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"context"
+	"fmt"
 	"net"
 
 	"golang.org/x/crypto/ssh"
@@ -11,6 +12,14 @@ import (
 	"github.com/coder/coder/peerbroker/proto"
 )
 
+// ReconnectingPTYRequest is sent from the client to the server
+// to pipe data to a PTY.
+type ReconnectingPTYRequest struct {
+	Data   string `json:"data"`
+	Height uint16 `json:"height"`
+	Width  uint16 `json:"width"`
+}
+
 // Conn wraps a peer connection with helper functions to
 // communicate with the agent.
 type Conn struct {
@@ -20,6 +29,18 @@ type Conn struct {
 	*peer.Conn
 }
 
+// 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.Dial(context.Background(), fmt.Sprintf("%s:%d:%d", id, height, width), &peer.ChannelOptions{
+		Protocol: "reconnecting-pty",
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("pty: %w", err)
+	}
+	return channel.NetConn(), nil
+}
+
 // SSH dials the built-in SSH server.
 func (c *Conn) SSH() (net.Conn, error) {
 	channel, err := c.Dial(context.Background(), "ssh", &peer.ChannelOptions{