feat: go back to the list after disconnecting

blocking_reader.go → blocking/reader.go

@@ -1,31 +1,32 @@
-package wishlist
+package blocking
 
 import (
 	"io"
 	"time"
 )
 
-// blockingReader is an io.Reader that blocks until the underlying reader
+// Reader is an io.Reader that blocks until the underlying reader until
 // returns something other than io.EOF.
 //
 // on EOF, it'll keep trying to read every 10ms.
 //
 // The purpose of this is to be used to "emulate a STDIN" (which never EOFs)
 // from another io.Reader, e.g. a bytes.Buffer.
-type blockingReader struct {
+type Reader struct {
 	r io.Reader
 }
 
-func newBlockingReader(r io.Reader) io.Reader {
-	return blockingReader{r}
+// New wraps a given io.Reader into a BlockingReader
+func New(r io.Reader) Reader {
+	return Reader{r: r}
 }
 
 type readResult struct {
 	n int
 	e error
 }
 
-func (r blockingReader) Read(data []byte) (int, error) {
+func (r Reader) Read(data []byte) (int, error) {
 	readch := make(chan readResult, 1)
 
 	go func() {

blocking_reader_test.go → blocking/reader_test.go

@@ -1,4 +1,4 @@
-package wishlist
+package blocking
 
 import (
 	"io"
@@ -9,7 +9,7 @@ import (
 )
 
 func TestBlockingReader(t *testing.T) {
-	r := newBlockingReader(&testReader{
+	r := New(&testReader{
 		results: []readResult{
 			{n: 0, e: io.EOF},              // EOF should be ignored
 			{n: 12, e: nil},                // return normally

client.go

@@ -6,12 +6,14 @@ import (
 	"io"
 	"log"
 
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/muesli/termenv"
 	gossh "golang.org/x/crypto/ssh"
 )
 
 // SSHClient is a SSH client.
 type SSHClient interface {
-	Connect(e *Endpoint) error
+	For(e *Endpoint) tea.ExecCommand
 }
 
 func createSession(conf *gossh.ClientConfig, e *Endpoint) (*gossh.Session, *gossh.Client, closers, error) {
@@ -72,3 +74,9 @@ func firstNonEmpty(ss ...string) string {
 	}
 	return ""
 }
+
+func resetPty(w io.Writer) {
+	fmt.Fprint(w, termenv.CSI+termenv.ExitAltScreenSeq)
+	fmt.Fprint(w, termenv.CSI+termenv.ResetSeq+"m")
+	fmt.Fprintf(w, termenv.CSI+termenv.EraseDisplaySeq, 2) // nolint:gomnd
+}

client_auth.go

@@ -23,6 +23,7 @@ var errNoRemoteAgent = fmt.Errorf("no agent forwarded")
 //
 // it first tries to use ssh-agent, if that's not available, it creates and uses a new key pair.
 func remoteBestAuthMethod(s ssh.Session) (gossh.AuthMethod, agent.Agent, closers, error) {
+	// TODO: we should probably make password protected keys work here too
 	method, agt, cls, err := tryRemoteAuthAgent(s)
 	if err != nil {
 		return nil, nil, nil, err
@@ -47,7 +48,7 @@ func remoteBestAuthMethod(s ssh.Session) (gossh.AuthMethod, agent.Agent, closers
 // It'll return a nil list if none of the methods is available.
 func localBestAuthMethod(e *Endpoint) ([]gossh.AuthMethod, error) {
 	var methods []gossh.AuthMethod
-	if method, err := tryLocalAgent(); err != nil || method != nil {
+	if method, err := tryLocalAgent(); err == nil && method != nil {
 		methods = append(methods, method)
 	}
 
@@ -80,7 +81,7 @@ func getLocalAgent() (agent.Agent, error) {
 	}
 	conn, err := net.Dial("unix", socket)
 	if err != nil {
-		return nil, fmt.Errorf("failed to connecto to SSH_AUTH_SOCK: %w", err)
+		return nil, fmt.Errorf("failed to connect to SSH_AUTH_SOCK: %w", err)
 	}
 	return agent.NewClient(conn), nil
 }

client_local.go

@@ -3,11 +3,14 @@ package wishlist
 import (
 	"context"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"os/user"
 	"path/filepath"
 
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/muesli/cancelreader"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 	"golang.org/x/term"
@@ -20,33 +23,73 @@ func NewLocalSSHClient() SSHClient {
 
 type localClient struct{}
 
-func (c *localClient) Connect(e *Endpoint) error {
+func (c *localClient) For(e *Endpoint) tea.ExecCommand {
+	return &localSession{
+		endpoint: e,
+	}
+}
+
+type localSession struct {
+	// endpoint we are connecting to
+	endpoint *Endpoint
+
+	stdin          io.Reader
+	stdout, stderr io.Writer
+}
+
+func (s *localSession) SetStdin(r io.Reader) {
+	s.stdin = r
+}
+
+func (s *localSession) SetStdout(w io.Writer) {
+	s.stdout = w
+}
+
+func (s *localSession) SetStderr(w io.Writer) {
+	s.stderr = w
+}
+
+func (s *localSession) Run() error {
+	resetPty(s.stdout)
+
 	user, err := user.Current()
 	if err != nil {
 		return fmt.Errorf("failed to get current username: %w", err)
 	}
 
-	methods, err := localBestAuthMethod(e)
+	methods, err := localBestAuthMethod(s.endpoint)
 	if err != nil {
 		return fmt.Errorf("failed to setup a authentication method: %w", err)
 	}
+
 	conf := &ssh.ClientConfig{
-		User:            firstNonEmpty(e.User, user.Username),
+		User:            firstNonEmpty(s.endpoint.User, user.Username),
 		Auth:            methods,
-		HostKeyCallback: hostKeyCallback(e, filepath.Join(user.HomeDir, ".ssh/known_hosts")),
+		HostKeyCallback: hostKeyCallback(s.endpoint, filepath.Join(user.HomeDir, ".ssh/known_hosts")),
 	}
 
-	session, client, cls, err := createSession(conf, e)
+	session, client, cls, err := createSession(conf, s.endpoint)
 	defer cls.close()
 	if err != nil {
 		return fmt.Errorf("failed to create session: %w", err)
 	}
+	defer closers{func() error {
+		rc, ok := session.Stdin.(cancelreader.CancelReader)
+		if ok && !rc.Cancel() {
+			return fmt.Errorf("could not cancel reader")
+		}
+		return nil
+	}}.close()
 
-	session.Stdout = os.Stdout
-	session.Stderr = os.Stderr
-	session.Stdin = os.Stdin
+	session.Stdout = s.stdout
+	session.Stderr = s.stderr
+	stdin, err := cancelreader.NewReader(s.stdin)
+	if err != nil {
+		return fmt.Errorf("could not create cancel reader")
+	}
+	session.Stdin = stdin
 
-	if e.ForwardAgent {
+	if s.endpoint.ForwardAgent {
 		log.Println("forwarding SSH agent")
 		agt, err := getLocalAgent()
 		if err != nil {
@@ -63,7 +106,7 @@ func (c *localClient) Connect(e *Endpoint) error {
 		}
 	}
 
-	if e.RequestTTY || e.RemoteCommand == "" {
+	if s.endpoint.RequestTTY || s.endpoint.RemoteCommand == "" {
 		fd := int(os.Stdout.Fd())
 		if !term.IsTerminal(fd) {
 			return fmt.Errorf("requested a TTY, but current session is not TTY, aborting")
@@ -92,13 +135,13 @@ func (c *localClient) Connect(e *Endpoint) error {
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
-		go c.notifyWindowChanges(ctx, session)
+		go s.notifyWindowChanges(ctx, session)
 	} else {
 		log.Println("did not request a tty")
 	}
 
-	if e.RemoteCommand == "" {
+	if s.endpoint.RemoteCommand == "" {
 		return shellAndWait(session)
 	}
-	return runAndWait(session, e.RemoteCommand)
+	return runAndWait(session, s.endpoint.RemoteCommand)
 }

client_remote.go

@@ -5,45 +5,78 @@ import (
 	"io"
 	"log"
 
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/wishlist/blocking"
 	"github.com/gliderlabs/ssh"
-	"github.com/muesli/termenv"
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 )
 
 type remoteClient struct {
+	// parent session
 	session ssh.Session
-	stdin   io.Reader
+
+	// stdin, which is usually multiplexed from the session stdin
+	stdin io.Reader
+
+	// whether to exhaust stdin first or not.
+	// if coming from the list, youll want to do that, otherwise you likely
+	// dont, as it might hang the connection waiting for something to read.
+	exhaust bool
+}
+
+func (c *remoteClient) For(e *Endpoint) tea.ExecCommand {
+	if c.exhaust {
+		_, _ = io.ReadAll(c.stdin)
+	}
+	return &remoteSession{
+		endpoint:      e,
+		parentSession: c.session,
+		stdin:         blocking.New(c.stdin),
+	}
 }
 
-func (c *remoteClient) Connect(e *Endpoint) error {
-	resetPty(c.session)
+type remoteSession struct {
+	// endpoint we are connecting to
+	endpoint *Endpoint
 
-	method, agt, closers, err := remoteBestAuthMethod(c.session)
-	defer closers.close()
+	// the parent session (ie the session running the listing)
+	parentSession ssh.Session
+
+	stdin io.Reader
+}
+
+func (s *remoteSession) SetStdin(r io.Reader)  {}
+func (s *remoteSession) SetStdout(w io.Writer) {}
+func (s *remoteSession) SetStderr(w io.Writer) {}
+
+func (s *remoteSession) Run() error {
+	resetPty(s.parentSession)
+
+	method, agt, closers, err := remoteBestAuthMethod(s.parentSession)
 	if err != nil {
 		return fmt.Errorf("failed to find an auth method: %w", err)
 	}
+	defer closers.close()
 
 	conf := &gossh.ClientConfig{
-		User:            firstNonEmpty(e.User, c.session.User()),
-		HostKeyCallback: hostKeyCallback(e, ".wishlist/known_hosts"),
+		User:            firstNonEmpty(s.endpoint.User, s.parentSession.User()),
+		HostKeyCallback: hostKeyCallback(s.endpoint, ".wishlist/known_hosts"),
 		Auth:            []gossh.AuthMethod{method},
 	}
-
-	session, client, cl, err := createSession(conf, e)
+	session, client, cl, err := createSession(conf, s.endpoint)
 	defer cl.close()
 	if err != nil {
 		return fmt.Errorf("failed to create session: %w", err)
 	}
 
-	log.Printf("%s connect to %q, %s", c.session.User(), e.Name, c.session.RemoteAddr().String())
+	log.Printf("%s connect to %q, %s", s.parentSession.User(), s.endpoint.Name, s.parentSession.RemoteAddr().String())
 
-	session.Stdout = c.session
-	session.Stderr = c.session.Stderr()
-	session.Stdin = c.stdin
+	session.Stdout = s.parentSession
+	session.Stderr = s.parentSession.Stderr()
+	session.Stdin = s.stdin
 
-	if e.ForwardAgent {
+	if s.endpoint.ForwardAgent {
 		log.Println("forwarding SSH agent")
 		if agt == nil {
 			return fmt.Errorf("requested ForwardAgent, but no agent is available")
@@ -56,9 +89,9 @@ func (c *remoteClient) Connect(e *Endpoint) error {
 		}
 	}
 
-	if e.RemoteCommand == "" || e.RequestTTY {
+	if s.endpoint.RemoteCommand == "" || s.endpoint.RequestTTY {
 		log.Println("requesting tty")
-		pty, winch, ok := c.session.Pty()
+		pty, winch, ok := s.parentSession.Pty()
 		if !ok {
 			return fmt.Errorf("requested a tty, but current session doesn't allow one")
 		}
@@ -69,16 +102,16 @@ func (c *remoteClient) Connect(e *Endpoint) error {
 
 		done := make(chan bool, 1)
 		defer func() { done <- true }()
-		go c.notifyWindowChanges(session, done, winch)
+		go s.notifyWindowChanges(session, done, winch)
 	}
 
-	if e.RemoteCommand == "" {
+	if s.endpoint.RemoteCommand == "" {
 		return shellAndWait(session)
 	}
-	return runAndWait(session, e.RemoteCommand)
+	return runAndWait(session, s.endpoint.RemoteCommand)
 }
 
-func (c *remoteClient) notifyWindowChanges(session *gossh.Session, done <-chan bool, winch <-chan ssh.Window) {
+func (s *remoteSession) notifyWindowChanges(session *gossh.Session, done <-chan bool, winch <-chan ssh.Window) {
 	for {
 		select {
 		case <-done:
@@ -95,9 +128,3 @@ func (c *remoteClient) notifyWindowChanges(session *gossh.Session, done <-chan b
 		}
 	}
 }
-
-func resetPty(w io.Writer) {
-	fmt.Fprint(w, termenv.CSI+termenv.ExitAltScreenSeq)
-	fmt.Fprint(w, termenv.CSI+termenv.ResetSeq+"m")
-	fmt.Fprintf(w, termenv.CSI+termenv.EraseDisplaySeq, 2) // nolint:gomnd
-}

client_signals_unix.go

@@ -14,7 +14,7 @@ import (
 	"golang.org/x/term"
 )
 
-func (c *localClient) notifyWindowChanges(ctx context.Context, session *ssh.Session) {
+func (s *localSession) notifyWindowChanges(ctx context.Context, session *ssh.Session) {
 	sig := make(chan os.Signal, 1)
 	signal.Notify(sig, syscall.SIGWINCH)
 	defer func() {

client_signals_windows.go

@@ -5,8 +5,9 @@ package wishlist
 
 import (
 	"context"
+
 	"golang.org/x/crypto/ssh"
 )
 
 // not available because windows does not implement siscall.SIGWINCH.
-func (c *localClient) notifyWindowChanges(ctx context.Context, session *ssh.Session) {}
+func (c *localSession) notifyWindowChanges(ctx context.Context, session *ssh.Session) {}