tui/flamegraph: tighten arrow escape parsing (task 423)

2 files changed, 27 insertions, 4 deletions

diff --git a/internal/tui/flamegraph/model.go b/internal/tui/flamegraph/model.go
index 2238ba0..a7b26f8 100644
--- a/internal/tui/flamegraph/model.go
+++ b/internal/tui/flamegraph/model.go
@@ -968,14 +968,18 @@ func keyMatchesDirection(keyName, plain string, ansiFinal byte) bool {
 }
 
 func isArrowEscapeSequence(value string, ansiFinal byte) bool {
-	if len(value) < 3 || value[0] != '\x1b' {
+	body, ok := strings.CutPrefix(value, "\x1b")
+	if !ok || len(body) < 2 {
 		return false
 	}
-	last := value[len(value)-1]
-	if last != ansiFinal {
+	switch body[0] {
+	case '[':
+		return body[len(body)-1] == ansiFinal
+	case 'O':
+		return len(body) == 2 && body[1] == ansiFinal
+	default:
 		return false
 	}
-	return value[1] == '[' || value[1] == 'O'
 }
 
 func (m Model) visibleRowOffset() int {
diff --git a/internal/tui/flamegraph/model_test.go b/internal/tui/flamegraph/model_test.go
index 23c44e1..c2626cd 100644
--- a/internal/tui/flamegraph/model_test.go
+++ b/internal/tui/flamegraph/model_test.go
@@ -612,6 +612,25 @@ func TestArrowEscapeSequencesAreRecognized(t *testing.T) {
 	}
 }
 
+func TestArrowEscapeSequencesRejectMalformedValues(t *testing.T) {
+	tests := []struct {
+		name     string
+		key      string
+		ansiCode byte
+	}{
+		{name: "missing escape", key: "up", ansiCode: 'A'},
+		{name: "too short", key: "\x1b[", ansiCode: 'A'},
+		{name: "wrong final", key: "\x1b[A", ansiCode: 'B'},
+		{name: "unsupported introducer", key: "\x1bPA", ansiCode: 'A'},
+		{name: "application mode with extra payload", key: "\x1bO1A", ansiCode: 'A'},
+	}
+	for _, tc := range tests {
+		if isArrowEscapeSequence(tc.key, tc.ansiCode) {
+			t.Fatalf("expected %s sequence %q to be rejected", tc.name, tc.key)
+		}
+	}
+}
+
 func TestFilteredNavigationSkipsHiddenBranches(t *testing.T) {
 	m := NewModel(nil)
 	m.frames = []tuiFrame{