Filter terminal reply traffic around tmux attach and detach

internal/session/ssh.go

@@ -14,11 +14,14 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/asheshgoplani/agent-deck/internal/termreply"
 	"github.com/asheshgoplani/agent-deck/internal/tmux"
 	"github.com/creack/pty"
 	"golang.org/x/term"
 )
 
+const sshAttachReplyQuarantine = 2 * time.Second
+
 // sshControlDir is the directory for SSH ControlMaster sockets.
 const sshControlDir = "/tmp/agent-deck-ssh"
 
@@ -206,6 +209,7 @@ func (r *SSHRunner) Attach(sessionID string) error {
 	case <-outputDone:
 	case <-time.After(50 * time.Millisecond):
 	}
+	termreply.QuarantineFor(sshAttachReplyQuarantine)
 
 	// Reset terminal styles that may have leaked from the remote session.
 	_, _ = os.Stdout.WriteString("\x1b]8;;\x1b\\\x1b[0m\x1b[24m\x1b[39m\x1b[49m")

internal/termreply/filter.go

@@ -0,0 +1,211 @@
+package termreply
+
+const (
+	escapeByte = 0x1b
+	bellByte   = 0x07
+
+	controlSequenceIntroducerByte = '['
+	singleShiftThreeByte          = 'O'
+
+	operatingSystemCommandByte    = ']'
+	deviceControlStringByte       = 'P'
+	applicationProgramCommandByte = '_'
+	privacyMessageByte            = '^'
+	startOfStringByte             = 'X'
+
+	stringTerminatorByte = '\\'
+
+	csiFinalArrowUpByte    = 'A'
+	csiFinalArrowDownByte  = 'B'
+	csiFinalArrowRightByte = 'C'
+	csiFinalArrowLeftByte  = 'D'
+	csiFinalEndByte        = 'F'
+	csiFinalHomeByte       = 'H'
+	csiFinalBacktabByte    = 'Z'
+	csiFinalTildeByte      = '~'
+	csiFinalKittyKeyByte   = 'u'
+)
+
+type filterMode uint8
+
+const (
+	filterModeIdle filterMode = iota
+	filterModeDiscardEscapeString
+	filterModeCollectCSI
+	filterModeCollectSS3
+)
+
+// Filter strips terminal-generated control replies from a byte stream while
+// preserving ordinary keyboard input. It is stateful so replies split across
+// reads are discarded without relying on terminal-specific payload strings.
+type Filter struct {
+	mode                filterMode
+	pendingEsc          bool
+	escapeSeenInDiscard bool
+	sequenceBuf         []byte
+}
+
+// Active reports whether the filter is carrying parser state across read boundaries.
+func (f *Filter) Active() bool {
+	return f.pendingEsc || f.mode != filterModeIdle || len(f.sequenceBuf) > 0
+}
+
+func isEscapeStringIntroducer(b byte) bool {
+	switch b {
+	case operatingSystemCommandByte,
+		deviceControlStringByte,
+		applicationProgramCommandByte,
+		privacyMessageByte,
+		startOfStringByte:
+		return true
+	default:
+		return false
+	}
+}
+
+func isSequenceFinalByte(b byte) bool {
+	return b >= 0x40 && b <= 0x7e
+}
+
+func isKeyboardCSIFinalByte(b byte) bool {
+	switch b {
+	case csiFinalArrowUpByte,
+		csiFinalArrowDownByte,
+		csiFinalArrowRightByte,
+		csiFinalArrowLeftByte,
+		csiFinalEndByte,
+		csiFinalHomeByte,
+		csiFinalBacktabByte,
+		csiFinalTildeByte,
+		csiFinalKittyKeyByte:
+		return true
+	default:
+		return false
+	}
+}
+
+func flushSequence(out []byte, seq []byte) []byte {
+	return append(out, seq...)
+}
+
+func (f *Filter) beginSequence(mode filterMode, prefix ...byte) {
+	f.mode = mode
+	f.sequenceBuf = append(f.sequenceBuf[:0], prefix...)
+}
+
+func (f *Filter) resetSequenceState() {
+	f.mode = filterModeIdle
+	f.sequenceBuf = f.sequenceBuf[:0]
+	f.escapeSeenInDiscard = false
+}
+
+// Consume filters a chunk of bytes. When armed is true, terminal-generated
+// control replies are discarded. If a reply started in a previous chunk, it
+// continues to be discarded until it terminates even if armed is now false.
+//
+// Terminal replies covered here:
+//   - escape-string families: OSC, DCS, APC, PM, SOS
+//   - CSI replies during the quarantine window, except for a small whitelist of
+//     keyboard-related CSI finals (arrows/home/end/backtab/~ keys/kitty CSI u)
+//
+// If final is true, any incomplete pending escape/CSI/SS3 sequence is flushed as
+// literal input, while an incomplete discarded escape-string reply is dropped.
+func (f *Filter) Consume(src []byte, armed bool, final bool) []byte {
+	out := make([]byte, 0, len(src))
+
+	for _, b := range src {
+		switch f.mode {
+		case filterModeDiscardEscapeString:
+			if f.escapeSeenInDiscard {
+				f.escapeSeenInDiscard = false
+				if b == stringTerminatorByte {
+					f.resetSequenceState()
+					continue
+				}
+				if b == escapeByte {
+					f.escapeSeenInDiscard = true
+				}
+				continue
+			}
+
+			if b == bellByte {
+				f.resetSequenceState()
+				continue
+			}
+			if b == escapeByte {
+				f.escapeSeenInDiscard = true
+			}
+			continue
+
+		case filterModeCollectCSI:
+			f.sequenceBuf = append(f.sequenceBuf, b)
+			if !isSequenceFinalByte(b) {
+				continue
+			}
+
+			if armed && !isKeyboardCSIFinalByte(b) {
+				f.resetSequenceState()
+				continue
+			}
+
+			out = flushSequence(out, f.sequenceBuf)
+			f.resetSequenceState()
+			continue
+
+		case filterModeCollectSS3:
+			f.sequenceBuf = append(f.sequenceBuf, b)
+			if !isSequenceFinalByte(b) {
+				continue
+			}
+
+			out = flushSequence(out, f.sequenceBuf)
+			f.resetSequenceState()
+			continue
+		}
+
+		if f.pendingEsc {
+			f.pendingEsc = false
+			switch {
+			case armed && isEscapeStringIntroducer(b):
+				f.mode = filterModeDiscardEscapeString
+				continue
+			case b == controlSequenceIntroducerByte:
+				f.beginSequence(filterModeCollectCSI, escapeByte, controlSequenceIntroducerByte)
+				continue
+			case b == singleShiftThreeByte:
+				f.beginSequence(filterModeCollectSS3, escapeByte, singleShiftThreeByte)
+				continue
+			case b == escapeByte:
+				out = append(out, escapeByte)
+				f.pendingEsc = true
+				continue
+			default:
+				out = append(out, escapeByte, b)
+				continue
+			}
+		}
+
+		if b == escapeByte {
+			f.pendingEsc = true
+			continue
+		}
+
+		out = append(out, b)
+	}
+
+	if final {
+		if f.pendingEsc {
+			out = append(out, escapeByte)
+		}
+		switch f.mode {
+		case filterModeCollectCSI, filterModeCollectSS3:
+			out = flushSequence(out, f.sequenceBuf)
+		case filterModeDiscardEscapeString:
+			// Drop incomplete escape-string replies on EOF.
+		}
+		f.pendingEsc = false
+		f.resetSequenceState()
+	}
+
+	return out
+}

