Files
coder/pty/pty_windows.go
Spike Curtis 36ffdce065 Return proper exit code on ssh with TTY (#3192)
* Return proper exit code on ssh with TTY

Signed-off-by: Spike Curtis <spike@coder.com>

* Fix revive lint

Signed-off-by: Spike Curtis <spike@coder.com>

* Fix Windows exit code for missing command

Signed-off-by: Spike Curtis <spike@coder.com>

* Fix close error handling on agent TTY

Signed-off-by: Spike Curtis <spike@coder.com>
2022-07-27 14:23:28 -05:00

144 lines
3.2 KiB
Go

//go:build windows
// +build windows
package pty
import (
"os"
"os/exec"
"sync"
"unsafe"
"golang.org/x/sys/windows"
"golang.org/x/xerrors"
)
var (
kernel32 = windows.NewLazySystemDLL("kernel32.dll")
procResizePseudoConsole = kernel32.NewProc("ResizePseudoConsole")
procCreatePseudoConsole = kernel32.NewProc("CreatePseudoConsole")
procClosePseudoConsole = kernel32.NewProc("ClosePseudoConsole")
)
// See: https://docs.microsoft.com/en-us/windows/console/creating-a-pseudoconsole-session
func newPty() (PTY, error) {
// We use the CreatePseudoConsole API which was introduced in build 17763
vsn := windows.RtlGetVersion()
if vsn.MajorVersion < 10 ||
vsn.BuildNumber < 17763 {
// If the CreatePseudoConsole API is not available, we fall back to a simpler
// implementation that doesn't create an actual PTY - just uses os.Pipe
return nil, xerrors.Errorf("pty not supported")
}
ptyWindows := &ptyWindows{}
var err error
ptyWindows.inputRead, ptyWindows.inputWrite, err = os.Pipe()
if err != nil {
return nil, err
}
ptyWindows.outputRead, ptyWindows.outputWrite, err = os.Pipe()
consoleSize := uintptr(80) + (uintptr(80) << 16)
ret, _, err := procCreatePseudoConsole.Call(
consoleSize,
uintptr(ptyWindows.inputRead.Fd()),
uintptr(ptyWindows.outputWrite.Fd()),
0,
uintptr(unsafe.Pointer(&ptyWindows.console)),
)
if int32(ret) < 0 {
return nil, xerrors.Errorf("create pseudo console (%d): %w", int32(ret), err)
}
return ptyWindows, nil
}
type ptyWindows struct {
console windows.Handle
outputWrite *os.File
outputRead *os.File
inputWrite *os.File
inputRead *os.File
closeMutex sync.Mutex
closed bool
}
type windowsProcess struct {
// cmdDone protects access to cmdErr: anything reading cmdErr should read from cmdDone first.
cmdDone chan any
cmdErr error
proc *os.Process
}
func (p *ptyWindows) Output() ReadWriter {
return ReadWriter{
Reader: p.outputRead,
Writer: p.outputWrite,
}
}
func (p *ptyWindows) Input() ReadWriter {
return ReadWriter{
Reader: p.inputRead,
Writer: p.inputWrite,
}
}
func (p *ptyWindows) Resize(height uint16, width uint16) error {
// Taken from: https://github.com/microsoft/hcsshim/blob/54a5ad86808d761e3e396aff3e2022840f39f9a8/internal/winapi/zsyscall_windows.go#L144
ret, _, err := procResizePseudoConsole.Call(uintptr(p.console), uintptr(*((*uint32)(unsafe.Pointer(&windows.Coord{
Y: int16(height),
X: int16(width),
})))))
if ret != 0 {
return err
}
return nil
}
func (p *ptyWindows) Close() error {
p.closeMutex.Lock()
defer p.closeMutex.Unlock()
if p.closed {
return nil
}
p.closed = true
_ = p.outputWrite.Close()
_ = p.outputRead.Close()
_ = p.inputWrite.Close()
_ = p.inputRead.Close()
ret, _, err := procClosePseudoConsole.Call(uintptr(p.console))
if ret < 0 {
return xerrors.Errorf("close pseudo console: %w", err)
}
return nil
}
func (p *windowsProcess) waitInternal() {
defer close(p.cmdDone)
state, err := p.proc.Wait()
if err != nil {
p.cmdErr = err
return
}
if !state.Success() {
p.cmdErr = &exec.ExitError{ProcessState: state}
return
}
}
func (p *windowsProcess) Wait() error {
<-p.cmdDone
return p.cmdErr
}
func (p *windowsProcess) Kill() error {
return p.proc.Kill()
}