Use ConPTY for Windows SSH sessions

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

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=

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) {

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}
 }