cmd_windows.go

//go:build windows
// +build windows

package pty

import (
	"bytes"
	"errors"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"unicode/utf16"
	"unsafe"

	"golang.org/x/sys/windows"
)

// copied from os/exec.Cmd for platform compatibility
// we need to use startupInfoEx for pty support, but os/exec.Cmd only have
// support for startupInfo on windows, so we have to rewrite some internal
// logic for windows while keep its behavior compatible with other platforms.

// windowExecCmd represents an external command being prepared or run.
//
// A cmd cannot be reused after calling its Run, Output or CombinedOutput
// methods.
type windowExecCmd struct {
	cmd        *exec.Cmd
	waitCalled bool
	conPty     *WindowsPty
	attrList   *windows.ProcThreadAttributeListContainer
}

var errProcessNotStarted = errors.New("exec: process has not started yet")

func (c *windowExecCmd) close() error {
	c.attrList.Delete()
	_ = c.conPty.Close()

	return nil
}

func (c *windowExecCmd) Run() error {
	err := c.Start()
	if err != nil {
		return err
	}

	return c.Wait()
}

func (c *windowExecCmd) Wait() error {
	if c.cmd.Process == nil {
		return errProcessNotStarted
	}

	var err error

	if c.waitCalled {
		return errors.New("exec: wait was already called")
	}

	c.waitCalled = true
	c.cmd.ProcessState, err = c.cmd.Process.Wait()
	if err != nil {
		return err
	}

	err = c.close()
	if err != nil {
		return err
	}

	if !c.cmd.ProcessState.Success() {
		return &exec.ExitError{ProcessState: c.cmd.ProcessState}
	}

	return nil
}

func (c *windowExecCmd) StdinPipe() (io.WriteCloser, error) {
	return nil, ErrUnsupported
}

func (c *windowExecCmd) StdoutPipe() (io.ReadCloser, error) {
	return nil, ErrUnsupported
}

func (c *windowExecCmd) StderrPipe() (io.ReadCloser, error) {
	return nil, ErrUnsupported
}

func (c *windowExecCmd) Output() ([]byte, error) {
	if c.cmd.Stdout != nil {
		return nil, errors.New("exec: Stdout already set")
	}

	var stdout bytes.Buffer
	c.cmd.Stdout = &stdout

	err := c.Run()
	return stdout.Bytes(), err
}

func (c *windowExecCmd) CombinedOutput() ([]byte, error) {
	if c.cmd.Stdout != nil {
		return nil, errors.New("exec: Stdout already set")
	}
	if c.cmd.Stderr != nil {
		return nil, errors.New("exec: Stderr already set")
	}
	var b bytes.Buffer
	c.cmd.Stdout = &b
	c.cmd.Stderr = &b
	err := c.Run()
	return b.Bytes(), err
}

func (c *windowExecCmd) argv() []string {
	if len(c.cmd.Args) > 0 {
		return c.cmd.Args
	}

	return []string{c.cmd.Path}
}

//
// Helpers for working with Windows. These are exact copies of the same utilities found in the go stdlib.
//

func lookExtensions(path, dir string) (string, error) {
	if filepath.Base(path) == path {
		path = filepath.Join(".", path)
	}

	if dir == "" {
		return exec.LookPath(path)
	}

	if filepath.VolumeName(path) != "" {
		return exec.LookPath(path)
	}

	if len(path) > 1 && os.IsPathSeparator(path[0]) {
		return exec.LookPath(path)
	}

	dirandpath := filepath.Join(dir, path)

	// We assume that LookPath will only add file extension.
	lp, err := exec.LookPath(dirandpath)
	if err != nil {
		return "", err
	}

	ext := strings.TrimPrefix(lp, dirandpath)

	return path + ext, nil
}

func dedupEnvCase(caseInsensitive bool, env []string) []string {
	// Construct the output in reverse order, to preserve the
	// last occurrence of each key.
	out := make([]string, 0, len(env))
	saw := make(map[string]bool, len(env))
	for n := len(env); n > 0; n-- {
		kv := env[n-1]

		i := strings.Index(kv, "=")
		if i == 0 {
			// We observe in practice keys with a single leading "=" on Windows.
			// TODO(#49886): Should we consume only the first leading "=" as part
			// of the key, or parse through arbitrarily many of them until a non-"="?
			i = strings.Index(kv[1:], "=") + 1
		}
		if i < 0 {
			if kv != "" {
				// The entry is not of the form "key=value" (as it is required to be).
				// Leave it as-is for now.
				// TODO(#52436): should we strip or reject these bogus entries?
				out = append(out, kv)
			}
			continue
		}
		k := kv[:i]
		if caseInsensitive {
			k = strings.ToLower(k)
		}
		if saw[k] {
			continue
		}

		saw[k] = true
		out = append(out, kv)
	}

	// Now reverse the slice to restore the original order.
	for i := 0; i < len(out)/2; i++ {
		j := len(out) - i - 1
		out[i], out[j] = out[j], out[i]
	}

	return out
}

