feat: use real pty (#197)

* feat: use pty Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: deps Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: deps Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: update example * fix: deps * docs: fix example * feat: pass lipgloss.Renderer down to the tea.App * fix: go mod tidy Signed-off-by: Carlos Alexandro Becker <[email protected]> * improvements Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: pty Co-authored-by: Ayman Bagabas <[email protected]> Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: better diff * chore: typo * fix: godocs * fix: allocate pty on macos Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: improvements based on #219 Signed-off-by: Carlos Alexandro Becker <[email protected]> Co-authored-by: Ayman Bagabas <[email protected]> * chore: godoc * fix: review Co-authored-by: Ayman Bagabas <[email protected]> Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: example Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: tea program handler Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: examples * refactor: improve p!=nil handling Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: ensure session envs are available to renderer (#223) * fix: rename func to makeoptions * fix: not too much breaking * chore: fix gitignore * fix: dep Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: update charmbracelet/ssh * fix: update dep Signed-off-by: Carlos Alexandro Becker <[email protected]> * chore: go mod tidy Signed-off-by: Carlos Alexandro Becker <[email protected]> * chore: update example --------- Signed-off-by: Carlos Alexandro Becker <[email protected]> Co-authored-by: Ayman Bagabas <[email protected]>
charmbracelet · Jan 18, 2024 · 7f44d6d · 7f44d6d
1 parent d2fab68
commit 7f44d6d
Show file tree

Hide file tree

Showing 12 changed files with 319 additions and 68 deletions.
diff --git a/.gitignore b/.gitignore
@@ -6,6 +6,8 @@ examples/git/.repos
 .repos
 .ssh
 coverage.txt
+id_ed25519
+id_ed25519.pub
 
 # MacOS specific
 .DS_Store
diff --git a/bubbletea/tea.go b/bubbletea/tea.go
@@ -3,6 +3,7 @@ package bubbletea
 
 import (
 	"context"
+	"strings"
 
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
@@ -34,72 +35,125 @@ type Handler func(ssh.Session) (tea.Model, []tea.ProgramOption)
 type ProgramHandler func(ssh.Session) *tea.Program
 
 // Middleware takes a Handler and hooks the input and output for the
-// ssh.Session into the tea.Program. It also captures window resize events and
-// sends them to the tea.Program as tea.WindowSizeMsgs. By default a 256 color
-// profile will be used when rendering with Lip Gloss.
+// ssh.Session into the tea.Program.
+//
+// It also captures window resize events and sends them to the tea.Program
+// as tea.WindowSizeMsgs.
 func Middleware(bth Handler) wish.Middleware {
-	return MiddlewareWithColorProfile(bth, termenv.ANSI256)
+	return MiddlewareWithProgramHandler(newDefaultProgramHandler(bth), termenv.Ascii)
 }
 
-// MiddlewareWithColorProfile allows you to specify the number of colors
-// returned by the server when using Lip Gloss. The number of colors supported
-// by an SSH client's terminal cannot be detected by the server but this will
-// allow for manually setting the color profile on all SSH connections.
-func MiddlewareWithColorProfile(bth Handler, cp termenv.Profile) wish.Middleware {
-	h := func(s ssh.Session) *tea.Program {
-		m, opts := bth(s)
-		if m == nil {
-			return nil
-		}
-		opts = append(opts, tea.WithInput(s), tea.WithOutput(s))
-		return tea.NewProgram(m, opts...)
-	}
-	return MiddlewareWithProgramHandler(h, cp)
+// MiddlewareWithColorProfile allows you to specify the minimum number of colors
+// this program needs to work properly.
+//
+// If the client's color profile has less colors than p, p will be forced.
+// Use with caution.
+func MiddlewareWithColorProfile(bth Handler, p termenv.Profile) wish.Middleware {
+	return MiddlewareWithProgramHandler(newDefaultProgramHandler(bth), p)
 }
 
 // MiddlewareWithProgramHandler allows you to specify the ProgramHandler to be
-// able to access the underlying tea.Program. This is useful for creating custom
-// middlewars that need access to tea.Program for instance to use p.Send() to
-// send messages to tea.Program.
+// able to access the underlying tea.Program, and the minimum supported color
+// profile.
+//
+// This is useful for creating custom middlewares that need access to
+// tea.Program for instance to use p.Send() to send messages to tea.Program.
 //
 // Make sure to set the tea.WithInput and tea.WithOutput to the ssh.Session
-// otherwise the program will not function properly.
-func MiddlewareWithProgramHandler(bth ProgramHandler, cp termenv.Profile) wish.Middleware {
-	// XXX: This is a hack to make sure the default Termenv output color
-	// profile is set before the program starts. Ideally, we want a Lip Gloss
-	// renderer per session.
-	lipgloss.SetColorProfile(cp)
-	return func(sh ssh.Handler) ssh.Handler {
+// otherwise the program will not function properly. The recommended way
+// of doing so is by using MakeOptions.
+//
+// If the client's color profile has less colors than p, p will be forced.
+// Use with caution.
+func MiddlewareWithProgramHandler(bth ProgramHandler, p termenv.Profile) wish.Middleware {
+	return func(h ssh.Handler) ssh.Handler {
 		return func(s ssh.Session) {
+			s.Context().SetValue(minColorProfileKey, p)
+			_, windowChanges, ok := s.Pty()
+			if !ok {
+				wish.Fatalln(s, "no active terminal, skipping")
+				return
+			}
 			p := bth(s)
-			if p != nil {
-				_, windowChanges, _ := s.Pty()
-				ctx, cancel := context.WithCancel(s.Context())
-				go func() {
-					for {
-						select {
-						case <-ctx.Done():
-							if p != nil {
-								p.Quit()
-								return
-							}
-						case w := <-windowChanges:
-							if p != nil {
-								p.Send(tea.WindowSizeMsg{Width: w.Width, Height: w.Height})
-							}
-						}
+			if p == nil {
+				h(s)
+				return
+			}
+			ctx, cancel := context.WithCancel(s.Context())
+			go func() {
+				for {
+					select {
+					case <-ctx.Done():
+						p.Quit()
+						return
+					case w := <-windowChanges:
+						p.Send(tea.WindowSizeMsg{Width: w.Width, Height: w.Height})
 					}
-				}()
-				if _, err := p.Run(); err != nil {
-					log.Error("app exit with error", "error", err)
 				}
-				// p.Kill() will force kill the program if it's still running,
-				// and restore the terminal to its original state in case of a
-				// tui crash
-				p.Kill()
-				cancel()
+			}()
+			if _, err := p.Run(); err != nil {
+				log.Error("app exit with error", "error", err)
 			}
-			sh(s)
+			// p.Kill() will force kill the program if it's still running,
+			// and restore the terminal to its original state in case of a
+			// tui crash
+			p.Kill()
+			cancel()
+			h(s)
+		}
+	}
+}
+
+var minColorProfileKey struct{}
+
+var profileNames = [4]string{"TrueColor", "ANSI256", "ANSI", "Ascii"}
+
+// MakeRenderer returns a lipgloss renderer for the current session.
+// This function handle PTYs as well, and should be used to style your application.
+func MakeRenderer(s ssh.Session) *lipgloss.Renderer {
+	cp, ok := s.Context().Value(minColorProfileKey).(termenv.Profile)
+	if !ok {
+		cp = termenv.Ascii
+	}
+	r := newRenderer(s)
+	if r.ColorProfile() > cp {
+		wish.Printf(s, "Warning: Client's terminal is %q, forcing %q\r\n", profileNames[r.ColorProfile()], profileNames[cp])
+		r.SetColorProfile(cp)
+	}
+	return r
+}
+
+// MakeOptions returns the tea.WithInput and tea.WithOutput program options
+// taking into account possible Emulated or Allocated PTYs.
+func MakeOptions(s ssh.Session) []tea.ProgramOption {
+	return makeOpts(s)
+}
+
+type sshEnviron []string
+
+var _ termenv.Environ = sshEnviron(nil)
+
+// Environ implements termenv.Environ.
+func (e sshEnviron) Environ() []string {
+	return e
+}
+
+// Getenv implements termenv.Environ.
+func (e sshEnviron) Getenv(k string) string {
+	for _, v := range e {
+		if strings.HasPrefix(v, k+"=") {
+			return v[len(k)+1:]
+		}
+	}
+	return ""
+}
+
+func newDefaultProgramHandler(bth Handler) ProgramHandler {
+	return func(s ssh.Session) *tea.Program {
+		m, opts := bth(s)
+		if m == nil {
+			return nil
 		}
+		return tea.NewProgram(m, append(opts, makeOpts(s)...)...)
 	}
 }
diff --git a/bubbletea/tea_other.go b/bubbletea/tea_other.go
@@ -0,0 +1,24 @@
+//go:build !linux && !darwin && !freebsd && !dragonfly && !netbsd && !openbsd && !solaris
+// +build !linux,!darwin,!freebsd,!dragonfly,!netbsd,!openbsd,!solaris
+
+package bubbletea
+
+import (
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/charmbracelet/ssh"
+	"github.com/muesli/termenv"
+)
+
+func makeOpts(s ssh.Session) []tea.ProgramOption {
+	return []tea.ProgramOption{
+		tea.WithInput(s),
+		tea.WithOutput(s),
+	}
+}
+
+func newRenderer(s ssh.Session) *lipgloss.Renderer {
+	pty, _, _ := s.Pty()
+	env := sshEnviron(append(s.Environ(), "TERM="+pty.Term))
+	return lipgloss.NewRenderer(s, termenv.WithEnvironment(env), termenv.WithUnsafe(), termenv.WithColorCache(true))
+}
diff --git a/bubbletea/tea_unix.go b/bubbletea/tea_unix.go
@@ -0,0 +1,35 @@
+//go:build darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
+// +build darwin dragonfly freebsd linux netbsd openbsd solaris
+
+package bubbletea
+
+import (
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/charmbracelet/ssh"
+	"github.com/muesli/termenv"
+)
+
+func makeOpts(s ssh.Session) []tea.ProgramOption {
+	pty, _, ok := s.Pty()
+	if !ok || s.EmulatedPty() {
+		return []tea.ProgramOption{
+			tea.WithInput(s),
+			tea.WithOutput(s),
+		}
+	}
+
+	return []tea.ProgramOption{
+		tea.WithInput(pty.Slave),
+		tea.WithOutput(pty.Slave),
+	}
+}
+
+func newRenderer(s ssh.Session) *lipgloss.Renderer {
+	pty, _, ok := s.Pty()
+	env := sshEnviron(append(s.Environ(), "TERM="+pty.Term))
+	if !ok || pty.Slave == nil {
+		return lipgloss.NewRenderer(s, termenv.WithEnvironment(env), termenv.WithUnsafe(), termenv.WithColorCache(true))
+	}
+	return lipgloss.NewRenderer(pty.Slave, termenv.WithEnvironment(env), termenv.WithColorCache(true))
+}
diff --git a/examples/.gitignore b/examples/.gitignore
@@ -0,0 +1,2 @@
+id_ed25519*
+file.txt
diff --git a/examples/bubbleteaprogram/main.go b/examples/bubbleteaprogram/main.go
@@ -18,7 +18,6 @@ import (
 	"github.com/charmbracelet/wish"
 	bm "github.com/charmbracelet/wish/bubbletea"
 	lm "github.com/charmbracelet/wish/logging"
-	"github.com/muesli/termenv"
 )
 
 const (
@@ -83,9 +82,9 @@ func myCustomBubbleteaMiddleware() wish.Middleware {
 			height: pty.Window.Height,
 			time:   time.Now(),
 		}
-		return newProg(m, tea.WithInput(s), tea.WithOutput(s), tea.WithAltScreen())
+		return newProg(m, append(bm.MakeOptions(s), tea.WithAltScreen())...)
 	}
-	return bm.MiddlewareWithProgramHandler(teaHandler, termenv.ANSI256)
+	return bm.MiddlewareWithProgramHandler(teaHandler)
 }
 
 // Just a generic tea.Model to demo terminal information of ssh.

diff --git a/examples/go.mod b/examples/go.mod
@@ -7,9 +7,8 @@ require (
 	github.com/charmbracelet/bubbletea v0.25.0
 	github.com/charmbracelet/lipgloss v0.9.1
 	github.com/charmbracelet/log v0.3.1
-	github.com/charmbracelet/ssh v0.0.0-20240104172912-e11ae277b249
+	github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371
 	github.com/charmbracelet/wish v0.5.0
-	github.com/muesli/termenv v0.15.2
 	github.com/spf13/cobra v1.8.0
 	golang.org/x/crypto v0.18.0
 )
@@ -22,8 +21,11 @@ require (
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/charmbracelet/keygen v0.5.0 // indirect
+	github.com/charmbracelet/x/errors v0.0.0-20240117030013-d31dba354651 // indirect
+	github.com/charmbracelet/x/exp/term v0.0.0-20240117031359-6e25c76a1efe // indirect
 	github.com/cloudflare/circl v1.3.7 // indirect
 	github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 // indirect
+	github.com/creack/pty v1.1.21 // indirect
 	github.com/cyphar/filepath-securejoin v0.2.4 // indirect
 	github.com/emirpasic/gods v1.18.1 // indirect
 	github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 // indirect
@@ -41,11 +43,13 @@ require (
 	github.com/muesli/ansi v0.0.0-20211018074035-2e021307bc4b // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
 	github.com/muesli/reflow v0.3.0 // indirect
+	github.com/muesli/termenv v0.15.2 // indirect
 	github.com/pjbgf/sha1cd v0.3.0 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/sergi/go-diff v1.1.0 // indirect
 	github.com/skeema/knownhosts v1.2.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/u-root/u-root v0.11.0 // indirect
 	github.com/xanzy/ssh-agent v0.3.3 // indirect
 	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
 	golang.org/x/mod v0.13.0 // indirect

diff --git a/examples/go.sum b/examples/go.sum
@@ -23,14 +23,20 @@ github.com/charmbracelet/lipgloss v0.9.1 h1:PNyd3jvaJbg4jRHKWXnCj1akQm4rh8dbEzN1
 github.com/charmbracelet/lipgloss v0.9.1/go.mod h1:1mPmG4cxScwUQALAAnacHaigiiHB9Pmr+v1VEawJl6I=
 github.com/charmbracelet/log v0.3.1 h1:TjuY4OBNbxmHWSwO3tosgqs5I3biyY8sQPny/eCMTYw=
 github.com/charmbracelet/log v0.3.1/go.mod h1:OR4E1hutLsax3ZKpXbgUqPtTjQfrh1pG3zwHGWuuq8g=
-github.com/charmbracelet/ssh v0.0.0-20240104172912-e11ae277b249 h1:M1Q/UIbi9Cfla0HK3A9puhQhb8ZPkA5HQxeSYcVrGpo=
-github.com/charmbracelet/ssh v0.0.0-20240104172912-e11ae277b249/go.mod h1:A1H384KV/cJcSKofWjdSIb+dfbikXiW6449EluL3qJI=
+github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371 h1:lgr2JbKDeq13Ar9gwmwELlJAF6R13pP3Kp37C9KLVM8=
+github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371/go.mod h1:l/6/exIt0QMHEW5IiqoY9HyPlSuXAm6xkngFIJgiEYI=
+github.com/charmbracelet/x/errors v0.0.0-20240117030013-d31dba354651 h1:3RXpZWGWTOeVXCTv0Dnzxdv/MhNUkBfEcbaTY0zrTQI=
+github.com/charmbracelet/x/errors v0.0.0-20240117030013-d31dba354651/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0=
+github.com/charmbracelet/x/exp/term v0.0.0-20240117031359-6e25c76a1efe h1:HeRgHWxOTu7l73rKsa5BRAeaUenmNyomiPCUHXv/y14=
+github.com/charmbracelet/x/exp/term v0.0.0-20240117031359-6e25c76a1efe/go.mod h1:kOOxxyxgAFQVcR5yQJWTuLjzt5dR2pcgwy3WaLEudjE=
 github.com/cloudflare/circl v1.3.3/go.mod h1:5XYMA4rFBvNIrhs50XuiBJ15vF2pZn4nnUKZrLbUZFA=
 github.com/cloudflare/circl v1.3.7 h1:qlCDlTPz2n9fu58M0Nh1J/JzcFpfgkFHHX3O35r5vcU=
 github.com/cloudflare/circl v1.3.7/go.mod h1:sRTcRWXGLrKw6yIGJ+l7amYJFfAXbZG0kBSc8r4zxgA=
 github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 h1:q2hJAaP1k2wIvVRd/hEHD7lacgqrCPS+k8g1MndzfWY=
 github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81/go.mod h1:YynlIjWYF8myEu6sdkwKIvGQq+cOckRm6So2avqoYAk=
 github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
+github.com/creack/pty v1.1.21 h1:1/QdRyBaHHJP61QkWMXlOIBfsgdDeeKfK8SYVUWJKf0=
+github.com/creack/pty v1.1.21/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/cyphar/filepath-securejoin v0.2.4 h1:Ugdm7cg7i6ZK6x3xDF1oEu1nfkyfH53EtKeQYTC3kyg=
 github.com/cyphar/filepath-securejoin v0.2.4/go.mod h1:aPGpWjXOXUn2NCNjFvBE6aRxGGx79pTxQpKOJNYHHl4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -106,6 +112,9 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/u-root/gobusybox/src v0.0.0-20221229083637-46b2883a7f90 h1:zTk5683I9K62wtZ6eUa6vu6IWwVHXPnoKK5n2unAwv0=
+github.com/u-root/u-root v0.11.0 h1:6gCZLOeRyevw7gbTwMj3fKxnr9+yHFlgF3N7udUVNO8=
+github.com/u-root/u-root v0.11.0/go.mod h1:DBkDtiZyONk9hzVEdB/PWI9B4TxDkElWlVTHseglrZY=
 github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM=
 github.com/xanzy/ssh-agent v0.3.3/go.mod h1:6dzNDKs0J9rVPHPhaGCukekBHKqfl+L3KghI1Bc68Uw=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=

diff --git a/examples/multichat/main.go b/examples/multichat/main.go
@@ -18,7 +18,6 @@ import (
 	"github.com/charmbracelet/wish"
 	bm "github.com/charmbracelet/wish/bubbletea"
 	lm "github.com/charmbracelet/wish/logging"
-	"github.com/muesli/termenv"
 )
 
 const (
@@ -41,16 +40,14 @@ func (a *app) send(msg tea.Msg) {
 
 func newApp() *app {
 	a := new(app)
-
 	s, err := wish.NewServer(
 		wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
 		wish.WithHostKeyPath(".ssh/term_info_ed25519"),
 		wish.WithMiddleware(
-			bm.MiddlewareWithProgramHandler(a.ProgramHandler, termenv.ANSI256),
+			bm.MiddlewareWithProgramHandler(a.ProgramHandler),
 			lm.Middleware(),
 		),
 	)
-
 	if err != nil {
 		log.Fatalln(err)
 	}
@@ -88,7 +85,7 @@ func (a *app) ProgramHandler(s ssh.Session) *tea.Program {
 	model.app = a
 	model.id = s.User()
 
-	p := tea.NewProgram(model, tea.WithOutput(s), tea.WithInput(s))
+	p := tea.NewProgram(model, bm.MakeOptions(s)...)
 	a.progs = append(a.progs, p)
 
 	return p