From 46efd31c6a81aecc0094374aefcc710853df284b Mon Sep 17 00:00:00 2001 From: Ayman Bagabas Date: Tue, 5 Apr 2022 12:19:20 -0400 Subject: [PATCH] feat: add bubbletea program middleware This allows passing a program handler that returns tea.Program to allow for more advanced use cases such as using p.Send() to send messages to tea.Program. Fixes: https://github.com/charmbracelet/wish/issues/37 --- bubbletea/tea.go | 34 +++++++- examples/bubbleteaprogram/main.go | 126 ++++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 4 deletions(-) create mode 100644 examples/bubbleteaprogram/main.go diff --git a/bubbletea/tea.go b/bubbletea/tea.go index f8d9dd7..886954e 100644 --- a/bubbletea/tea.go +++ b/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 { diff --git a/examples/bubbleteaprogram/main.go b/examples/bubbleteaprogram/main.go new file mode 100644 index 0000000..c43016e --- /dev/null +++ b/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) +}