19906577602

Committed 03 Dec 2025 07:41PM UTC coverage: 53.737%. Remained the same

Build # 19906577602

Build Type

push

github

Committed by

web-flow

Commit Message

ssh/tui: fix channel events received by the tui out of order (#5964)

This fixes a bug where channel open/close messages could be handled by
the tunnel status model out of order if they were sent in rapid
succession.

Run Details

2 of 19 new or added lines in 2 files covered. (10.53%)

16 existing lines in 9 files now uncovered.

28979 of 53927 relevant lines covered (53.74%)

91.15 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

44.58

/pkg/ssh/cli.go

package ssh

import (
        "errors"
        "fmt"
        "io"
        "strings"

        tea "charm.land/bubbletea/v2"
        "github.com/muesli/termenv"
        "github.com/spf13/cobra"

        "github.com/pomerium/envoy-custom/api/extensions/filters/network/ssh"
        "github.com/pomerium/pomerium/config"
        "github.com/pomerium/pomerium/pkg/ssh/tui"
)

type CLI struct {
        *cobra.Command
        tui      *tea.Program
        tuiDone  chan struct{}
        msgQueue chan tea.Msg
        ptyInfo  *ssh.SSHDownstreamPTYInfo
        username string
        stdin    io.Reader
        stdout   io.Writer
        stderr   io.Writer
}

func NewCLI(
        cfg *config.Config,
        ctrl ChannelControlInterface,
        ptyInfo *ssh.SSHDownstreamPTYInfo,
        stdin io.Reader,
        stdout io.Writer,
        stderr io.Writer,
) *CLI {
        cmd := &cobra.Command{
                Use: "pomerium",
                PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
                        _, cmdIsInteractive := cmd.Annotations["interactive"]
                        switch {
                        case (ptyInfo == nil) && cmdIsInteractive:
                                return fmt.Errorf("\x1b[31m'%s' is an interactive command and requires a TTY (try passing '-t' to ssh)\x1b[0m", cmd.Use)
                        }
                        return nil
                },
        }

        cmd.CompletionOptions.DisableDefaultCmd = true
        // set a non-nil args list, otherwise it will read from os.Args by default
        cmd.SetArgs([]string{})
        cmd.SetIn(stdin)
        cmd.SetOut(stderr)
        cmd.SetErr(stderr)
        cmd.SilenceUsage = true

        cli := &CLI{
                Command:  cmd,
                tui:      nil,
                tuiDone:  make(chan struct{}),
                msgQueue: make(chan tea.Msg, 256),
                ptyInfo:  ptyInfo,
                username: *ctrl.Username(),
                stdin:    stdin,
                stdout:   stdout,
                stderr:   stderr,
        }

        if cfg.Options.IsRuntimeFlagSet(config.RuntimeFlagSSHRoutesPortal) {
                cli.AddPortalCommand(ctrl)
        }
        cli.AddTunnelCommand(ctrl)
        cli.AddLogoutCommand(ctrl)
        cli.AddWhoamiCommand(ctrl)

        return cli
}

func (cli *CLI) AddLogoutCommand(_ ChannelControlInterface) {
        cli.AddCommand(&cobra.Command{
                Use:           "logout",
                Short:         "Log out",
                SilenceErrors: true,
                RunE: func(_ *cobra.Command, _ []string) error {
                        _, _ = cli.stderr.Write([]byte("Logged out successfully\n"))
                        return ErrDeleteSessionOnExit
                },
        })
}

func (cli *CLI) AddWhoamiCommand(ctrl ChannelControlInterface) {
        cli.AddCommand(&cobra.Command{
                Use:   "whoami",
                Short: "Show details for the current session",
                RunE: func(cmd *cobra.Command, _ []string) error {
                        s, err := ctrl.FormatSession(cmd.Context())
                        if err != nil {
                                return fmt.Errorf("couldn't fetch session: %w", err)
                        }
                        _, _ = cli.stderr.Write(s)
                        return nil
                },
        })
}