internal/termreply/filter_test.go

@@ -0,0 +1,37 @@
+package termreply
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFilterDiscardsStringRepliesAcrossChunks(t *testing.T) {
+	var f Filter
+
+	got := f.Consume([]byte("\x1b]11;rgb:d3d3/f5f5/f5f5"), true, false)
+	require.Empty(t, got)
+	require.True(t, f.Active())
+
+	got = f.Consume([]byte("\x07j"), true, false)
+	require.Equal(t, []byte("j"), got)
+	require.False(t, f.Active())
+}
+
+func TestFilterDiscardsGenericCSIReplies(t *testing.T) {
+	var f Filter
+
+	got := f.Consume([]byte("\x1b[?1;2c"), true, false)
+	require.Empty(t, got)
+	require.False(t, f.Active())
+}
+
+func TestFilterPreservesKeyboardCSIAndSS3Input(t *testing.T) {
+	var f Filter
+
+	require.Equal(t, []byte("\x1b[A"), f.Consume([]byte("\x1b[A"), true, false))
+	require.False(t, f.Active())
+
+	require.Equal(t, []byte("\x1bOA"), f.Consume([]byte("\x1bOA"), true, false))
+	require.False(t, f.Active())
+}

internal/termreply/guard.go

@@ -0,0 +1,36 @@
+package termreply
+
+import (
+	"sync/atomic"
+	"time"
+)
+
+var quarantineUntilUnixNano atomic.Int64
+
+// QuarantineFor drops terminal reply traffic until the later of the existing
+// deadline or now+duration.
+func QuarantineFor(duration time.Duration) {
+	if duration <= 0 {
+		return
+	}
+	target := time.Now().Add(duration).UnixNano()
+	for {
+		current := quarantineUntilUnixNano.Load()
+		if current >= target {
+			return
+		}
+		if quarantineUntilUnixNano.CompareAndSwap(current, target) {
+			return
+		}
+	}
+}
+
+// Active reports whether terminal replies should currently be discarded.
+func Active() bool {
+	return time.Now().UnixNano() < quarantineUntilUnixNano.Load()
+}
+
+// Clear removes any active quarantine window. Intended for tests.
+func Clear() {
+	quarantineUntilUnixNano.Store(0)
+}