Extend session with active SSH connection

cmd/breakpoint/wait.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"time"
 
 	"github.com/dustin/go-humanize"
 	"github.com/muesli/reflow/wordwrap"
@@ -75,6 +76,10 @@ func newWaitCmd() *cobra.Command {
 
 				_, _ = w.Write(ww.Bytes())
 			},
+			WriteNotify: func() {
+				// FIXME: mgr.ExtendWait sends down `m.updated` channel, which announces. Introduce new method?
+				mgr.ExtendWait(5 * time.Minute)
+			},
 		})
 		if err != nil {
 			return err

pkg/sshd/sshd.go

@@ -19,6 +19,8 @@ import (
 	"golang.org/x/exp/slices"
 )
 
+const notifyDelay = 1 * time.Minute
+
 type SSHServerOpts struct {
 	AllowedUsers   []string
 	AuthorizedKeys map[string]string // Key to owner
@@ -27,6 +29,7 @@ type SSHServerOpts struct {
 	Dir            string
 
 	InteractiveMOTD func(io.Writer)
+	WriteNotify     func()
 }
 
 type sshKey struct {
@@ -94,23 +97,26 @@ func MakeServer(ctx context.Context, opts SSHServerOpts) (*SSHServer, error) {
 			// Make sure that the connection with the client is kept alive.
 			go keepAlive(ctx, sessionLog, session)
 
+			// Wrapping the session lets us know when writes are happening.
+			nsess := newNotifyingSession(ctx, session, opts.WriteNotify)
+
 			if isPty {
 				// Print MOTD only if no command was provided
-				if opts.InteractiveMOTD != nil && session.RawCommand() == "" {
-					opts.InteractiveMOTD(session)
+				if opts.InteractiveMOTD != nil && nsess.RawCommand() == "" {
+					opts.InteractiveMOTD(nsess)
 				}
 
-				if err := handlePty(session, ptyReq, winCh, cmd); err != nil {
+				if err := handlePty(nsess, ptyReq, winCh, cmd); err != nil {
 					sessionLog.Err(err).Msg("pty start failed")
-					session.Exit(1)
+					nsess.Exit(1)
 					return
 				}
 			} else {
-				cmd.Stdout = session
-				cmd.Stderr = session
+				cmd.Stdout = nsess
+				cmd.Stderr = nsess
 				if err := cmd.Start(); err != nil {
 					sessionLog.Err(err).Msg("start failed")
-					session.Exit(1)
+					nsess.Exit(1)
 					return
 				}
 			}
@@ -182,3 +188,49 @@ func lookupKey(allowed []sshKey, key ssh.PublicKey) (sshKey, bool) {
 	}
 	return sshKey{}, false
 }
+
+type notifyingSession struct {
+	ssh.Session
+	notifyCh chan struct{}
+	notify   func()
+}
+
+func newNotifyingSession(ctx context.Context, s ssh.Session, notify func()) ssh.Session {
+	if notify == nil {
+		return s
+	}
+
+	sess := notifyingSession{
+		Session:  s,
+		notifyCh: make(chan struct{}),
+		notify:   notify,
+	}
+	go sess.listen(ctx)
+	return sess
+}
+
+func (s notifyingSession) listen(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-s.notifyCh:
+		}
+
+		s.notify()
+
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(notifyDelay):
+		}
+	}
+}
+
+func (s notifyingSession) Write(p []byte) (int, error) {
+	select {
+	case s.notifyCh <- struct{}{}:
+	default: // avoid blocking
+	}
+	return s.Session.Write(p)
+}