type sshEnviron struct {
        Env map[string]string
}

// Environ implements termenv.Environ.
func (s *sshEnviron) Environ() []string {
        kv := make([]string, 0, len(s.Env))
        for k, v := range s.Env {
                kv = append(kv, fmt.Sprintf("%s=%s", k, v))
        }
        return kv
}

// Getenv implements termenv.Environ.
func (s *sshEnviron) Getenv(key string) string {
        return s.Env[key]
}

var _ termenv.Environ = (*sshEnviron)(nil)

const (
        ptyWidthMax  = 512
        ptyHeightMax = 512
)

func (cli *CLI) AddTunnelCommand(ctrl ChannelControlInterface) {
        cli.AddCommand(&cobra.Command{
                Use:    "tunnel",
                Short:  "tunnel status",
                Hidden: true,
                Annotations: map[string]string{
                        "interactive": "",
                },
                RunE: func(cmd *cobra.Command, _ []string) error {
                        env := &sshEnviron{
                                Env: map[string]string{
                                        "TERM":      cli.ptyInfo.TermEnv,
                                        "TTY_FORCE": "1",

                                        // Important: disables synchronized output querying which I think
                                        // might be causing the renderer to get stuck
                                        "SSH_TTY": "1",
                                },
                        }

                        prog := tui.NewTunnelStatusProgram(cmd.Context(),
                                tea.WithInput(cli.stdin),
                                tea.WithWindowSize(int(min(cli.ptyInfo.WidthColumns, ptyWidthMax)), int(min(cli.ptyInfo.HeightRows, ptyHeightMax))),
                                tea.WithOutput(termenv.NewOutput(cli.stdout, termenv.WithEnvironment(env), termenv.WithUnsafe())),
                                tea.WithEnvironment(env.Environ()),
                        )
                        cli.tui = prog.Program
                        defer close(cli.tuiDone)

                        initDone := make(chan struct{})
                        go func() {
                                ctrl.AddPortForwardUpdateListener(prog)
                                defer ctrl.RemovePortForwardUpdateListener(prog)
                                close(initDone)
                                for {
                                        select {
                                        case <-cli.tuiDone:
                                                return
                                        case msg := <-cli.msgQueue:
                                                cli.tui.Send(msg)
                                        }
                                }
                        }()
                        _, err := prog.Run()
                        <-initDone
                        if err != nil {
                                return err
                        }
                        return nil
                },
        })
}

// ErrHandoff is a sentinel error to indicate that the command triggered a handoff,
// and we should not automatically disconnect
var ErrHandoff = errors.New("handoff")

// ErrDeleteSessionOnExit is a sentinel error to indicate that the authorized
// session should be deleted once the SSH connection ends.
var ErrDeleteSessionOnExit = errors.New("delete_session_on_exit")

