diff --git a/go.mod b/go.mod index 8b472ea..384b3ae 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.20 require ( github.com/MicahParks/keyfunc v1.9.0 + github.com/charmbracelet/x/conpty v0.1.0 github.com/creack/pty v1.1.18 github.com/dustin/go-humanize v1.0.1 github.com/gliderlabs/ssh v0.3.5 @@ -23,6 +24,7 @@ require ( golang.org/x/crypto v0.7.0 golang.org/x/exp v0.0.0-20221205204356-47842c84f3db golang.org/x/sync v0.2.0 + golang.org/x/sys v0.20.0 google.golang.org/grpc v1.55.0 google.golang.org/protobuf v1.30.0 inet.af/tcpproxy v0.0.0-20221017015627-91f861402626 @@ -33,6 +35,7 @@ require ( github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86 // indirect github.com/cloudflare/circl v1.3.3 // indirect github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect github.com/golang/protobuf v1.5.3 // indirect @@ -57,7 +60,6 @@ require ( golang.org/x/mod v0.11.0 // indirect golang.org/x/net v0.10.0 // indirect golang.org/x/oauth2 v0.7.0 // indirect - golang.org/x/sys v0.8.0 // indirect golang.org/x/text v0.9.0 // indirect golang.org/x/tools v0.9.1 // indirect google.golang.org/appengine v1.6.7 // indirect diff --git a/go.sum b/go.sum index 6ee3224..b2467d5 100644 --- a/go.sum +++ b/go.sum @@ -14,6 +14,10 @@ github.com/bwesterb/go-ristretto v1.2.0/go.mod h1:fUIoIZaG73pV5biE2Blr2xEzDoMj7N github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/charmbracelet/x/conpty v0.1.0 h1:4zc8KaIcbiL4mghEON8D72agYtSeIgq8FSThSPQIb+U= +github.com/charmbracelet/x/conpty v0.1.0/go.mod h1:rMFsDJoDwVmiYM10aD4bH2XiRgwI7NYJtQgl5yskjEQ= +github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86 h1:JSt3B+U9iqk37QUU2Rvb6DSBYRLtWqFqfxf8l5hOZUA= +github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= @@ -214,8 +218,8 @@ golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220825204002-c680a09ffe64/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU= -golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y= +golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.0.0-20220722155259-a9ba230a4035/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= diff --git a/pkg/sshd/pty_unix.go b/pkg/sshd/pty_unix.go index 6ac15e4..03e2b87 100644 --- a/pkg/sshd/pty_unix.go +++ b/pkg/sshd/pty_unix.go @@ -3,6 +3,7 @@ package sshd import ( + "context" "fmt" "io" "os" @@ -14,22 +15,23 @@ import ( "github.com/gliderlabs/ssh" ) -func handlePty(session io.ReadWriter, ptyReq ssh.Pty, winCh <-chan ssh.Window, cmd *exec.Cmd) error { +func startPty(_ context.Context, session io.ReadWriter, ptyReq ssh.Pty, winCh <-chan ssh.Window, cmd *exec.Cmd) (func() error, error) { cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term)) ptyFile, err := pty.Start(cmd) if err != nil { - return err + return nil, err } - defer ptyFile.Close() - go syncWinSize(ptyFile, winCh) go func() { _, _ = io.Copy(ptyFile, session) // stdin }() - _, _ = io.Copy(session, ptyFile) // stdout - return nil + return func() error { + _, _ = io.Copy(session, ptyFile) // stdout + _ = ptyFile.Close() + return cmd.Wait() + }, nil } func syncWinSize(ptyFile *os.File, winCh <-chan ssh.Window) { diff --git a/pkg/sshd/pty_windows.go b/pkg/sshd/pty_windows.go index 35d0224..e00ab34 100644 --- a/pkg/sshd/pty_windows.go +++ b/pkg/sshd/pty_windows.go @@ -3,15 +3,166 @@ package sshd import ( + "context" + "errors" + "fmt" "io" + "math" "os/exec" + "sync" + "syscall" + "github.com/charmbracelet/x/conpty" "github.com/gliderlabs/ssh" + "golang.org/x/sys/windows" ) -func handlePty(session io.ReadWriter, _ ssh.Pty, _ <-chan ssh.Window, cmd *exec.Cmd) error { - cmd.Stdin = session - cmd.Stdout = session - cmd.Stderr = session - return cmd.Start() +type conPTYSession struct { + mu sync.Mutex + *conpty.ConPty + stopped bool +} + +func startPty(ctx context.Context, session io.ReadWriter, ptyReq ssh.Pty, winCh <-chan ssh.Window, cmd *exec.Cmd) (func() error, error) { + if err := validateConPTYSize(ptyReq.Window.Width, ptyReq.Window.Height); err != nil { + return nil, err + } + + pseudoconsole, err := conpty.New(ptyReq.Window.Width, ptyReq.Window.Height, 0) + if err != nil { + return nil, err + } + conPTY := &conPTYSession{ConPty: pseudoconsole} + + _, processHandle, err := pseudoconsole.Spawn(cmd.Path, cmd.Args, &syscall.ProcAttr{ + Dir: cmd.Dir, + Env: cmd.Env, + Sys: cmd.SysProcAttr, + }) + if err != nil { + _ = pseudoconsole.Close() + return nil, err + } + process := windows.Handle(processHandle) + + sessionClosed := make(chan struct{}) + go resizeConPTY(conPTY, winCh, sessionClosed) + go func() { + _, _ = io.Copy(pseudoconsole.InPipe(), session) + }() + outputDone := make(chan struct{}) + go func() { + defer close(outputDone) + drainConPTY(session, pseudoconsole.OutPipe()) + }() + + waitDone := make(chan processResult, 1) + go func() { + waitDone <- waitForProcess(process) + }() + + return func() error { + defer windows.CloseHandle(process) + disconnected := false + var result processResult + select { + case result = <-waitDone: + case <-ctx.Done(): + disconnected = true + case <-sessionClosed: + disconnected = true + } + + conPTY.stop() + if disconnected { + _ = pseudoconsole.InPipe().Close() + _ = pseudoconsole.OutPipe().Close() + _ = windows.TerminateProcess(process, 1) + result = <-waitDone + } + + closeErr := pseudoconsole.Close() + <-outputDone + if result.err != nil { + return result.err + } + if !disconnected && result.exitCode != 0 { + return fmt.Errorf("process exited with code %d", result.exitCode) + } + if !disconnected && closeErr != nil { + return closeErr + } + return nil + }, nil +} + +func (c *conPTYSession) resize(width, height int) error { + if err := validateConPTYSize(width, height); err != nil { + return err + } + + c.mu.Lock() + defer c.mu.Unlock() + if c.stopped { + return errors.New("ConPTY session is closed") + } + return c.ConPty.Resize(width, height) +} + +func (c *conPTYSession) stop() { + c.mu.Lock() + defer c.mu.Unlock() + c.stopped = true +} + +func resizeConPTY(pseudoconsole *conPTYSession, winCh <-chan ssh.Window, sessionClosed chan<- struct{}) { + defer close(sessionClosed) + for win := range winCh { + if win.Width > 0 && win.Height > 0 { + _ = pseudoconsole.resize(win.Width, win.Height) + } + } +} + +func validateConPTYSize(width, height int) error { + if width <= 0 || width > math.MaxInt16 || height <= 0 || height > math.MaxInt16 { + return fmt.Errorf("invalid ConPTY size: %dx%d", width, height) + } + return nil +} + +func drainConPTY(dst io.Writer, src io.Reader) { + buf := make([]byte, 32*1024) + for { + n, err := src.Read(buf) + if n > 0 && dst != nil { + if _, writeErr := dst.Write(buf[:n]); writeErr != nil { + dst = nil + } + } + if err != nil { + return + } + } +} + +type processResult struct { + exitCode uint32 + err error +} + +func waitForProcess(process windows.Handle) processResult { + status, err := windows.WaitForSingleObject(process, windows.INFINITE) + if err != nil { + return processResult{err: err} + } + if status != windows.WAIT_OBJECT_0 { + return processResult{err: fmt.Errorf("unexpected process wait status: %d", status)} + } + + var exitCode uint32 + if err := windows.GetExitCodeProcess(process, &exitCode); err != nil { + return processResult{err: err} + } + return processResult{exitCode: exitCode} } diff --git a/pkg/sshd/pty_windows_test.go b/pkg/sshd/pty_windows_test.go index e027dd6..b7cf781 100644 --- a/pkg/sshd/pty_windows_test.go +++ b/pkg/sshd/pty_windows_test.go @@ -3,16 +3,18 @@ package sshd import ( - "bytes" "context" "crypto/rand" "crypto/rsa" + "io" "net" "os" + "os/exec" "strings" "testing" "time" + "github.com/gliderlabs/ssh" gossh "golang.org/x/crypto/ssh" ) @@ -71,19 +73,130 @@ func TestWindowsPTYSession(t *testing.T) { t.Fatal(err) } - const marker = "breakpoint-windows-pty-ok" - var stdout, stderr bytes.Buffer - session.Stdin = strings.NewReader("echo " + marker + "\r\nexit\r\n") - session.Stdout = &stdout - session.Stderr = &stderr + stdin, err := session.StdinPipe() + if err != nil { + t.Fatal(err) + } + stdout, err := session.StdoutPipe() + if err != nil { + t.Fatal(err) + } + output := make(chan []byte, 16) + go copyOutput(stdout, output) if err := session.Shell(); err != nil { t.Fatal(err) } - if err := session.Wait(); err != nil { - t.Fatalf("Windows shell failed: %v\nstdout:\n%s\nstderr:\n%s", err, stdout.String(), stderr.String()) + + if _, err := io.WriteString(stdin, "set /a 6*7\r"); err != nil { + t.Fatal(err) + } + waitForOutput(t, output, "42") + + if err := session.WindowChange(37, 113); err != nil { + t.Fatal(err) + } + if _, err := io.WriteString(stdin, `powershell -NoProfile -Command "$d=(Get-Date).AddSeconds(5); do {$s=[Console]::WindowWidth.ToString()+'x'+[Console]::WindowHeight; Start-Sleep -Milliseconds 20} until ($s -eq '113x37' -or (Get-Date) -gt $d); Write-Output ('SIZE='+$s)"`+"\r"); err != nil { + t.Fatal(err) + } + waitForOutput(t, output, "SIZE=113x37") + + if _, err := io.WriteString(stdin, "exit\r"); err != nil { + t.Fatal(err) + } + waitDone := make(chan error, 1) + go func() { + waitDone <- session.Wait() + }() + select { + case err := <-waitDone: + if err != nil { + t.Fatal(err) + } + case <-time.After(10 * time.Second): + t.Fatal("Windows shell did not exit") + } +} + +func copyOutput(src io.Reader, output chan<- []byte) { + defer close(output) + buf := make([]byte, 4096) + for { + n, err := src.Read(buf) + if n > 0 { + chunk := append([]byte(nil), buf[:n]...) + output <- chunk + } + if err != nil { + return + } + } +} + +func waitForOutput(t *testing.T, output <-chan []byte, expected string) { + t.Helper() + timer := time.NewTimer(10 * time.Second) + defer timer.Stop() + + var received strings.Builder + for { + select { + case chunk, ok := <-output: + if !ok { + t.Fatalf("output closed before %q was received; output:\n%s", expected, received.String()) + } + received.Write(chunk) + if strings.Contains(received.String(), expected) { + return + } + case <-timer.C: + t.Fatalf("timed out waiting for %q; output:\n%s", expected, received.String()) + } + } +} + +func TestWindowsPTYStopsWhenSessionCloses(t *testing.T) { + inputReader, inputWriter := io.Pipe() + defer inputWriter.Close() + + winCh := make(chan ssh.Window) + waitPty, err := startPty( + context.Background(), + struct { + io.Reader + io.Writer + }{inputReader, io.Discard}, + ssh.Pty{Window: ssh.Window{Width: 80, Height: 24}}, + winCh, + exec.Command(os.Getenv("COMSPEC")), + ) + if err != nil { + t.Fatal(err) + } + + close(winCh) + waitDone := make(chan error, 1) + go func() { + waitDone <- waitPty() + }() + select { + case err := <-waitDone: + if err != nil { + t.Fatal(err) + } + case <-time.After(10 * time.Second): + t.Fatal("Windows shell did not stop after its SSH session closed") + } +} + +func TestValidateConPTYSize(t *testing.T) { + if err := validateConPTYSize(80, 24); err != nil { + t.Fatal(err) + } + if err := validateConPTYSize(1<<15, 24); err == nil { + t.Fatal("oversized ConPTY width was accepted") } - if !strings.Contains(stdout.String(), marker) { - t.Fatalf("Windows shell stdout did not contain %q:\n%s\nstderr:\n%s", marker, stdout.String(), stderr.String()) + if err := validateConPTYSize(80, 1<<15); err == nil { + t.Fatal("oversized ConPTY height was accepted") } } diff --git a/pkg/sshd/sshd.go b/pkg/sshd/sshd.go index 62f851f..302cd57 100644 --- a/pkg/sshd/sshd.go +++ b/pkg/sshd/sshd.go @@ -100,20 +100,24 @@ func MakeServer(ctx context.Context, opts SSHServerOpts) (*SSHServer, error) { opts.InteractiveMOTD(session) } - if err := handlePty(session, ptyReq, winCh, cmd); err != nil { + waitPty, err := startPty(ctx, session, ptyReq, winCh, cmd) + if err != nil { sessionLog.Err(err).Msg("pty start failed") session.Exit(1) return } - } else { - cmd.Stdin = session - cmd.Stdout = session - cmd.Stderr = session - if err := cmd.Start(); err != nil { - sessionLog.Err(err).Msg("start failed") - session.Exit(1) - return - } + err = waitPty() + sessionLog.Info().Err(err).Msg("ssh session end") + return + } + + cmd.Stdin = session + cmd.Stdout = session + cmd.Stderr = session + if err := cmd.Start(); err != nil { + sessionLog.Err(err).Msg("start failed") + session.Exit(1) + return } // XXX pass exit code to caller?