feat: add bubbletea program middleware

bubbletea/tea.go

@@ -22,6 +22,15 @@ type BubbleTeaHandler = Handler // nolint: revive
 // start it with the tea.ProgramOptions returned.
 type Handler func(ssh.Session) (tea.Model, []tea.ProgramOption)
 
+// ProgramHandler is the function Bubble Tea apps implement to hook into the SSH
+// Middleware. This should return a new tea.Program. This handler is different
+// from the default handler in that it returns a tea.Program instead of
+// (tea.Model, tea.ProgramOptions).
+//
+// Make sure to set the tea.WithInput and tea.WithOutput to the ssh.Session
+// otherwise the program will not function properly.
+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
@@ -35,14 +44,31 @@ func Middleware(bth Handler) wish.Middleware {
 // 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)
+}
+
+// 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.
+//
+// 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 {
 	return func(sh ssh.Handler) ssh.Handler {
 		lipgloss.SetColorProfile(cp)
 		return func(s ssh.Session) {
 			errc := make(chan error, 1)
-			m, opts := bth(s)
-			if m != nil {
-				opts = append(opts, tea.WithInput(s), tea.WithOutput(s))
-				p := tea.NewProgram(m, opts...)
+			p := bth(s)
+			if p != nil {
 				_, windowChanges, _ := s.Pty()
 				go func() {
 					for {

examples/bubbleteaprogram/main.go

@@ -0,0 +1,126 @@
+package main
+
+// An example Bubble Tea server. This will put an ssh session into alt screen
+// and continually print up to date terminal information.
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/wish"
+	bm "github.com/charmbracelet/wish/bubbletea"
+	lm "github.com/charmbracelet/wish/logging"
+	"github.com/gliderlabs/ssh"
+	"github.com/muesli/termenv"
+)
+
+const host = "localhost"
+const port = 23234
+
+func main() {
+	s, err := wish.NewServer(
+		wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
+		wish.WithHostKeyPath(".ssh/term_info_ed25519"),
+		wish.WithMiddleware(
+			myCustomBubbleteaMiddleware(),
+			lm.Middleware(),
+		),
+	)
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	log.Printf("Starting SSH server on %s:%d", host, port)
+	go func() {
+		if err = s.ListenAndServe(); err != nil {
+			log.Fatalln(err)
+		}
+	}()
+
+	<-done
+	log.Println("Stopping SSH server")
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer func() { cancel() }()
+	if err := s.Shutdown(ctx); err != nil {
+		log.Fatalln(err)
+	}
+}
+
+// You can write your own custom bubbletea middleware that wraps tea.Program.
+// Make sure you set the program input and output to ssh.Session.
+func myCustomBubbleteaMiddleware() wish.Middleware {
+	newProg := func(m tea.Model, opts ...tea.ProgramOption) *tea.Program {
+		p := tea.NewProgram(m, opts...)
+		go func() {
+			for {
+				select {
+				case <-time.After(1 * time.Second):
+					p.Send(timeMsg(time.Now()))
+				}
+			}
+		}()
+		return p
+	}
+	teaHandler := func(s ssh.Session) *tea.Program {
+		pty, _, active := s.Pty()
+		if !active {
+			fmt.Println("no active terminal, skipping")
+			s.Exit(1)
+			return nil
+		}
+		m := model{
+			term:   pty.Term,
+			width:  pty.Window.Width,
+			height: pty.Window.Height,
+			time:   time.Now(),
+		}
+		return newProg(m, tea.WithInput(s), tea.WithOutput(s), tea.WithAltScreen())
+	}
+	return bm.MiddlewareWithProgramHandler(teaHandler, termenv.ANSI256)
+}
+
+// Just a generic tea.Model to demo terminal information of ssh.
+type model struct {
+	term   string
+	width  int
+	height int
+	time   time.Time
+}
+
+type timeMsg time.Time
+
+func (m model) Init() tea.Cmd {
+	return nil
+}
+
+func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case timeMsg:
+		m.time = time.Time(msg)
+	case tea.WindowSizeMsg:
+		m.height = msg.Height
+		m.width = msg.Width
+	case tea.KeyMsg:
+		switch msg.String() {
+		case "q", "ctrl+c":
+			return m, tea.Quit
+		}
+	}
+	return m, nil
+}
+
+func (m model) View() string {
+	s := "Your term is %s\n"
+	s += "Your window size is x: %d y: %d\n"
+	s += "Time: " + m.time.Format(time.RFC1123) + "\n\n"
+	s += "Press 'q' to quit\n"
+	return fmt.Sprintf(s, m.term, m.width, m.height)
+}