func (cli *CLI) AddPortalCommand(ctrl ChannelControlInterface) {
        cli.AddCommand(&cobra.Command{
                Use:   "portal",
                Short: "Interactive route portal",
                Annotations: map[string]string{
                        "interactive": "",
                },
                RunE: func(cmd *cobra.Command, _ []string) error {
                        var routes []string
                        for r := range ctrl.AllSSHRoutes() {
                                routes = append(routes, fmt.Sprintf("%s@%s", *ctrl.Username(), strings.TrimPrefix(r.From, "ssh://")))
                        }
                        env := &sshEnviron{
                                Env: map[string]string{
                                        "TERM":      cli.ptyInfo.TermEnv,
                                        "TTY_FORCE": "1",
                                        "SSH_TTY":   "1",
                                },
                        }
                        signedWidth := int(min(cli.ptyInfo.WidthColumns, ptyWidthMax))
                        signedHeight := int(min(cli.ptyInfo.HeightRows, ptyHeightMax))
                        prog := tui.NewPortalProgram(cmd.Context(), routes, max(0, signedWidth-2), max(0, signedHeight-2),
                                tea.WithInput(cli.stdin),
                                tea.WithWindowSize(signedWidth, signedHeight),
                                tea.WithOutput(termenv.NewOutput(cli.stdout, termenv.WithEnvironment(env), termenv.WithUnsafe())),
                                tea.WithEnvironment(env.Environ()),
                        )
                        cli.tui = prog.Program

                        choice, err := prog.Run()
                        if err != nil {
                                return err
                        }
                        if choice == "" {
                                return nil // quit/ctrl+c
                        }

                        username, hostname, _ := strings.Cut(choice, "@")
                        // Perform authorize check for this route
                        if username != cli.username {
                                panic("bug: username mismatch")
                        }
                        if hostname == "" {
                                panic("bug: hostname is empty")
                        }
                        handoffMsg, err := ctrl.PrepareHandoff(cmd.Context(), hostname, cli.ptyInfo)
                        if err != nil {
                                return err
                        }
                        if err := ctrl.SendControlAction(handoffMsg); err != nil {
                                return err
                        }
                        return ErrHandoff
                },
        })
}

func (cli *CLI) SendTeaMsg(msg tea.Msg) {
        select {
        case <-cli.tuiDone:
        case cli.msgQueue <- msg:
        }
}

