fix(ssh): allocate pty and use the latest tea middleware

Ayman Bagabas created

Change summary

pkg/ssh/middleware.go             | 102 ++++++++++++++++----------------
pkg/ssh/session.go                |  33 +---------
pkg/ssh/ssh.go                    |   1 
testscript/script_test.go         |   8 +
testscript/testdata/ui-home.txtar |   2 
5 files changed, 62 insertions(+), 84 deletions(-)

Detailed changes

pkg/ssh/middleware.go 🔗

@@ -78,67 +78,65 @@ var cliCommandCounter = promauto.NewCounterVec(prometheus.CounterOpts{
 // This middleware must be run after the ContextMiddleware.
 func CommandMiddleware(sh ssh.Handler) ssh.Handler {
 	return func(s ssh.Session) {
-		func() {
-			_, _, ptyReq := s.Pty()
-			if ptyReq {
-				return
-			}
-
-			ctx := s.Context()
-			cfg := config.FromContext(ctx)
+		_, _, ptyReq := s.Pty()
+		if ptyReq {
+			sh(s)
+			return
+		}
 
-			args := s.Command()
-			cliCommandCounter.WithLabelValues(cmd.CommandName(args)).Inc()
-			rootCmd := &cobra.Command{
-				Short:        "Soft Serve is a self-hostable Git server for the command line.",
-				SilenceUsage: true,
-			}
-			rootCmd.CompletionOptions.DisableDefaultCmd = true
+		ctx := s.Context()
+		cfg := config.FromContext(ctx)
 
-			rootCmd.SetUsageTemplate(cmd.UsageTemplate)
-			rootCmd.SetUsageFunc(cmd.UsageFunc)
+		args := s.Command()
+		cliCommandCounter.WithLabelValues(cmd.CommandName(args)).Inc()
+		rootCmd := &cobra.Command{
+			Short:        "Soft Serve is a self-hostable Git server for the command line.",
+			SilenceUsage: true,
+		}
+		rootCmd.CompletionOptions.DisableDefaultCmd = true
+
+		rootCmd.SetUsageTemplate(cmd.UsageTemplate)
+		rootCmd.SetUsageFunc(cmd.UsageFunc)
+		rootCmd.AddCommand(
+			cmd.GitUploadPackCommand(),
+			cmd.GitUploadArchiveCommand(),
+			cmd.GitReceivePackCommand(),
+			cmd.RepoCommand(),
+			cmd.SettingsCommand(),
+			cmd.UserCommand(),
+			cmd.InfoCommand(),
+			cmd.PubkeyCommand(),
+			cmd.SetUsernameCommand(),
+			cmd.JWTCommand(),
+			cmd.TokenCommand(),
+		)
+
+		if cfg.LFS.Enabled {
 			rootCmd.AddCommand(
-				cmd.GitUploadPackCommand(),
-				cmd.GitUploadArchiveCommand(),
-				cmd.GitReceivePackCommand(),
-				cmd.RepoCommand(),
-				cmd.SettingsCommand(),
-				cmd.UserCommand(),
-				cmd.InfoCommand(),
-				cmd.PubkeyCommand(),
-				cmd.SetUsernameCommand(),
-				cmd.JWTCommand(),
-				cmd.TokenCommand(),
+				cmd.GitLFSAuthenticateCommand(),
 			)
 
-			if cfg.LFS.Enabled {
+			if cfg.LFS.SSHEnabled {
 				rootCmd.AddCommand(
-					cmd.GitLFSAuthenticateCommand(),
+					cmd.GitLFSTransfer(),
 				)
-
-				if cfg.LFS.SSHEnabled {
-					rootCmd.AddCommand(
-						cmd.GitLFSTransfer(),
-					)
-				}
 			}
+		}
 
-			rootCmd.SetArgs(args)
-			if len(args) == 0 {
-				// otherwise it'll default to os.Args, which is not what we want.
-				rootCmd.SetArgs([]string{"--help"})
-			}
-			rootCmd.SetIn(s)
-			rootCmd.SetOut(s)
-			rootCmd.SetErr(s.Stderr())
-			rootCmd.SetContext(ctx)
-
-			if err := rootCmd.ExecuteContext(ctx); err != nil {
-				s.Exit(1) // nolint: errcheck
-				return
-			}
-		}()
-		sh(s)
+		rootCmd.SetArgs(args)
+		if len(args) == 0 {
+			// otherwise it'll default to os.Args, which is not what we want.
+			rootCmd.SetArgs([]string{"--help"})
+		}
+		rootCmd.SetIn(s)
+		rootCmd.SetOut(s)
+		rootCmd.SetErr(s.Stderr())
+		rootCmd.SetContext(ctx)
+
+		if err := rootCmd.ExecuteContext(ctx); err != nil {
+			s.Exit(1) // nolint: errcheck
+			return
+		}
 	}
 }
 

pkg/ssh/session.go 🔗

@@ -1,11 +1,9 @@
 package ssh
 
 import (
-	"strings"
 	"time"
 
 	tea "github.com/charmbracelet/bubbletea"
-	"github.com/charmbracelet/lipgloss"
 	"github.com/charmbracelet/soft-serve/pkg/access"
 	"github.com/charmbracelet/soft-serve/pkg/backend"
 	"github.com/charmbracelet/soft-serve/pkg/config"
@@ -13,7 +11,7 @@ import (
 	"github.com/charmbracelet/soft-serve/pkg/ui/common"
 	"github.com/charmbracelet/ssh"
 	"github.com/charmbracelet/wish"
-	"github.com/muesli/termenv"
+	bm "github.com/charmbracelet/wish/bubbletea"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
@@ -54,19 +52,18 @@ func SessionHandler(s ssh.Session) *tea.Program {
 		}
 	}
 
-	envs := &sessionEnv{s}
-	output := lipgloss.NewRenderer(s, termenv.WithColorCache(true), termenv.WithEnvironment(envs))
+	output := bm.MakeRenderer(s)
 	c := common.NewCommon(ctx, output, pty.Window.Width, pty.Window.Height)
 	c.SetValue(common.ConfigKey, cfg)
 	m := NewUI(c, initialRepo)
-	p := tea.NewProgram(m,
-		tea.WithInput(s),
-		tea.WithOutput(s),
+	opts := bm.MakeOptions(s)
+	opts = append(opts,
 		tea.WithAltScreen(),
 		tea.WithoutCatchPanics(),
 		tea.WithMouseCellMotion(),
 		tea.WithContext(ctx),
 	)
+	p := tea.NewProgram(m, opts...)
 
 	tuiSessionCounter.WithLabelValues(initialRepo, pty.Term).Inc()
 
@@ -78,23 +75,3 @@ func SessionHandler(s ssh.Session) *tea.Program {
 
 	return p
 }
-
-var _ termenv.Environ = &sessionEnv{}
-
-type sessionEnv struct {
-	ssh.Session
-}
-
-func (s *sessionEnv) Environ() []string {
-	pty, _, _ := s.Pty()
-	return append(s.Session.Environ(), "TERM="+pty.Term)
-}
-
-func (s *sessionEnv) Getenv(key string) string {
-	for _, env := range s.Environ() {
-		if strings.HasPrefix(env, key+"=") {
-			return strings.TrimPrefix(env, key+"=")
-		}
-	}
-	return ""
-}

pkg/ssh/ssh.go 🔗

@@ -88,6 +88,7 @@ func NewSSHServer(ctx context.Context) (*SSHServer, error) {
 	s.srv, err = wish.NewServer(
 		ssh.PublicKeyAuth(s.PublicKeyHandler),
 		ssh.KeyboardInteractiveAuth(s.KeyboardInteractiveHandler),
+		ssh.AllocatePty(),
 		wish.WithAddress(cfg.SSH.ListenAddr),
 		wish.WithHostKeyPath(cfg.SSH.KeyPath),
 		wish.WithMiddleware(mw...),

testscript/script_test.go 🔗

@@ -238,6 +238,10 @@ func cmdUI(key ssh.Signer) func(ts *testscript.TestScript, neg bool, args []stri
 		stdin, err := sess.StdinPipe()
 		check(ts, err, neg)
 
+		err = sess.RequestPty("dumb", 40, 80, ssh.TerminalModes{})
+		check(ts, err, neg)
+		check(ts, sess.Start(""), neg)
+
 		in, err := strconv.Unquote(args[0])
 		check(ts, err, neg)
 		reader := strings.NewReader(in)
@@ -256,9 +260,7 @@ func cmdUI(key ssh.Signer) func(ts *testscript.TestScript, neg bool, args []stri
 			}
 		}()
 
-		err = sess.RequestPty("dumb", 40, 80, ssh.TerminalModes{})
-		check(ts, err, neg)
-		check(ts, sess.Run(""), neg)
+		check(ts, sess.Wait(), neg)
 	}
 }
 

testscript/testdata/ui-home.txtar 🔗

@@ -10,7 +10,7 @@ ui '"    q"'
 cp stdout home.txt
 grep 'Test Soft Serve' home.txt
 grep '• Repositories' home.txt
-grep 'No items found' home.txt
+grep 'No items' home.txt
 
 # test about tab
 ui '"\t    q"'