@@ -30,6 +30,7 @@ import (
3030 "github.com/spf13/afero"
3131 "go.uber.org/atomic"
3232 gossh "golang.org/x/crypto/ssh"
33+ "golang.org/x/exp/slices"
3334 "golang.org/x/xerrors"
3435 "tailscale.com/net/speedtest"
3536 "tailscale.com/tailcfg"
@@ -90,7 +91,7 @@ func New(options Options) io.Closer {
9091 }
9192 }
9293 ctx , cancelFunc := context .WithCancel (context .Background ())
93- server := & agent {
94+ a := & agent {
9495 reconnectingPTYTimeout : options .ReconnectingPTYTimeout ,
9596 logger : options .Logger ,
9697 closeCancel : cancelFunc ,
@@ -101,8 +102,8 @@ func New(options Options) io.Closer {
101102 filesystem : options .Filesystem ,
102103 tempDir : options .TempDir ,
103104 }
104- server .init (ctx )
105- return server
105+ a .init (ctx )
106+ return a
106107}
107108
108109type agent struct {
@@ -215,8 +216,16 @@ func (a *agent) run(ctx context.Context) error {
215216 return xerrors .Errorf ("create tailnet: %w" , err )
216217 }
217218 a .closeMutex .Lock ()
218- a .network = network
219+ // Re-check if agent was closed while initializing the network.
220+ closed := a .isClosed ()
221+ if ! closed {
222+ a .network = network
223+ }
219224 a .closeMutex .Unlock ()
225+ if closed {
226+ _ = network .Close ()
227+ return xerrors .New ("agent is closed" )
228+ }
220229 } else {
221230 // Update the DERP map!
222231 network .SetDERPMap (metadata .DERPMap )
@@ -246,27 +255,20 @@ func (a *agent) trackConnGoroutine(fn func()) error {
246255}
247256
248257func (a * agent ) createTailnet (ctx context.Context , derpMap * tailcfg.DERPMap ) (_ * tailnet.Conn , err error ) {
249- a .closeMutex .Lock ()
250- if a .isClosed () {
251- a .closeMutex .Unlock ()
252- return nil , xerrors .New ("closed" )
253- }
254258 network , err := tailnet .NewConn (& tailnet.Options {
255259 Addresses : []netip.Prefix {netip .PrefixFrom (codersdk .TailnetIP , 128 )},
256260 DERPMap : derpMap ,
257261 Logger : a .logger .Named ("tailnet" ),
258262 EnableTrafficStats : true ,
259263 })
260264 if err != nil {
261- a .closeMutex .Unlock ()
262265 return nil , xerrors .Errorf ("create tailnet: %w" , err )
263266 }
264267 defer func () {
265268 if err != nil {
266269 network .Close ()
267270 }
268271 }()
269- a .closeMutex .Unlock ()
270272
271273 sshListener , err := network .Listen ("tcp" , ":" + strconv .Itoa (codersdk .TailnetSSHPort ))
272274 if err != nil {
@@ -299,10 +301,12 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
299301 }
300302 }()
301303 if err = a .trackConnGoroutine (func () {
304+ logger := a .logger .Named ("reconnecting-pty" )
305+
302306 for {
303307 conn , err := reconnectingPTYListener .Accept ()
304308 if err != nil {
305- a . logger .Debug (ctx , "accept pty failed" , slog .Error (err ))
309+ logger .Debug (ctx , "accept pty failed" , slog .Error (err ))
306310 return
307311 }
308312 // This cannot use a JSON decoder, since that can
@@ -323,7 +327,9 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
323327 if err != nil {
324328 continue
325329 }
326- go a .handleReconnectingPTY (ctx , msg , conn )
330+ go func () {
331+ _ = a .handleReconnectingPTY (ctx , logger , msg , conn )
332+ }()
327333 }
328334 }); err != nil {
329335 return nil , err
@@ -797,38 +803,56 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
797803 return cmd .Wait ()
798804}
799805
800- func (a * agent ) handleReconnectingPTY (ctx context.Context , msg codersdk.ReconnectingPTYInit , conn net.Conn ) {
806+ func (a * agent ) handleReconnectingPTY (ctx context.Context , logger slog. Logger , msg codersdk.ReconnectingPTYInit , conn net.Conn ) ( retErr error ) {
801807 defer conn .Close ()
802808
803809 connectionID := uuid .NewString ()
810+ logger = logger .With (slog .F ("id" , msg .ID ), slog .F ("connection_id" , connectionID ))
811+
812+ defer func () {
813+ if err := retErr ; err != nil {
814+ a .closeMutex .Lock ()
815+ closed := a .isClosed ()
816+ a .closeMutex .Unlock ()
817+
818+ // If the agent is closed, we don't want to
819+ // log this as an error since it's expected.
820+ if closed {
821+ logger .Debug (ctx , "session error after agent close" , slog .Error (err ))
822+ } else {
823+ logger .Error (ctx , "session error" , slog .Error (err ))
824+ }
825+ }
826+ logger .Debug (ctx , "session closed" )
827+ }()
828+
804829 var rpty * reconnectingPTY
805830 rawRPTY , ok := a .reconnectingPTYs .Load (msg .ID )
806831 if ok {
832+ logger .Debug (ctx , "connecting to existing session" )
807833 rpty , ok = rawRPTY .(* reconnectingPTY )
808834 if ! ok {
809- a .logger .Error (ctx , "found invalid type in reconnecting pty map" , slog .F ("id" , msg .ID ))
810- return
835+ return xerrors .Errorf ("found invalid type in reconnecting pty map: %T" , rawRPTY )
811836 }
812837 } else {
838+ logger .Debug (ctx , "creating new session" )
839+
813840 // Empty command will default to the users shell!
814841 cmd , err := a .createCommand (ctx , msg .Command , nil )
815842 if err != nil {
816- a .logger .Error (ctx , "create reconnecting pty command" , slog .Error (err ))
817- return
843+ return xerrors .Errorf ("create command: %w" , err )
818844 }
819845 cmd .Env = append (cmd .Env , "TERM=xterm-256color" )
820846
821847 // Default to buffer 64KiB.
822848 circularBuffer , err := circbuf .NewBuffer (64 << 10 )
823849 if err != nil {
824- a .logger .Error (ctx , "create circular buffer" , slog .Error (err ))
825- return
850+ return xerrors .Errorf ("create circular buffer: %w" , err )
826851 }
827852
828853 ptty , process , err := pty .Start (cmd )
829854 if err != nil {
830- a .logger .Error (ctx , "start reconnecting pty command" , slog .F ("id" , msg .ID ), slog .Error (err ))
831- return
855+ return xerrors .Errorf ("start command: %w" , err )
832856 }
833857
834858 ctx , cancelFunc := context .WithCancel (ctx )
@@ -872,7 +896,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
872896 _ , err = rpty .circularBuffer .Write (part )
873897 rpty .circularBufferMutex .Unlock ()
874898 if err != nil {
875- a . logger .Error (ctx , "reconnecting pty write buffer" , slog .Error (err ), slog . F ( "id" , msg . ID ))
899+ logger .Error (ctx , "write to circular buffer" , slog .Error (err ))
876900 break
877901 }
878902 rpty .activeConnsMutex .Lock ()
@@ -888,23 +912,27 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
888912 rpty .Close ()
889913 a .reconnectingPTYs .Delete (msg .ID )
890914 }); err != nil {
891- a .logger .Error (ctx , "start reconnecting pty routine" , slog .F ("id" , msg .ID ), slog .Error (err ))
892- return
915+ return xerrors .Errorf ("start routine: %w" , err )
893916 }
894917 }
895918 // Resize the PTY to initial height + width.
896919 err := rpty .ptty .Resize (msg .Height , msg .Width )
897920 if err != nil {
898921 // We can continue after this, it's not fatal!
899- a . logger .Error (ctx , "resize reconnecting pty" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
922+ logger .Error (ctx , "resize" , slog .Error (err ))
900923 }
901924 // Write any previously stored data for the TTY.
902925 rpty .circularBufferMutex .RLock ()
903- _ , err = conn . Write (rpty .circularBuffer .Bytes ())
926+ prevBuf := slices . Clone (rpty .circularBuffer .Bytes ())
904927 rpty .circularBufferMutex .RUnlock ()
928+ // Note that there is a small race here between writing buffered
929+ // data and storing conn in activeConns. This is likely a very minor
930+ // edge case, but we should look into ways to avoid it. Holding
931+ // activeConnsMutex would be one option, but holding this mutex
932+ // while also holding circularBufferMutex seems dangerous.
933+ _ , err = conn .Write (prevBuf )
905934 if err != nil {
906- a .logger .Warn (ctx , "write reconnecting pty buffer" , slog .F ("id" , msg .ID ), slog .Error (err ))
907- return
935+ return xerrors .Errorf ("write buffer to conn: %w" , err )
908936 }
909937 // Multiple connections to the same TTY are permitted.
910938 // This could easily be used for terminal sharing, but
@@ -945,16 +973,16 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
945973 for {
946974 err = decoder .Decode (& req )
947975 if xerrors .Is (err , io .EOF ) {
948- return
976+ return nil
949977 }
950978 if err != nil {
951- a . logger .Warn (ctx , "reconnecting pty buffer read error" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
952- return
979+ logger .Warn (ctx , "read conn" , slog .Error (err ))
980+ return nil
953981 }
954982 _ , err = rpty .ptty .Input ().Write ([]byte (req .Data ))
955983 if err != nil {
956- a . logger .Warn (ctx , "write to reconnecting pty" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
957- return
984+ logger .Warn (ctx , "write to pty" , slog .Error (err ))
985+ return nil
958986 }
959987 // Check if a resize needs to happen!
960988 if req .Height == 0 || req .Width == 0 {
@@ -963,7 +991,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
963991 err = rpty .ptty .Resize (req .Height , req .Width )
964992 if err != nil {
965993 // We can continue after this, it's not fatal!
966- a . logger .Error (ctx , "resize reconnecting pty" , slog . F ( "id" , msg . ID ) , slog .Error (err ))
994+ logger .Error (ctx , "resize" , slog .Error (err ))
967995 }
968996 }
969997}
0 commit comments