func addCriticalEnv(env []string) []string {
	for _, kv := range env {
		eq := strings.Index(kv, "=")
		if eq < 0 {
			continue
		}
		k := kv[:eq]
		if strings.EqualFold(k, "SYSTEMROOT") {
			// We already have it.
			return env
		}
	}

	return append(env, "SYSTEMROOT="+os.Getenv("SYSTEMROOT"))
}

func execEnvDefault(sys *syscall.SysProcAttr) (env []string, err error) {
	if sys == nil || sys.Token == 0 {
		return syscall.Environ(), nil
	}

	var block *uint16
	err = windows.CreateEnvironmentBlock(&block, windows.Token(sys.Token), false)
	if err != nil {
		return nil, err
	}

	defer windows.DestroyEnvironmentBlock(block)
	blockp := uintptr(unsafe.Pointer(block))

	for {

		// find NUL terminator
		end := unsafe.Pointer(blockp)
		for *(*uint16)(end) != 0 {
			end = unsafe.Pointer(uintptr(end) + 2)
		}

		n := (uintptr(end) - uintptr(unsafe.Pointer(blockp))) / 2
		if n == 0 {
			// environment block ends with empty string
			break
		}

		entry := (*[(1 << 30) - 1]uint16)(unsafe.Pointer(blockp))[:n:n]
		env = append(env, string(utf16.Decode(entry)))
		blockp += 2 * (uintptr(len(entry)) + 1)
	}
	return
}

func createEnvBlock(envv []string) *uint16 {
	if len(envv) == 0 {
		return &utf16.Encode([]rune("\x00\x00"))[0]
	}
	length := 0
	for _, s := range envv {
		length += len(s) + 1
	}
	length += 1

	b := make([]byte, length)
	i := 0
	for _, s := range envv {
		l := len(s)
		copy(b[i:i+l], []byte(s))
		copy(b[i+l:i+l+1], []byte{0})
		i = i + l + 1
	}
	copy(b[i:i+1], []byte{0})

	return &utf16.Encode([]rune(string(b)))[0]
}

func makeCmdLine(args []string) string {
	var s string
	for _, v := range args {
		if s != "" {
			s += " "
		}
		s += windows.EscapeArg(v)
	}
	return s
}

func isSlash(c uint8) bool {
	return c == '\\' || c == '/'
}

func normalizeDir(dir string) (name string, err error) {
	ndir, err := syscall.FullPath(dir)
	if err != nil {
		return "", err
	}
	if len(ndir) > 2 && isSlash(ndir[0]) && isSlash(ndir[1]) {
		// dir cannot have \\server\share\path form
		return "", syscall.EINVAL
	}
	return ndir, nil
}

func volToUpper(ch int) int {
	if 'a' <= ch && ch <= 'z' {
		ch += 'A' - 'a'
	}
	return ch
}

func joinExeDirAndFName(dir, p string) (name string, err error) {
	if len(p) == 0 {
		return "", syscall.EINVAL
	}
	if len(p) > 2 && isSlash(p[0]) && isSlash(p[1]) {
		// \\server\share\path form
		return p, nil
	}
	if len(p) > 1 && p[1] == ':' {
		// has drive letter
		if len(p) == 2 {
			return "", syscall.EINVAL
		}
		if isSlash(p[2]) {
			return p, nil
		} else {
			d, err := normalizeDir(dir)
			if err != nil {
				return "", err
			}
			if volToUpper(int(p[0])) == volToUpper(int(d[0])) {
				return syscall.FullPath(d + "\\" + p[2:])
			} else {
				return syscall.FullPath(p)
			}
		}
	} else {
		// no drive letter
		d, err := normalizeDir(dir)
		if err != nil {
			return "", err
		}

		if isSlash(p[0]) {
			return windows.FullPath(d[:2] + p)
		} else {
			return windows.FullPath(d + "\\" + p)
		}
	}
}