Aug 14, 2023 · Jul 18, 2023 · Aug 8, 2023 · Aug 8, 2023 · Aug 8, 2023 · Aug 8, 2023
diff --git a/agent/agent.go b/agent/agent.go
 "sync"
 "time"

 "github.com/armon/circbuf"
 "github.com/go-chi/chi/v5"
 "github.com/google/uuid"
 "github.com/prometheus/client_golang/prometheus"

 "cdr.dev/slog"
 "github.com/coder/coder/agent/agentssh"
 "github.com/coder/coder/agent/reconnectingpty"
 "github.com/coder/coder/buildinfo"
 "github.com/coder/coder/coderd/database"
 "github.com/coder/coder/coderd/gitauth"
 "github.com/coder/coder/codersdk"
 "github.com/coder/coder/codersdk/agentsdk"
 "github.com/coder/coder/pty"
 "github.com/coder/coder/tailnet"
 "github.com/coder/retry"
 )
 }

 func New(options Options) Agent {
 if options.ReconnectingPTYTimeout == 0 {
 options.ReconnectingPTYTimeout = 5 * time.Minute
 }
 if options.Filesystem == nil {
 options.Filesystem = afero.NewOsFs()
 }
 defer a.connCountReconnectingPTY.Add(-1)

 connectionID := uuid.NewString()
 logger= logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
 logger.Debug(ctx, "starting handler")
 connLogger := logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
 connLogger.Debug(ctx, "starting handler")

 defer func() {
 if err := retErr; err != nil {
 // If the agent is closed, we don't want to
 // log this as an error since it's expected.
 if closed {
 logger.Debug(ctx, "reconnectingPTY failed withsession error (agent closed)", slog.Error(err))
 connLogger.Debug(ctx, "reconnectingpty failed withattach error (agent closed)", slog.Error(err))
 } else {
 logger.Error(ctx, "reconnectingPTY failed withsession error", slog.Error(err))
 connLogger.Error(ctx, "reconnectingpty failed withattach error", slog.Error(err))
 }
 }
 logger.Debug(ctx, "session closed")
 connLogger.Debug(ctx, "reconnecting pty connection closed")
 }()

 var rpty*reconnectingPTY
 sendConnected := make(chan*reconnectingPTY, 1)
 var rptyreconnectingpty.ReconnectingPTY
 sendConnected := make(chanreconnectingpty.ReconnectingPTY, 1)
 // On store, reserve this ID to prevent multiple concurrent new connections.
 waitReady, ok := a.reconnectingPTYs.LoadOrStore(msg.ID, sendConnected)
 if ok {
 close(sendConnected) // Unused.
 logger.Debug(ctx, "connecting to existingsession")
 c, ok := waitReady.(chan*reconnectingPTY)
 connLogger.Debug(ctx, "connecting to existingreconnecting pty")
 c, ok := waitReady.(chanreconnectingpty.ReconnectingPTY)
 if !ok {
 return xerrors.Errorf("found invalid type in reconnecting pty map: %T", waitReady)
 }
 }
 c <- rpty // Put it back for the next reconnect.
 } else {
 logger.Debug(ctx, "creating newsession")
 connLogger.Debug(ctx, "creating newreconnecting pty")

 connected := false
 defer func() {
 a.metrics.reconnectingPTYErrors.WithLabelValues("create_command").Add(1)
 return xerrors.Errorf("create command: %w", err)
 }
 cmd.Env = append(cmd.Env, "TERM=xterm-256color")

 // Default to buffer 64KiB.
 circularBuffer, err := circbuf.NewBuffer(64 << 10)
 if err != nil {
 return xerrors.Errorf("create circular buffer: %w", err)
 }

 ptty, process, err := pty.Start(cmd)
 if err != nil {
 a.metrics.reconnectingPTYErrors.WithLabelValues("start_command").Add(1)
 return xerrors.Errorf("start command: %w", err)
 }
 rpty = reconnectingpty.New(ctx, cmd, &reconnectingpty.Options{
 Timeout: a.reconnectingPTYTimeout,
 Metrics: a.metrics.reconnectingPTYErrors,
 }, logger.With(slog.F("message_id", msg.ID)))

 ctx, cancel := context.WithCancel(ctx)
 rpty = &reconnectingPTY{
 activeConns: map[string]net.Conn{
 // We have to put the connection in the map instantly otherwise
 // the connection won't be closed if the process instantly dies.
 connectionID: conn,
 },
 ptty: ptty,
 // Timeouts created with an after func can be reset!
 timeout:        time.AfterFunc(a.reconnectingPTYTimeout, cancel),
 circularBuffer: circularBuffer,
 }
 // We don't need to separately monitor for the process exiting.
 // When it exits, our ptty.OutputReader() will return EOF after
 // reading all process output.
 if err = a.trackConnGoroutine(func() {
 buffer := make([]byte, 1024)
 for {
 read, err := rpty.ptty.OutputReader().Read(buffer)
 if err != nil {
 // When the PTY is closed, this is triggered.
 // Error is typically a benign EOF, so only log for debugging.
 if errors.Is(err, io.EOF) {
 logger.Debug(ctx, "unable to read pty output, command might have exited", slog.Error(err))
 } else {
 logger.Warn(ctx, "unable to read pty output, command might have exited", slog.Error(err))
 a.metrics.reconnectingPTYErrors.WithLabelValues("output_reader").Add(1)
 }
 break
 }
 part := buffer[:read]
 rpty.circularBufferMutex.Lock()
 _, err = rpty.circularBuffer.Write(part)
 rpty.circularBufferMutex.Unlock()
 if err != nil {
 logger.Error(ctx, "write to circular buffer", slog.Error(err))
 break
 }
 rpty.activeConnsMutex.Lock()
 for cid, conn := range rpty.activeConns {
 _, err = conn.Write(part)
 if err != nil {
 logger.Warn(ctx,
 "error writing to active conn",
 slog.F("other_conn_id", cid),
 slog.Error(err),
 )
 a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
 }
 }
 rpty.activeConnsMutex.Unlock()
 }

 // Cleanup the process, PTY, and delete it's
 // ID from memory.
 _ = process.Kill()
 rpty.Close()
 rpty.Wait()
 a.reconnectingPTYs.Delete(msg.ID)
 }); err != nil {
 _ = process.Kill()
 _ = ptty.Close()
 rpty.Close(err.Error())
 return xerrors.Errorf("start routine: %w", err)
 }

 connected = true
 sendConnected <- rpty
 }
 // Resize the PTY to initial height + width.
 err := rpty.ptty.Resize(msg.Height, msg.Width)
 if err != nil {
 // We can continue after this, it's not fatal!
 logger.Error(ctx, "reconnecting PTY initial resize failed, but will continue", slog.Error(err))
 a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
 }
 // Write any previously stored data for the TTY.
 rpty.circularBufferMutex.RLock()
 prevBuf := slices.Clone(rpty.circularBuffer.Bytes())
 rpty.circularBufferMutex.RUnlock()
 // Note that there is a small race here between writing buffered
 // data and storing conn in activeConns. This is likely a very minor
 // edge case, but we should look into ways to avoid it. Holding
 // activeConnsMutex would be one option, but holding this mutex
 // while also holding circularBufferMutex seems dangerous.
 _, err = conn.Write(prevBuf)
 if err != nil {
 a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
 return xerrors.Errorf("write buffer to conn: %w", err)
 }
 // 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 codersdk.ReconnectingPTYRequest
 for {
 err = decoder.Decode(&req)
 if xerrors.Is(err, io.EOF) {
 return nil
 }
 if err != nil {
 logger.Warn(ctx, "reconnecting PTY failed with read error", slog.Error(err))
 return nil
 }
 _, err = rpty.ptty.InputWriter().Write([]byte(req.Data))
 if err != nil {
 logger.Warn(ctx, "reconnecting PTY failed with write error", slog.Error(err))
 a.metrics.reconnectingPTYErrors.WithLabelValues("input_writer").Add(1)
 return nil
 }
 // 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!
 logger.Error(ctx, "reconnecting PTY resize failed, but will continue", slog.Error(err))
 a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
 }
 }
 return rpty.Attach(ctx, connectionID, conn, msg.Height, msg.Width, connLogger)
 }

 // startReportingConnectionStats runs the connection stats reporting goroutine.
 return nil
 }

 type reconnectingPTY struct {
 activeConnsMutex sync.Mutex
 activeConns      map[string]net.Conn

 circularBuffer      *circbuf.Buffer
 circularBufferMutex sync.RWMutex
 timeout             *time.Timer
 ptty                pty.PTYCmd
 }

 // 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.circularBufferMutex.Lock()
 r.circularBuffer.Reset()
 r.circularBufferMutex.Unlock()
 r.timeout.Stop()
 }

 // userHomeDir returns the home directory of the current user, giving
 // priority to the $HOME environment variable.
 func userHomeDir() (string, error) {
Original file line number	Diff line number	Diff line change
Expand Up		@@ -21,7 +21,6 @@ import (
		"sync"
		"time"

		"github.com/armon/circbuf"
		"github.com/go-chi/chi/v5"
		"github.com/google/uuid"
		"github.com/prometheus/client_golang/prometheus"
Expand All		@@ -36,12 +35,12 @@ import (

		"cdr.dev/slog"
		"github.com/coder/coder/agent/agentssh"
		"github.com/coder/coder/agent/reconnectingpty"
		"github.com/coder/coder/buildinfo"
		"github.com/coder/coder/coderd/database"
		"github.com/coder/coder/coderd/gitauth"
		"github.com/coder/coder/codersdk"
		"github.com/coder/coder/codersdk/agentsdk"
		"github.com/coder/coder/pty"
		"github.com/coder/coder/tailnet"
		"github.com/coder/retry"
		)
Expand DownExpand Up		@@ -92,9 +91,6 @@ type Agent interface {
		}

		func New(options Options) Agent {
		if options.ReconnectingPTYTimeout == 0 {
		options.ReconnectingPTYTimeout = 5 * time.Minute
		}
		if options.Filesystem == nil {
		options.Filesystem = afero.NewOsFs()
		}
Expand DownExpand Up		@@ -1075,8 +1071,8 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
		defer a.connCountReconnectingPTY.Add(-1)

		connectionID := uuid.NewString()
		logger= logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
		logger.Debug(ctx, "starting handler")
		connLogger := logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
		connLogger.Debug(ctx, "starting handler")

		defer func() {
		if err := retErr; err != nil {
Expand All		@@ -1087,22 +1083,22 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
		// If the agent is closed, we don't want to
		// log this as an error since it's expected.
		if closed {
		logger.Debug(ctx, "reconnectingPTY failed withsession error (agent closed)", slog.Error(err))
		connLogger.Debug(ctx, "reconnectingpty failed withattach error (agent closed)", slog.Error(err))
		} else {
		logger.Error(ctx, "reconnectingPTY failed withsession error", slog.Error(err))
		connLogger.Error(ctx, "reconnectingpty failed withattach error", slog.Error(err))
		}
		}
		logger.Debug(ctx, "session closed")
		connLogger.Debug(ctx, "reconnecting pty connection closed")
		}()

		var rpty*reconnectingPTY
		sendConnected := make(chan*reconnectingPTY, 1)
		var rptyreconnectingpty.ReconnectingPTY
		sendConnected := make(chanreconnectingpty.ReconnectingPTY, 1)
		// On store, reserve this ID to prevent multiple concurrent new connections.
		waitReady, ok := a.reconnectingPTYs.LoadOrStore(msg.ID, sendConnected)
		if ok {
		close(sendConnected) // Unused.
		logger.Debug(ctx, "connecting to existingsession")
		c, ok := waitReady.(chan*reconnectingPTY)
		connLogger.Debug(ctx, "connecting to existingreconnecting pty")
		c, ok := waitReady.(chanreconnectingpty.ReconnectingPTY)
		if !ok {
		return xerrors.Errorf("found invalid type in reconnecting pty map: %T", waitReady)
		}
Expand All		@@ -1112,7 +1108,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
		}
		c <- rpty // Put it back for the next reconnect.
		} else {
		logger.Debug(ctx, "creating newsession")
		connLogger.Debug(ctx, "creating newreconnecting pty")

		connected := false
		defer func() {
Expand All		@@ -1128,169 +1124,24 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
		a.metrics.reconnectingPTYErrors.WithLabelValues("create_command").Add(1)
		return xerrors.Errorf("create command: %w", err)
		}
		cmd.Env = append(cmd.Env, "TERM=xterm-256color")

		// Default to buffer 64KiB.
		circularBuffer, err := circbuf.NewBuffer(64 << 10)
		if err != nil {
		return xerrors.Errorf("create circular buffer: %w", err)
		}

		ptty, process, err := pty.Start(cmd)
		if err != nil {
		a.metrics.reconnectingPTYErrors.WithLabelValues("start_command").Add(1)
		return xerrors.Errorf("start command: %w", err)
		}
		rpty = reconnectingpty.New(ctx, cmd, &reconnectingpty.Options{
		Timeout: a.reconnectingPTYTimeout,
		Metrics: a.metrics.reconnectingPTYErrors,
		}, logger.With(slog.F("message_id", msg.ID)))

		ctx, cancel := context.WithCancel(ctx)
		rpty = &reconnectingPTY{
		activeConns: map[string]net.Conn{
		// We have to put the connection in the map instantly otherwise
		// the connection won't be closed if the process instantly dies.
		connectionID: conn,
		},
		ptty: ptty,
		// Timeouts created with an after func can be reset!
		timeout: time.AfterFunc(a.reconnectingPTYTimeout, cancel),
		circularBuffer: circularBuffer,
		}
		// We don't need to separately monitor for the process exiting.
		// When it exits, our ptty.OutputReader() will return EOF after
		// reading all process output.
		if err = a.trackConnGoroutine(func() {
		buffer := make([]byte, 1024)
		for {
		read, err := rpty.ptty.OutputReader().Read(buffer)
		if err != nil {
		// When the PTY is closed, this is triggered.
		// Error is typically a benign EOF, so only log for debugging.
		if errors.Is(err, io.EOF) {
		logger.Debug(ctx, "unable to read pty output, command might have exited", slog.Error(err))
		} else {
		logger.Warn(ctx, "unable to read pty output, command might have exited", slog.Error(err))
		a.metrics.reconnectingPTYErrors.WithLabelValues("output_reader").Add(1)
		}
		break
		}
		part := buffer[:read]
		rpty.circularBufferMutex.Lock()
		_, err = rpty.circularBuffer.Write(part)
		rpty.circularBufferMutex.Unlock()
		if err != nil {
		logger.Error(ctx, "write to circular buffer", slog.Error(err))
		break
		}
		rpty.activeConnsMutex.Lock()
		for cid, conn := range rpty.activeConns {
		_, err = conn.Write(part)
		if err != nil {
		logger.Warn(ctx,
		"error writing to active conn",
		slog.F("other_conn_id", cid),
		slog.Error(err),
		)
		a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
		}
		}
		rpty.activeConnsMutex.Unlock()
		}

		// Cleanup the process, PTY, and delete it's
		// ID from memory.
		_ = process.Kill()
		rpty.Close()
		rpty.Wait()
		a.reconnectingPTYs.Delete(msg.ID)
		}); err != nil {
		_ = process.Kill()
		_ = ptty.Close()
		rpty.Close(err.Error())
		return xerrors.Errorf("start routine: %w", err)
		}

		connected = true
		sendConnected <- rpty
		}
		// Resize the PTY to initial height + width.
		err := rpty.ptty.Resize(msg.Height, msg.Width)
		if err != nil {
		// We can continue after this, it's not fatal!
		logger.Error(ctx, "reconnecting PTY initial resize failed, but will continue", slog.Error(err))
		a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
		}
		// Write any previously stored data for the TTY.
		rpty.circularBufferMutex.RLock()
		prevBuf := slices.Clone(rpty.circularBuffer.Bytes())
		rpty.circularBufferMutex.RUnlock()
		// Note that there is a small race here between writing buffered
		// data and storing conn in activeConns. This is likely a very minor
		// edge case, but we should look into ways to avoid it. Holding
		// activeConnsMutex would be one option, but holding this mutex
		// while also holding circularBufferMutex seems dangerous.
		_, err = conn.Write(prevBuf)
		if err != nil {
		a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
		return xerrors.Errorf("write buffer to conn: %w", err)
		}
		// 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 codersdk.ReconnectingPTYRequest
		for {
		err = decoder.Decode(&req)
		if xerrors.Is(err, io.EOF) {
		return nil
		}
		if err != nil {
		logger.Warn(ctx, "reconnecting PTY failed with read error", slog.Error(err))
		return nil
		}
		_, err = rpty.ptty.InputWriter().Write([]byte(req.Data))
		if err != nil {
		logger.Warn(ctx, "reconnecting PTY failed with write error", slog.Error(err))
		a.metrics.reconnectingPTYErrors.WithLabelValues("input_writer").Add(1)
		return nil
		}
		// 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!
		logger.Error(ctx, "reconnecting PTY resize failed, but will continue", slog.Error(err))
		a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
		}
		}
		return rpty.Attach(ctx, connectionID, conn, msg.Height, msg.Width, connLogger)
		}

		// startReportingConnectionStats runs the connection stats reporting goroutine.
Expand DownExpand Up		@@ -1541,31 +1392,6 @@ lifecycleWaitLoop:
		return nil
		}

		type reconnectingPTY struct {
		activeConnsMutex sync.Mutex
		activeConns map[string]net.Conn

		circularBuffer *circbuf.Buffer
		circularBufferMutex sync.RWMutex
		timeout *time.Timer
		ptty pty.PTYCmd
		}

		// 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.circularBufferMutex.Lock()
		r.circularBuffer.Reset()
		r.circularBufferMutex.Unlock()
		r.timeout.Stop()
		}

		// userHomeDir returns the home directory of the current user, giving
		// priority to the $HOME environment variable.
		func userHomeDir() (string, error) {
Expand Down