1	package ssh
2
3	import (
4	"errors"
5	"fmt"
6	"io"
7	"strings"
8
9	tea "charm.land/bubbletea/v2"
10	"github.com/muesli/termenv"
11	"github.com/spf13/cobra"
12
13	"github.com/pomerium/envoy-custom/api/extensions/filters/network/ssh"
14	"github.com/pomerium/pomerium/config"
15	"github.com/pomerium/pomerium/pkg/ssh/tui"
16	)
17
18	type CLI struct {
19	*cobra.Command
20	tui *tea.Program
21	tuiDone chan struct{}
22	msgQueue chan tea.Msg
23	ptyInfo *ssh.SSHDownstreamPTYInfo
24	username string
25	stdin io.Reader
26	stdout io.Writer
27	stderr io.Writer
28	}
29
30	func NewCLI(
31	cfg *config.Config,
32	ctrl ChannelControlInterface,
33	ptyInfo *ssh.SSHDownstreamPTYInfo,
34	stdin io.Reader,
35	stdout io.Writer,
36	stderr io.Writer,
37	) *CLI {	29✔
38	cmd := &cobra.Command{	29✔
39	Use: "pomerium",	29✔
40	PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {	52✔
41	_, cmdIsInteractive := cmd.Annotations["interactive"]	23✔
42	switch {	23✔
43	case (ptyInfo == nil) && cmdIsInteractive:	1✔
44	return fmt.Errorf("\x1b[31m'%s' is an interactive command and requires a TTY (try passing '-t' to ssh)\x1b[0m", cmd.Use)	1✔
45	}
46	return nil	22✔
47	},
48	}
49
50	cmd.CompletionOptions.DisableDefaultCmd = true	29✔
51	// set a non-nil args list, otherwise it will read from os.Args by default	29✔
52	cmd.SetArgs([]string{})	29✔
53	cmd.SetIn(stdin)	29✔
54	cmd.SetOut(stderr)	29✔
55	cmd.SetErr(stderr)	29✔
56	cmd.SilenceUsage = true	29✔
57		29✔
58	cli := &CLI{	29✔
59	Command: cmd,	29✔
60	tui: nil,	29✔
61	tuiDone: make(chan struct{}),	29✔
62	msgQueue: make(chan tea.Msg, 256),	29✔
63	ptyInfo: ptyInfo,	29✔
64	username: *ctrl.Username(),	29✔
65	stdin: stdin,	29✔
66	stdout: stdout,	29✔
67	stderr: stderr,	29✔
68	}	29✔
69		29✔
70	if cfg.Options.IsRuntimeFlagSet(config.RuntimeFlagSSHRoutesPortal) {	49✔
71	cli.AddPortalCommand(ctrl)	20✔
72	}	20✔
73	cli.AddTunnelCommand(ctrl)	29✔
74	cli.AddLogoutCommand(ctrl)	29✔
75	cli.AddWhoamiCommand(ctrl)	29✔
76		29✔
77	return cli	29✔
78	}
79
80	func (cli *CLI) AddLogoutCommand(_ ChannelControlInterface) {	29✔
81	cli.AddCommand(&cobra.Command{	29✔
82	Use: "logout",	29✔
83	Short: "Log out",	29✔
84	SilenceErrors: true,	29✔
85	RunE: func(_ *cobra.Command, _ []string) error {	39✔
86	_, _ = cli.stderr.Write([]byte("Logged out successfully\n"))	10✔
87	return ErrDeleteSessionOnExit	10✔
88	},	10✔
89	})
90	}
91
92	func (cli *CLI) AddWhoamiCommand(ctrl ChannelControlInterface) {	29✔
93	cli.AddCommand(&cobra.Command{	29✔
94	Use: "whoami",	29✔
95	Short: "Show details for the current session",	29✔
96	RunE: func(cmd *cobra.Command, _ []string) error {	41✔
97	s, err := ctrl.FormatSession(cmd.Context())	12✔
98	if err != nil {	14✔
99	return fmt.Errorf("couldn't fetch session: %w", err)	2✔
100	}	2✔
101	_, _ = cli.stderr.Write(s)	10✔
102	return nil	10✔
103	},
104	})
105	}
106
107	type sshEnviron struct {
108	Env map[string]string
109	}
110
111	// Environ implements termenv.Environ.
112	func (s *sshEnviron) Environ() []string {	×
113	kv := make([]string, 0, len(s.Env))	×
114	for k, v := range s.Env {	×
115	kv = append(kv, fmt.Sprintf("%s=%s", k, v))	×
116	}	×
117	return kv	×
118	}
119
120	// Getenv implements termenv.Environ.
121	func (s *sshEnviron) Getenv(key string) string {	×
122	return s.Env[key]	×
123	}	×
124
125	var _ termenv.Environ = (*sshEnviron)(nil)
126
127	const (
128	ptyWidthMax = 512
129	ptyHeightMax = 512
130	)
131
132	func (cli *CLI) AddTunnelCommand(ctrl ChannelControlInterface) {	29✔
133	cli.AddCommand(&cobra.Command{	29✔
134	Use: "tunnel",	29✔
135	Short: "tunnel status",	29✔
136	Hidden: true,	29✔
137	Annotations: map[string]string{	29✔
138	"interactive": "",	29✔
139	},	29✔
140	RunE: func(cmd *cobra.Command, _ []string) error {	29✔
141	env := &sshEnviron{	×
142	Env: map[string]string{	×
143	"TERM": cli.ptyInfo.TermEnv,	×
144	"TTY_FORCE": "1",	×
145		×
146	// Important: disables synchronized output querying which I think	×
147	// might be causing the renderer to get stuck	×
148	"SSH_TTY": "1",	×
149	},	×
150	}	×
151		×
152	prog := tui.NewTunnelStatusProgram(cmd.Context(),	×
153	tea.WithInput(cli.stdin),	×
154	tea.WithWindowSize(int(min(cli.ptyInfo.WidthColumns, ptyWidthMax)), int(min(cli.ptyInfo.HeightRows, ptyHeightMax))),	×
155	tea.WithOutput(termenv.NewOutput(cli.stdout, termenv.WithEnvironment(env), termenv.WithUnsafe())),	×
156	tea.WithEnvironment(env.Environ()),	×
157	)	×
158	cli.tui = prog.Program	×
NEW 159	defer close(cli.tuiDone)	×
160		×
161	initDone := make(chan struct{})	×
162	go func() {	×
163	ctrl.AddPortForwardUpdateListener(prog)	×
NEW 164	defer ctrl.RemovePortForwardUpdateListener(prog)	×
NEW 165	close(initDone)	×
NEW 166	for {	×
NEW 167	select {	×
NEW 168	case <-cli.tuiDone:	×
NEW 169	return	×
NEW 170	case msg := <-cli.msgQueue:	×
NEW 171	cli.tui.Send(msg)	×
172	}
173	}
174	}()
UNCOV 175	_, err := prog.Run()	×
176	<-initDone	×
177	if err != nil {	×
178	return err	×
179	}	×
180	return nil	×
181	},
182	})
183	}
184
185	// ErrHandoff is a sentinel error to indicate that the command triggered a handoff,
186	// and we should not automatically disconnect
187	var ErrHandoff = errors.New("handoff")
188
189	// ErrDeleteSessionOnExit is a sentinel error to indicate that the authorized
190	// session should be deleted once the SSH connection ends.
191	var ErrDeleteSessionOnExit = errors.New("delete_session_on_exit")
192
193	func (cli *CLI) AddPortalCommand(ctrl ChannelControlInterface) {	20✔
194	cli.AddCommand(&cobra.Command{	20✔
195	Use: "portal",	20✔
196	Short: "Interactive route portal",	20✔
197	Annotations: map[string]string{	20✔
198	"interactive": "",	20✔
199	},	20✔
200	RunE: func(cmd *cobra.Command, _ []string) error {	20✔
201	var routes []string	×
202	for r := range ctrl.AllSSHRoutes() {	×
203	routes = append(routes, fmt.Sprintf("%s@%s", *ctrl.Username(), strings.TrimPrefix(r.From, "ssh://")))	×
204	}	×
205	env := &sshEnviron{	×
206	Env: map[string]string{	×
207	"TERM": cli.ptyInfo.TermEnv,	×
208	"TTY_FORCE": "1",	×
209	"SSH_TTY": "1",	×
210	},	×
211	}	×
212	signedWidth := int(min(cli.ptyInfo.WidthColumns, ptyWidthMax))	×
213	signedHeight := int(min(cli.ptyInfo.HeightRows, ptyHeightMax))	×
214	prog := tui.NewPortalProgram(cmd.Context(), routes, max(0, signedWidth-2), max(0, signedHeight-2),	×
215	tea.WithInput(cli.stdin),	×
216	tea.WithWindowSize(signedWidth, signedHeight),	×
217	tea.WithOutput(termenv.NewOutput(cli.stdout, termenv.WithEnvironment(env), termenv.WithUnsafe())),	×
218	tea.WithEnvironment(env.Environ()),	×
219	)	×
220	cli.tui = prog.Program	×
221		×
222	choice, err := prog.Run()	×
223	if err != nil {	×
224	return err	×
225	}	×
226	if choice == "" {	×
227	return nil // quit/ctrl+c	×
228	}	×
229
230	username, hostname, _ := strings.Cut(choice, "@")	×
231	// Perform authorize check for this route	×
232	if username != cli.username {	×
233	panic("bug: username mismatch")	×
234	}
235	if hostname == "" {	×
236	panic("bug: hostname is empty")	×
237	}
238	handoffMsg, err := ctrl.PrepareHandoff(cmd.Context(), hostname, cli.ptyInfo)	×
239	if err != nil {	×
240	return err	×
241	}	×
242	if err := ctrl.SendControlAction(handoffMsg); err != nil {	×
243	return err	×
244	}	×
245	return ErrHandoff	×
246	},
247	})
248	}
249
250	func (cli *CLI) SendTeaMsg(msg tea.Msg) {	×
NEW 251	select {	×
NEW 252	case <-cli.tuiDone:	×
NEW 253	case cli.msgQueue <- msg:	×
254	}
255	}

pomerium / pomerium / 19906577602

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous