fix(snapshot): address follow-up review findings
- commit
- f4b9a613b57b3f8728f6b3a3aece0cfab483070a
- author
- brice · · 2026-07-06 11:08 UTC
- parent
- e2b324b
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
5 files · +140 -89
M internal/platform/proc_linux.go +10 -2
| @@ -8,15 +8,23 @@ import ( | |||
| 8 | 8 | ||
| 9 | 9 | // ProcessCwd returns the current working directory for a Linux process. | |
| 10 | 10 | func ProcessCwd(pid int) (string, error) { | |
| 11 | - | return os.Readlink(fmt.Sprintf("/proc/%d/cwd", pid)) | |
| 11 | + | return processCwd("/proc", pid) | |
| 12 | + | } | |
| 13 | + | ||
| 14 | + | func processCwd(root string, pid int) (string, error) { | |
| 15 | + | return os.Readlink(fmt.Sprintf("%s/%d/cwd", root, pid)) | |
| 12 | 16 | } | |
| 13 | 17 | ||
| 14 | 18 | // ProcessComm returns the command name from /proc/<pid>/comm. | |
| 15 | 19 | func ProcessComm(pid int) (string, error) { | |
| 20 | + | return processComm("/proc", pid) | |
| 21 | + | } | |
| 22 | + | ||
| 23 | + | func processComm(root string, pid int) (string, error) { | |
| 16 | 24 | if pid <= 0 { | |
| 17 | 25 | return "", fmt.Errorf("process comm: invalid pid %d", pid) | |
| 18 | 26 | } | |
| 19 | - | data, err := os.ReadFile(fmt.Sprintf("/proc/%d/comm", pid)) | |
| 27 | + | data, err := os.ReadFile(fmt.Sprintf("%s/%d/comm", root, pid)) | |
| 20 | 28 | if err != nil { | |
| 21 | 29 | return "", fmt.Errorf("process comm: read pid %d: %w", pid, err) | |
| 22 | 30 | } | |
| @@ -8,15 +8,23 @@ import ( | |||||
| 8 | 8 | ||||
| 9 | // ProcessCwd returns the current working directory for a Linux process. | 9 | // ProcessCwd returns the current working directory for a Linux process. | ||
| 10 | func ProcessCwd(pid int) (string, error) { | 10 | func ProcessCwd(pid int) (string, error) { | ||
| 11 | - | return os.Readlink(fmt.Sprintf("/proc/%d/cwd", pid)) | 11 | + | return processCwd("/proc", pid) |
| 12 | + | } | |||
| 13 | + | ||||
| 14 | + | func processCwd(root string, pid int) (string, error) { | |||
| 15 | + | return os.Readlink(fmt.Sprintf("%s/%d/cwd", root, pid)) | |||
| 12 | } | 16 | } | ||
| 13 | 17 | ||||
| 14 | // ProcessComm returns the command name from /proc/<pid>/comm. | 18 | // ProcessComm returns the command name from /proc/<pid>/comm. | ||
| 15 | func ProcessComm(pid int) (string, error) { | 19 | func ProcessComm(pid int) (string, error) { | ||
| 20 | + | return processComm("/proc", pid) | |||
| 21 | + | } | |||
| 22 | + | ||||
| 23 | + | func processComm(root string, pid int) (string, error) { | |||
| 16 | if pid <= 0 { | 24 | if pid <= 0 { | ||
| 17 | return "", fmt.Errorf("process comm: invalid pid %d", pid) | 25 | return "", fmt.Errorf("process comm: invalid pid %d", pid) | ||
| 18 | } | 26 | } | ||
| 19 | - | data, err := os.ReadFile(fmt.Sprintf("/proc/%d/comm", pid)) | 27 | + | data, err := os.ReadFile(fmt.Sprintf("%s/%d/comm", root, pid)) |
| 20 | if err != nil { | 28 | if err != nil { | ||
| 21 | return "", fmt.Errorf("process comm: read pid %d: %w", pid, err) | 29 | return "", fmt.Errorf("process comm: read pid %d: %w", pid, err) | ||
| 22 | } | 30 | } | ||
M internal/platform/process_linux.go +10 -4
| @@ -28,9 +28,15 @@ func NewProcessInspector() *ProcessInspector { return newProcessInspector("/proc | |||
| 28 | 28 | ||
| 29 | 29 | func newProcessInspector(root string) *ProcessInspector { return &ProcessInspector{root: root} } | |
| 30 | 30 | ||
| 31 | - | func (*ProcessInspector) Cwd(pid int) (string, error) { return ProcessCwd(pid) } | |
| 32 | - | func (*ProcessInspector) Comm(pid int) (string, error) { return ProcessComm(pid) } | |
| 33 | - | func (*ProcessInspector) Argv(pid int) ([]string, error) { return ProcessArgv(pid) } | |
| 31 | + | func (p *ProcessInspector) Cwd(pid int) (string, error) { | |
| 32 | + | return processCwd(p.root, pid) | |
| 33 | + | } | |
| 34 | + | func (p *ProcessInspector) Comm(pid int) (string, error) { | |
| 35 | + | return processComm(p.root, pid) | |
| 36 | + | } | |
| 37 | + | func (p *ProcessInspector) Argv(pid int) ([]string, error) { | |
| 38 | + | return processArgv(filepath.Join(p.root, strconv.Itoa(pid), "cmdline"), pid) | |
| 39 | + | } | |
| 34 | 40 | func (p *ProcessInspector) GroupArgv(pgid int, shellPid int) ([]string, error) { | |
| 35 | 41 | recs, err := p.processRecords() | |
| 36 | 42 | if err != nil { | |
| @@ -40,7 +46,7 @@ func (p *ProcessInspector) GroupArgv(pgid int, shellPid int) ([]string, error) { | |||
| 40 | 46 | if !ok { | |
| 41 | 47 | return nil, nil | |
| 42 | 48 | } | |
| 43 | - | return ProcessArgv(pid) | |
| 49 | + | return p.Argv(pid) | |
| 44 | 50 | } | |
| 45 | 51 | ||
| 46 | 52 | func (p *ProcessInspector) processRecords() ([]processRecord, error) { | |
| @@ -28,9 +28,15 @@ func NewProcessInspector() *ProcessInspector { return newProcessInspector("/proc | |||||
| 28 | 28 | ||||
| 29 | func newProcessInspector(root string) *ProcessInspector { return &ProcessInspector{root: root} } | 29 | func newProcessInspector(root string) *ProcessInspector { return &ProcessInspector{root: root} } | ||
| 30 | 30 | ||||
| 31 | - | func (*ProcessInspector) Cwd(pid int) (string, error) { return ProcessCwd(pid) } | 31 | + | func (p *ProcessInspector) Cwd(pid int) (string, error) { |
| 32 | - | func (*ProcessInspector) Comm(pid int) (string, error) { return ProcessComm(pid) } | 32 | + | return processCwd(p.root, pid) |
| 33 | - | func (*ProcessInspector) Argv(pid int) ([]string, error) { return ProcessArgv(pid) } | 33 | + | } |
| 34 | + | func (p *ProcessInspector) Comm(pid int) (string, error) { | |||
| 35 | + | return processComm(p.root, pid) | |||
| 36 | + | } | |||
| 37 | + | func (p *ProcessInspector) Argv(pid int) ([]string, error) { | |||
| 38 | + | return processArgv(filepath.Join(p.root, strconv.Itoa(pid), "cmdline"), pid) | |||
| 39 | + | } | |||
| 34 | func (p *ProcessInspector) GroupArgv(pgid int, shellPid int) ([]string, error) { | 40 | func (p *ProcessInspector) GroupArgv(pgid int, shellPid int) ([]string, error) { | ||
| 35 | recs, err := p.processRecords() | 41 | recs, err := p.processRecords() | ||
| 36 | if err != nil { | 42 | if err != nil { | ||
| @@ -40,7 +46,7 @@ func (p *ProcessInspector) GroupArgv(pgid int, shellPid int) ([]string, error) { | |||||
| 40 | if !ok { | 46 | if !ok { | ||
| 41 | return nil, nil | 47 | return nil, nil | ||
| 42 | } | 48 | } | ||
| 43 | - | return ProcessArgv(pid) | 49 | + | return p.Argv(pid) |
| 44 | } | 50 | } | ||
| 45 | 51 | ||||
| 46 | func (p *ProcessInspector) processRecords() ([]processRecord, error) { | 52 | func (p *ProcessInspector) processRecords() ([]processRecord, error) { | ||
M internal/platform/process_linux_test.go +30 -0
| @@ -66,6 +66,36 @@ func TestSelectProcessGroupPID(t *testing.T) { | |||
| 66 | 66 | } | |
| 67 | 67 | } | |
| 68 | 68 | ||
| 69 | + | func TestProcessInspectorUsesConfiguredRoot(t *testing.T) { | |
| 70 | + | root := t.TempDir() | |
| 71 | + | writeProcStat(t, root, 10, 7) | |
| 72 | + | pidDir := filepath.Join(root, "10") | |
| 73 | + | cwdTarget := t.TempDir() | |
| 74 | + | if err := os.Symlink(cwdTarget, filepath.Join(pidDir, "cwd")); err != nil { | |
| 75 | + | t.Fatal(err) | |
| 76 | + | } | |
| 77 | + | if err := os.WriteFile(filepath.Join(pidDir, "comm"), []byte("custom\n"), 0o600); err != nil { | |
| 78 | + | t.Fatal(err) | |
| 79 | + | } | |
| 80 | + | if err := os.WriteFile(filepath.Join(pidDir, "cmdline"), []byte("custom\x00arg\x00"), 0o600); err != nil { | |
| 81 | + | t.Fatal(err) | |
| 82 | + | } | |
| 83 | + | ins := newProcessInspector(root) | |
| 84 | + | ||
| 85 | + | cwd, err := ins.Cwd(10) | |
| 86 | + | if err != nil || cwd != cwdTarget { | |
| 87 | + | t.Fatalf("Cwd = %q, %v; want %q, nil", cwd, err, cwdTarget) | |
| 88 | + | } | |
| 89 | + | comm, err := ins.Comm(10) | |
| 90 | + | if err != nil || comm != "custom" { | |
| 91 | + | t.Fatalf("Comm = %q, %v; want custom, nil", comm, err) | |
| 92 | + | } | |
| 93 | + | argv, err := ins.Argv(10) | |
| 94 | + | if err != nil || !reflect.DeepEqual(argv, []string{"custom", "arg"}) { | |
| 95 | + | t.Fatalf("Argv = %#v, %v; want custom argv, nil", argv, err) | |
| 96 | + | } | |
| 97 | + | } | |
| 98 | + | ||
| 69 | 99 | func TestProcessInspectorCachesProcessRecordsBriefly(t *testing.T) { | |
| 70 | 100 | root := t.TempDir() | |
| 71 | 101 | writeProcStat(t, root, 10, 7) | |
| @@ -66,6 +66,36 @@ func TestSelectProcessGroupPID(t *testing.T) { | |||||
| 66 | } | 66 | } | ||
| 67 | } | 67 | } | ||
| 68 | 68 | ||||
| 69 | + | func TestProcessInspectorUsesConfiguredRoot(t *testing.T) { | |||
| 70 | + | root := t.TempDir() | |||
| 71 | + | writeProcStat(t, root, 10, 7) | |||
| 72 | + | pidDir := filepath.Join(root, "10") | |||
| 73 | + | cwdTarget := t.TempDir() | |||
| 74 | + | if err := os.Symlink(cwdTarget, filepath.Join(pidDir, "cwd")); err != nil { | |||
| 75 | + | t.Fatal(err) | |||
| 76 | + | } | |||
| 77 | + | if err := os.WriteFile(filepath.Join(pidDir, "comm"), []byte("custom\n"), 0o600); err != nil { | |||
| 78 | + | t.Fatal(err) | |||
| 79 | + | } | |||
| 80 | + | if err := os.WriteFile(filepath.Join(pidDir, "cmdline"), []byte("custom\x00arg\x00"), 0o600); err != nil { | |||
| 81 | + | t.Fatal(err) | |||
| 82 | + | } | |||
| 83 | + | ins := newProcessInspector(root) | |||
| 84 | + | ||||
| 85 | + | cwd, err := ins.Cwd(10) | |||
| 86 | + | if err != nil || cwd != cwdTarget { | |||
| 87 | + | t.Fatalf("Cwd = %q, %v; want %q, nil", cwd, err, cwdTarget) | |||
| 88 | + | } | |||
| 89 | + | comm, err := ins.Comm(10) | |||
| 90 | + | if err != nil || comm != "custom" { | |||
| 91 | + | t.Fatalf("Comm = %q, %v; want custom, nil", comm, err) | |||
| 92 | + | } | |||
| 93 | + | argv, err := ins.Argv(10) | |||
| 94 | + | if err != nil || !reflect.DeepEqual(argv, []string{"custom", "arg"}) { | |||
| 95 | + | t.Fatalf("Argv = %#v, %v; want custom argv, nil", argv, err) | |||
| 96 | + | } | |||
| 97 | + | } | |||
| 98 | + | ||||
| 69 | func TestProcessInspectorCachesProcessRecordsBriefly(t *testing.T) { | 99 | func TestProcessInspectorCachesProcessRecordsBriefly(t *testing.T) { | ||
| 70 | root := t.TempDir() | 100 | root := t.TempDir() | ||
| 71 | writeProcStat(t, root, 10, 7) | 101 | writeProcStat(t, root, 10, 7) | ||
M internal/usecase/daemon/identity.go +5 -1
| @@ -6,8 +6,12 @@ import ( | |||
| 6 | 6 | "fmt" | |
| 7 | 7 | "io" | |
| 8 | 8 | "strings" | |
| 9 | + | "sync/atomic" | |
| 10 | + | "time" | |
| 9 | 11 | ) | |
| 10 | 12 | ||
| 13 | + | var fallbackStableIDCounter atomic.Uint64 | |
| 14 | + | ||
| 11 | 15 | func newStableID(prefix string) (string, error) { | |
| 12 | 16 | var b [16]byte | |
| 13 | 17 | if _, err := io.ReadFull(rand.Reader, b[:]); err != nil { | |
| @@ -32,7 +36,7 @@ func (d *Daemon) newTabPaneStableIDs() (string, string, error) { | |||
| 32 | 36 | func fallbackStableID(prefix string) string { | |
| 33 | 37 | id, err := newStableID(prefix) | |
| 34 | 38 | if err != nil { | |
| 35 | - | return prefix + "_unknown" | |
| 39 | + | return fmt.Sprintf("%s_fallback_%x_%x", prefix, time.Now().UnixNano(), fallbackStableIDCounter.Add(1)) | |
| 36 | 40 | } | |
| 37 | 41 | return id | |
| 38 | 42 | } | |
| @@ -6,8 +6,12 @@ import ( | |||||
| 6 | "fmt" | 6 | "fmt" | ||
| 7 | "io" | 7 | "io" | ||
| 8 | "strings" | 8 | "strings" | ||
| 9 | + | "sync/atomic" | |||
| 10 | + | "time" | |||
| 9 | ) | 11 | ) | ||
| 10 | 12 | ||||
| 13 | + | var fallbackStableIDCounter atomic.Uint64 | |||
| 14 | + | ||||
| 11 | func newStableID(prefix string) (string, error) { | 15 | func newStableID(prefix string) (string, error) { | ||
| 12 | var b [16]byte | 16 | var b [16]byte | ||
| 13 | if _, err := io.ReadFull(rand.Reader, b[:]); err != nil { | 17 | if _, err := io.ReadFull(rand.Reader, b[:]); err != nil { | ||
| @@ -32,7 +36,7 @@ func (d *Daemon) newTabPaneStableIDs() (string, string, error) { | |||||
| 32 | func fallbackStableID(prefix string) string { | 36 | func fallbackStableID(prefix string) string { | ||
| 33 | id, err := newStableID(prefix) | 37 | id, err := newStableID(prefix) | ||
| 34 | if err != nil { | 38 | if err != nil { | ||
| 35 | - | return prefix + "_unknown" | 39 | + | return fmt.Sprintf("%s_fallback_%x_%x", prefix, time.Now().UnixNano(), fallbackStableIDCounter.Add(1)) |
| 36 | } | 40 | } | ||
| 37 | return id | 41 | return id | ||
| 38 | } | 42 | } | ||
M internal/usecase/daemon/snapshot_restore_test.go +85 -82
| @@ -111,97 +111,100 @@ func processRestoreTestStore(t *testing.T, name string, proc *snapcodec.Process) | |||
| 111 | 111 | })}}} | |
| 112 | 112 | } | |
| 113 | 113 | ||
| 114 | - | func TestRestoreSnapshotsWritesAllowlistedProcessCommands(t *testing.T) { | |
| 115 | - | store := processRestoreTestStore(t, "proc", &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric}) | |
| 116 | - | factory := &restorePTYFactory{} | |
| 117 | - | d := newTestDaemon(t, factory, stubClock{}) | |
| 118 | - | WithSnapshotStore(store)(d) | |
| 119 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"less"}}}) | |
| 120 | - | ||
| 121 | - | d.restoreSnapshots(context.Background()) | |
| 122 | - | ||
| 123 | - | require.Len(t, factory.opens, 1) | |
| 124 | - | require.Equal(t, []string{"less README.md\n"}, factory.opens[0].pty.writes) | |
| 125 | - | } | |
| 126 | - | ||
| 127 | - | func TestRestoreSnapshotsSkipsDeniedProcessCommands(t *testing.T) { | |
| 128 | - | store := processRestoreTestStore(t, "proc", &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric}) | |
| 129 | - | factory := &restorePTYFactory{} | |
| 130 | - | d := newTestDaemon(t, factory, stubClock{}) | |
| 131 | - | WithSnapshotStore(store)(d) | |
| 132 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"vim"}}}) | |
| 133 | - | ||
| 134 | - | d.restoreSnapshots(context.Background()) | |
| 114 | + | func TestRestoreSnapshotsProcessRestoreCommands(t *testing.T) { | |
| 115 | + | lessProc := &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric} | |
| 116 | + | agentProc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc123"}} | |
| 117 | + | controlAgentProc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc\nmalicious"}} | |
| 135 | 118 | ||
| 136 | - | require.Len(t, factory.opens, 1) | |
| 137 | - | require.Empty(t, factory.opens[0].pty.writes) | |
| 138 | - | } | |
| 139 | - | ||
| 140 | - | func TestRestoreSnapshotsKeepsProcessCommandsScopedToPane(t *testing.T) { | |
| 141 | - | store := &restoreSnapshotStore{blobs: []ports.SnapshotBlob{{Name: "scoped", Data: mustSnapshotBytes(t, snapcodec.Session{ | |
| 142 | - | Name: "scoped", | |
| 143 | - | Tabs: []snapcodec.Tab{ | |
| 144 | - | { | |
| 145 | - | Cols: 80, | |
| 146 | - | Rows: 24, | |
| 147 | - | Tree: layout.NewTree("pane-1"), | |
| 148 | - | Panes: []snapcodec.Pane{{ | |
| 149 | - | ID: "pane-1", | |
| 150 | - | Cwd: "/one", | |
| 151 | - | Process: &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric}, | |
| 152 | - | }}, | |
| 153 | - | }, | |
| 154 | - | { | |
| 155 | - | Cols: 80, | |
| 156 | - | Rows: 24, | |
| 157 | - | Tree: layout.NewTree("pane-1"), | |
| 158 | - | Panes: []snapcodec.Pane{{ | |
| 159 | - | ID: "pane-1", | |
| 160 | - | Cwd: "/two", | |
| 161 | - | }}, | |
| 162 | - | }, | |
| 119 | + | tests := []struct { | |
| 120 | + | name string | |
| 121 | + | store *restoreSnapshotStore | |
| 122 | + | allow []string | |
| 123 | + | wantWrites [][]string | |
| 124 | + | }{ | |
| 125 | + | { | |
| 126 | + | name: "allowlisted generic command", | |
| 127 | + | store: processRestoreTestStore(t, "allowlisted", lessProc), | |
| 128 | + | allow: []string{"less"}, | |
| 129 | + | wantWrites: [][]string{{"less README.md\n"}}, | |
| 163 | 130 | }, | |
| 164 | - | })}}} | |
| 165 | - | factory := &restorePTYFactory{} | |
| 166 | - | d := newTestDaemon(t, factory, stubClock{}) | |
| 167 | - | WithSnapshotStore(store)(d) | |
| 168 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"less"}}}) | |
| 169 | - | ||
| 170 | - | d.restoreSnapshots(context.Background()) | |
| 131 | + | { | |
| 132 | + | name: "denied generic command", | |
| 133 | + | store: processRestoreTestStore(t, "denied", lessProc), | |
| 134 | + | allow: []string{"vim"}, | |
| 135 | + | wantWrites: [][]string{{}}, | |
| 136 | + | }, | |
| 137 | + | { | |
| 138 | + | name: "pane scoped command", | |
| 139 | + | store: &restoreSnapshotStore{blobs: []ports.SnapshotBlob{{Name: "scoped", Data: mustSnapshotBytes(t, snapcodec.Session{ | |
| 140 | + | Name: "scoped", | |
| 141 | + | Tabs: []snapcodec.Tab{ | |
| 142 | + | { | |
| 143 | + | Cols: 80, | |
| 144 | + | Rows: 24, | |
| 145 | + | Tree: layout.NewTree("pane-1"), | |
| 146 | + | Panes: []snapcodec.Pane{{ | |
| 147 | + | ID: "pane-1", | |
| 148 | + | Cwd: "/one", | |
| 149 | + | Process: lessProc, | |
| 150 | + | }}, | |
| 151 | + | }, | |
| 152 | + | { | |
| 153 | + | Cols: 80, | |
| 154 | + | Rows: 24, | |
| 155 | + | Tree: layout.NewTree("pane-1"), | |
| 156 | + | Panes: []snapcodec.Pane{{ | |
| 157 | + | ID: "pane-1", | |
| 158 | + | Cwd: "/two", | |
| 159 | + | }}, | |
| 160 | + | }, | |
| 161 | + | }, | |
| 162 | + | })}}}, | |
| 163 | + | allow: []string{"less"}, | |
| 164 | + | wantWrites: [][]string{{"less README.md\n"}, {}}, | |
| 165 | + | }, | |
| 166 | + | { | |
| 167 | + | name: "agent ID command", | |
| 168 | + | store: processRestoreTestStore(t, "agent", agentProc), | |
| 169 | + | allow: []string{"pi"}, | |
| 170 | + | wantWrites: [][]string{{"pi --resume abc123\n"}}, | |
| 171 | + | }, | |
| 172 | + | { | |
| 173 | + | name: "agent ID with control byte rejected", | |
| 174 | + | store: processRestoreTestStore(t, "agent-control", controlAgentProc), | |
| 175 | + | allow: []string{"pi"}, | |
| 176 | + | wantWrites: [][]string{{}}, | |
| 177 | + | }, | |
| 178 | + | } | |
| 171 | 179 | ||
| 172 | - | require.Len(t, factory.opens, 2) | |
| 173 | - | require.Equal(t, []string{"less README.md\n"}, factory.opens[0].pty.writes) | |
| 174 | - | require.Empty(t, factory.opens[1].pty.writes) | |
| 180 | + | for _, tt := range tests { | |
| 181 | + | t.Run(tt.name, func(t *testing.T) { | |
| 182 | + | factory := &restorePTYFactory{} | |
| 183 | + | d := newTestDaemon(t, factory, stubClock{}) | |
| 184 | + | WithSnapshotStore(tt.store)(d) | |
| 185 | + | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: tt.allow}}) | |
| 186 | + | ||
| 187 | + | d.restoreSnapshots(context.Background()) | |
| 188 | + | ||
| 189 | + | require.Len(t, factory.opens, len(tt.wantWrites)) | |
| 190 | + | for i, want := range tt.wantWrites { | |
| 191 | + | if len(want) == 0 { | |
| 192 | + | require.Empty(t, factory.opens[i].pty.writes) | |
| 193 | + | continue | |
| 194 | + | } | |
| 195 | + | require.Equal(t, want, factory.opens[i].pty.writes) | |
| 196 | + | } | |
| 197 | + | }) | |
| 198 | + | } | |
| 175 | 199 | } | |
| 176 | 200 | ||
| 177 | - | func TestRestoreSnapshotsWritesAgentIDCommandAndWriteFailureIsNonFatal(t *testing.T) { | |
| 201 | + | func TestRestoreSnapshotsProcessCommandWriteFailureIsNonFatal(t *testing.T) { | |
| 178 | 202 | proc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc123"}} | |
| 179 | - | store := processRestoreTestStore(t, "agent", proc) | |
| 180 | - | factory := &restorePTYFactory{} | |
| 203 | + | factory := &restorePTYFactory{writeErr: errors.New("boom")} | |
| 181 | 204 | d := newTestDaemon(t, factory, stubClock{}) | |
| 182 | - | WithSnapshotStore(store)(d) | |
| 183 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"pi"}}}) | |
| 184 | - | ||
| 185 | - | d.restoreSnapshots(context.Background()) | |
| 186 | - | ||
| 187 | - | require.Len(t, factory.opens, 1) | |
| 188 | - | require.Equal(t, []string{"pi --resume abc123\n"}, factory.opens[0].pty.writes) | |
| 189 | - | factory.writeErr = errors.New("boom") | |
| 190 | - | require.NoError(t, d.restoreSession(context.Background(), snapcodec.Session{Name: "agent2", Tabs: []snapcodec.Tab{{Cols: 80, Rows: 24, Tree: layout.NewTree("pane-1"), Panes: []snapcodec.Pane{{ID: "pane-1", Cwd: "/tmp", Process: proc}}}}})) | |
| 191 | - | } | |
| 192 | - | ||
| 193 | - | func TestRestoreSnapshotsRejectsAgentIDWithControlBytes(t *testing.T) { | |
| 194 | - | proc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc\nmalicious"}} | |
| 195 | - | store := processRestoreTestStore(t, "agent-control", proc) | |
| 196 | - | factory := &restorePTYFactory{} | |
| 197 | - | d := newTestDaemon(t, factory, stubClock{}) | |
| 198 | - | WithSnapshotStore(store)(d) | |
| 199 | 205 | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"pi"}}}) | |
| 200 | 206 | ||
| 201 | - | d.restoreSnapshots(context.Background()) | |
| 202 | - | ||
| 203 | - | require.Len(t, factory.opens, 1) | |
| 204 | - | require.Empty(t, factory.opens[0].pty.writes) | |
| 207 | + | require.NoError(t, d.restoreSession(context.Background(), snapcodec.Session{Name: "agent", Tabs: []snapcodec.Tab{{Cols: 80, Rows: 24, Tree: layout.NewTree("pane-1"), Panes: []snapcodec.Pane{{ID: "pane-1", Cwd: "/tmp", Process: proc}}}}})) | |
| 205 | 208 | } | |
| 206 | 209 | ||
| 207 | 210 | func TestRestoreSnapshotsOpensCollapsedStackPanesWithValidPTYSize(t *testing.T) { | |
| @@ -111,97 +111,100 @@ func processRestoreTestStore(t *testing.T, name string, proc *snapcodec.Process) | |||||
| 111 | })}}} | 111 | })}}} | ||
| 112 | } | 112 | } | ||
| 113 | 113 | ||||
| 114 | - | func TestRestoreSnapshotsWritesAllowlistedProcessCommands(t *testing.T) { | 114 | + | func TestRestoreSnapshotsProcessRestoreCommands(t *testing.T) { |
| 115 | - | store := processRestoreTestStore(t, "proc", &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric}) | 115 | + | lessProc := &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric} |
| 116 | - | factory := &restorePTYFactory{} | 116 | + | agentProc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc123"}} |
| 117 | - | d := newTestDaemon(t, factory, stubClock{}) | 117 | + | controlAgentProc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc\nmalicious"}} |
| 118 | - | WithSnapshotStore(store)(d) | |||
| 119 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"less"}}}) | |||
| 120 | - | ||||
| 121 | - | d.restoreSnapshots(context.Background()) | |||
| 122 | - | ||||
| 123 | - | require.Len(t, factory.opens, 1) | |||
| 124 | - | require.Equal(t, []string{"less README.md\n"}, factory.opens[0].pty.writes) | |||
| 125 | - | } | |||
| 126 | - | ||||
| 127 | - | func TestRestoreSnapshotsSkipsDeniedProcessCommands(t *testing.T) { | |||
| 128 | - | store := processRestoreTestStore(t, "proc", &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric}) | |||
| 129 | - | factory := &restorePTYFactory{} | |||
| 130 | - | d := newTestDaemon(t, factory, stubClock{}) | |||
| 131 | - | WithSnapshotStore(store)(d) | |||
| 132 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"vim"}}}) | |||
| 133 | - | ||||
| 134 | - | d.restoreSnapshots(context.Background()) | |||
| 135 | 118 | ||||
| 136 | - | require.Len(t, factory.opens, 1) | 119 | + | tests := []struct { |
| 137 | - | require.Empty(t, factory.opens[0].pty.writes) | 120 | + | name string |
| 138 | - | } | 121 | + | store *restoreSnapshotStore |
| 139 | - | 122 | + | allow []string | |
| 140 | - | func TestRestoreSnapshotsKeepsProcessCommandsScopedToPane(t *testing.T) { | 123 | + | wantWrites [][]string |
| 141 | - | store := &restoreSnapshotStore{blobs: []ports.SnapshotBlob{{Name: "scoped", Data: mustSnapshotBytes(t, snapcodec.Session{ | 124 | + | }{ |
| 142 | - | Name: "scoped", | 125 | + | { |
| 143 | - | Tabs: []snapcodec.Tab{ | 126 | + | name: "allowlisted generic command", |
| 144 | - | { | 127 | + | store: processRestoreTestStore(t, "allowlisted", lessProc), |
| 145 | - | Cols: 80, | 128 | + | allow: []string{"less"}, |
| 146 | - | Rows: 24, | 129 | + | wantWrites: [][]string{{"less README.md\n"}}, |
| 147 | - | Tree: layout.NewTree("pane-1"), | |||
| 148 | - | Panes: []snapcodec.Pane{{ | |||
| 149 | - | ID: "pane-1", | |||
| 150 | - | Cwd: "/one", | |||
| 151 | - | Process: &snapcodec.Process{Argv: []string{"less", "README.md"}, Strategy: processStrategyGeneric}, | |||
| 152 | - | }}, | |||
| 153 | - | }, | |||
| 154 | - | { | |||
| 155 | - | Cols: 80, | |||
| 156 | - | Rows: 24, | |||
| 157 | - | Tree: layout.NewTree("pane-1"), | |||
| 158 | - | Panes: []snapcodec.Pane{{ | |||
| 159 | - | ID: "pane-1", | |||
| 160 | - | Cwd: "/two", | |||
| 161 | - | }}, | |||
| 162 | - | }, | |||
| 163 | }, | 130 | }, | ||
| 164 | - | })}}} | 131 | + | { |
| 165 | - | factory := &restorePTYFactory{} | 132 | + | name: "denied generic command", |
| 166 | - | d := newTestDaemon(t, factory, stubClock{}) | 133 | + | store: processRestoreTestStore(t, "denied", lessProc), |
| 167 | - | WithSnapshotStore(store)(d) | 134 | + | allow: []string{"vim"}, |
| 168 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"less"}}}) | 135 | + | wantWrites: [][]string{{}}, |
| 169 | - | 136 | + | }, | |
| 170 | - | d.restoreSnapshots(context.Background()) | 137 | + | { |
| 138 | + | name: "pane scoped command", | |||
| 139 | + | store: &restoreSnapshotStore{blobs: []ports.SnapshotBlob{{Name: "scoped", Data: mustSnapshotBytes(t, snapcodec.Session{ | |||
| 140 | + | Name: "scoped", | |||
| 141 | + | Tabs: []snapcodec.Tab{ | |||
| 142 | + | { | |||
| 143 | + | Cols: 80, | |||
| 144 | + | Rows: 24, | |||
| 145 | + | Tree: layout.NewTree("pane-1"), | |||
| 146 | + | Panes: []snapcodec.Pane{{ | |||
| 147 | + | ID: "pane-1", | |||
| 148 | + | Cwd: "/one", | |||
| 149 | + | Process: lessProc, | |||
| 150 | + | }}, | |||
| 151 | + | }, | |||
| 152 | + | { | |||
| 153 | + | Cols: 80, | |||
| 154 | + | Rows: 24, | |||
| 155 | + | Tree: layout.NewTree("pane-1"), | |||
| 156 | + | Panes: []snapcodec.Pane{{ | |||
| 157 | + | ID: "pane-1", | |||
| 158 | + | Cwd: "/two", | |||
| 159 | + | }}, | |||
| 160 | + | }, | |||
| 161 | + | }, | |||
| 162 | + | })}}}, | |||
| 163 | + | allow: []string{"less"}, | |||
| 164 | + | wantWrites: [][]string{{"less README.md\n"}, {}}, | |||
| 165 | + | }, | |||
| 166 | + | { | |||
| 167 | + | name: "agent ID command", | |||
| 168 | + | store: processRestoreTestStore(t, "agent", agentProc), | |||
| 169 | + | allow: []string{"pi"}, | |||
| 170 | + | wantWrites: [][]string{{"pi --resume abc123\n"}}, | |||
| 171 | + | }, | |||
| 172 | + | { | |||
| 173 | + | name: "agent ID with control byte rejected", | |||
| 174 | + | store: processRestoreTestStore(t, "agent-control", controlAgentProc), | |||
| 175 | + | allow: []string{"pi"}, | |||
| 176 | + | wantWrites: [][]string{{}}, | |||
| 177 | + | }, | |||
| 178 | + | } | |||
| 171 | 179 | ||||
| 172 | - | require.Len(t, factory.opens, 2) | 180 | + | for _, tt := range tests { |
| 173 | - | require.Equal(t, []string{"less README.md\n"}, factory.opens[0].pty.writes) | 181 | + | t.Run(tt.name, func(t *testing.T) { |
| 174 | - | require.Empty(t, factory.opens[1].pty.writes) | 182 | + | factory := &restorePTYFactory{} |
| 183 | + | d := newTestDaemon(t, factory, stubClock{}) | |||
| 184 | + | WithSnapshotStore(tt.store)(d) | |||
| 185 | + | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: tt.allow}}) | |||
| 186 | + | ||||
| 187 | + | d.restoreSnapshots(context.Background()) | |||
| 188 | + | ||||
| 189 | + | require.Len(t, factory.opens, len(tt.wantWrites)) | |||
| 190 | + | for i, want := range tt.wantWrites { | |||
| 191 | + | if len(want) == 0 { | |||
| 192 | + | require.Empty(t, factory.opens[i].pty.writes) | |||
| 193 | + | continue | |||
| 194 | + | } | |||
| 195 | + | require.Equal(t, want, factory.opens[i].pty.writes) | |||
| 196 | + | } | |||
| 197 | + | }) | |||
| 198 | + | } | |||
| 175 | } | 199 | } | ||
| 176 | 200 | ||||
| 177 | - | func TestRestoreSnapshotsWritesAgentIDCommandAndWriteFailureIsNonFatal(t *testing.T) { | 201 | + | func TestRestoreSnapshotsProcessCommandWriteFailureIsNonFatal(t *testing.T) { |
| 178 | proc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc123"}} | 202 | proc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc123"}} | ||
| 179 | - | store := processRestoreTestStore(t, "agent", proc) | 203 | + | factory := &restorePTYFactory{writeErr: errors.New("boom")} |
| 180 | - | factory := &restorePTYFactory{} | |||
| 181 | d := newTestDaemon(t, factory, stubClock{}) | 204 | d := newTestDaemon(t, factory, stubClock{}) | ||
| 182 | - | WithSnapshotStore(store)(d) | |||
| 183 | - | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"pi"}}}) | |||
| 184 | - | ||||
| 185 | - | d.restoreSnapshots(context.Background()) | |||
| 186 | - | ||||
| 187 | - | require.Len(t, factory.opens, 1) | |||
| 188 | - | require.Equal(t, []string{"pi --resume abc123\n"}, factory.opens[0].pty.writes) | |||
| 189 | - | factory.writeErr = errors.New("boom") | |||
| 190 | - | require.NoError(t, d.restoreSession(context.Background(), snapcodec.Session{Name: "agent2", Tabs: []snapcodec.Tab{{Cols: 80, Rows: 24, Tree: layout.NewTree("pane-1"), Panes: []snapcodec.Pane{{ID: "pane-1", Cwd: "/tmp", Process: proc}}}}})) | |||
| 191 | - | } | |||
| 192 | - | ||||
| 193 | - | func TestRestoreSnapshotsRejectsAgentIDWithControlBytes(t *testing.T) { | |||
| 194 | - | proc := &snapcodec.Process{Argv: []string{"pi"}, Strategy: processStrategyPi, Opts: snapcodec.ProcessOpts{AgentSessionID: "abc\nmalicious"}} | |||
| 195 | - | store := processRestoreTestStore(t, "agent-control", proc) | |||
| 196 | - | factory := &restorePTYFactory{} | |||
| 197 | - | d := newTestDaemon(t, factory, stubClock{}) | |||
| 198 | - | WithSnapshotStore(store)(d) | |||
| 199 | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"pi"}}}) | 205 | d.ApplyConfig(domain.Config{Snapshot: domain.SnapshotConfig{RestoreProcessesSet: true, RestoreProcesses: []string{"pi"}}}) | ||
| 200 | 206 | ||||
| 201 | - | d.restoreSnapshots(context.Background()) | 207 | + | require.NoError(t, d.restoreSession(context.Background(), snapcodec.Session{Name: "agent", Tabs: []snapcodec.Tab{{Cols: 80, Rows: 24, Tree: layout.NewTree("pane-1"), Panes: []snapcodec.Pane{{ID: "pane-1", Cwd: "/tmp", Process: proc}}}}})) |
| 202 | - | ||||
| 203 | - | require.Len(t, factory.opens, 1) | |||
| 204 | - | require.Empty(t, factory.opens[0].pty.writes) | |||
| 205 | } | 208 | } | ||
| 206 | 209 | ||||
| 207 | func TestRestoreSnapshotsOpensCollapsedStackPanesWithValidPTYSize(t *testing.T) { | 210 | func TestRestoreSnapshotsOpensCollapsedStackPanesWithValidPTYSize(t *testing.T) { | ||