// Copyright (c) Tailscale Inc & AUTHORS
// SPDX-License-Identifier: BSD-3-Clause

package winutil

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"
	"unicode/utf16"
	"unsafe"

	"github.com/dblohm7/wingoes"
	"golang.org/x/sys/windows"
	"tailscale.com/types/logger"
	"tailscale.com/util/multierr"
)

var (
	// ErrDefunctProcess is returned by (*UniqueProcess).AsRestartableProcess
	// when the process no longer exists.
	ErrDefunctProcess = errors.New("process is defunct")
	// ErrProcessNotRestartable is returned by (*UniqueProcess).AsRestartableProcess
	// when the process has previously indicated that it must not be restarted
	// during a patch/upgrade.
	ErrProcessNotRestartable = errors.New("process is not restartable")
)

// Implementation note: the code in this file will be invoked from within
// MSI custom actions, so please try to return windows.Errno error codes
// whenever possible; this makes the action return more accurate errors to
// the installer engine.

const (
	_RESTART_NO_CRASH  = 1
	_RESTART_NO_HANG   = 2
	_RESTART_NO_PATCH  = 4
	_RESTART_NO_REBOOT = 8
)

func registerForRestart(opts RegisterForRestartOpts) error {
	var flags uint32

	if !opts.RestartOnCrash {
		flags |= _RESTART_NO_CRASH
	}
	if !opts.RestartOnHang {
		flags |= _RESTART_NO_HANG
	}
	if !opts.RestartOnUpgrade {
		flags |= _RESTART_NO_PATCH
	}
	if !opts.RestartOnReboot {
		flags |= _RESTART_NO_REBOOT
	}

	var cmdLine *uint16
	if opts.UseCmdLineArgs {
		if len(opts.CmdLineArgs) == 0 {
			// re-use our current args, excluding the exe name itself
			opts.CmdLineArgs = os.Args[1:]
		}

		var b strings.Builder
		for _, arg := range opts.CmdLineArgs {
			if b.Len() > 0 {
				b.WriteByte(' ')
			}
			b.WriteString(windows.EscapeArg(arg))
		}

		if b.Len() > 0 {
			var err error
			cmdLine, err = windows.UTF16PtrFromString(b.String())
			if err != nil {
				return err
			}
		}
	}

	hr := registerApplicationRestart(cmdLine, flags)
	if e := wingoes.ErrorFromHRESULT(hr); e.Failed() {
		return e
	}

	return nil
}

type _RMHANDLE uint32

// See https://web.archive.org/web/20231128212837/https://learn.microsoft.com/en-us/windows/win32/rstmgr/using-restart-manager-with-a-secondary-installer
const _INVALID_RMHANDLE = ^_RMHANDLE(0)

type _RM_UNIQUE_PROCESS struct {
	PID              uint32
	ProcessStartTime windows.Filetime
}

type _RM_APP_TYPE int32

const (
	_RmUnknownApp  _RM_APP_TYPE = 0
	_RmMainWindow  _RM_APP_TYPE = 1
	_RmOtherWindow _RM_APP_TYPE = 2
	_RmService     _RM_APP_TYPE = 3
	_RmExplorer    _RM_APP_TYPE = 4
	_RmConsole     _RM_APP_TYPE = 5
	_RmCritical    _RM_APP_TYPE = 1000
)

type _RM_APP_STATUS uint32

const (
	//lint:ignore U1000 maps to a win32 API
	_RmStatusUnknown        _RM_APP_STATUS = 0x0
	_RmStatusRunning        _RM_APP_STATUS = 0x1
	_RmStatusStopped        _RM_APP_STATUS = 0x2
	_RmStatusStoppedOther   _RM_APP_STATUS = 0x4
	_RmStatusRestarted      _RM_APP_STATUS = 0x8
	_RmStatusErrorOnStop    _RM_APP_STATUS = 0x10
	_RmStatusErrorOnRestart _RM_APP_STATUS = 0x20
	_RmStatusShutdownMasked _RM_APP_STATUS = 0x40
	_RmStatusRestartMasked  _RM_APP_STATUS = 0x80
)

type _RM_PROCESS_INFO struct {
	Process          _RM_UNIQUE_PROCESS
	AppName          [256]uint16
	ServiceShortName [64]uint16
	AppType          _RM_APP_TYPE
	AppStatus        _RM_APP_STATUS
	TSSessionID      uint32
	Restartable      int32 // Win32 BOOL
}

// RestartManagerSession represents an open Restart Manager session.
type RestartManagerSession interface {
	io.Closer
	// AddPaths adds the fully-qualified paths in fqPaths to the set of binaries
	// that will be monitored by this restart manager session. NOTE: This
	// method is expensive to call, so it is better to make a single call with
	// a larger slice than to make multiple calls with smaller slices.
	AddPaths(fqPaths []string) error
	// AffectedProcesses returns the UniqueProcess information for all running
	// processes that utilize the binaries previously specified by calls to
	// AddPaths.
	AffectedProcesses() ([]UniqueProcess, error)
	// Key returns the session key associated with this instance.
	Key() string
}

// rmSession encapsulates the necessary information to represent an open
// restart manager session.
//
// Implementation note: rmSession methods that return errors should use
// windows.Errno codes whenever possible, as we call them from the custom
// action DLL. MSI custom actions are expected to return windows.Errno values;
// to ensure our compliance with this expectation, we should also use those
// values. Failure to do so will result in a generic windows.Errno being
// returned to the Windows Installer, which obviously is less than ideal.
type rmSession struct {
	session _RMHANDLE
	key     string
	logf    logger.Logf
}

const _CCH_RM_SESSION_KEY = 32 // (excludes NUL terminator)

// NewRestartManagerSession creates a new RestartManagerSession that utilizes
// logf for logging.
func NewRestartManagerSession(logf logger.Logf) (RestartManagerSession, error) {
	var sessionKeyBuf [_CCH_RM_SESSION_KEY + 1]uint16
	result := rmSession{
		logf: logf,
	}
	if err := rmStartSession(&result.session, 0, &sessionKeyBuf[0]); err != nil {
		return nil, err
	}

	result.key = windows.UTF16ToString(sessionKeyBuf[:_CCH_RM_SESSION_KEY])
	return &result, nil
}

// AttachRestartManagerSession opens a connection to an existing session
// specified by sessionKey, using logf for logging.
func AttachRestartManagerSession(logf logger.Logf, sessionKey string) (RestartManagerSession, error) {
	sessionKey16, err := windows.UTF16PtrFromString(sessionKey)
	if err != nil {
		return nil, err
	}

	result := rmSession{
		key:  sessionKey,
		logf: logf,
	}
	if err := rmJoinSession(&result.session, sessionKey16); err != nil {
		return nil, err
	}
	return &result, nil
}

func (rms *rmSession) Close() error {
	if rms == nil || rms.session == _INVALID_RMHANDLE {
		return nil
	}
	if err := rmEndSession(rms.session); err != nil {
		return err
	}
	rms.session = _INVALID_RMHANDLE
	return nil
}

func (rms *rmSession) Key() string {
	return rms.key
}

func (rms *rmSession) AffectedProcesses() ([]UniqueProcess, error) {
	infos, err := rms.processList()
	if err != nil {
		return nil, err
	}

	result := make([]UniqueProcess, 0, len(infos))
	for _, info := range infos {
		result = append(result, UniqueProcess{
			_RM_UNIQUE_PROCESS: info.Process,
			CanReceiveGUIMsgs:  info.AppType == _RmMainWindow || info.AppType == _RmOtherWindow,
		})
	}

	return result, nil
}

func (rms *rmSession) processList() ([]_RM_PROCESS_INFO, error) {
	const maxAttempts = 5
	var avail, rebootReasons uint32
	needed := uint32(1)

	var buf []_RM_PROCESS_INFO
	err := error(windows.ERROR_MORE_DATA)
	numAttempts := 0
	for err == windows.ERROR_MORE_DATA && numAttempts < maxAttempts {
		numAttempts++
		buf = make([]_RM_PROCESS_INFO, needed)
		avail = needed
		err = rmGetList(rms.session, &needed, &avail, unsafe.SliceData(buf), &rebootReasons)
	}

	if err != nil {
		if err == windows.ERROR_SESSION_CREDENTIAL_CONFLICT {
			// Add some more context about the meaning of this error.
			err = fmt.Errorf("%w (the Restart Manager does not permit calling RmGetList from a process that did not originally create the session)", err)
		}
		return nil, err
	}

	return buf[:avail], nil
}

func (rms *rmSession) AddPaths(fqPaths []string) error {
	if len(fqPaths) == 0 {
		return nil
	}

	fqPaths16 := make([]*uint16, 0, len(fqPaths))
	for _, fqPath := range fqPaths {
		if !filepath.IsAbs(fqPath) {
			return fmt.Errorf("%w: paths must be fully-qualified", windows.ERROR_BAD_PATHNAME)
		}

		fqPath16, err := windows.UTF16PtrFromString(fqPath)
		if err != nil {
			return err
		}

		fqPaths16 = append(fqPaths16, fqPath16)
	}

	return rmRegisterResources(rms.session, uint32(len(fqPaths16)), unsafe.SliceData(fqPaths16), 0, nil, 0, nil)
}

// UniqueProcess contains the necessary information to uniquely identify a
// process in the face of potential PID reuse.
type UniqueProcess struct {
	_RM_UNIQUE_PROCESS
	// CanReceiveGUIMsgs is true when the process has open top-level windows.
	CanReceiveGUIMsgs bool
}

// AsRestartableProcess obtains a RestartableProcess populated using the
// information obtained from up.
func (up *UniqueProcess) AsRestartableProcess() (*RestartableProcess, error) {
	// We need PROCESS_QUERY_INFORMATION instead of PROCESS_QUERY_LIMITED_INFORMATION
	// in order for ProcessImageName to be able to work from within a privileged
	// Windows Installer process.
	// We need PROCESS_VM_READ for GetApplicationRestartSettings.
	// We need PROCESS_TERMINATE and SYNCHRONIZE to terminate the process and
	// to be able to wait for the terminated process's handle to signal.
	access := uint32(windows.PROCESS_QUERY_INFORMATION | windows.PROCESS_TERMINATE | windows.PROCESS_VM_READ | windows.SYNCHRONIZE)
	h, err := windows.OpenProcess(access, false, up.PID)
	if err != nil {
		return nil, fmt.Errorf("OpenProcess(%d[%#X]): %w", up.PID, up.PID, err)
	}
	defer func() {
		if h == 0 {
			return
		}
		windows.CloseHandle(h)
	}()

	var creationTime, exitTime, kernelTime, userTime windows.Filetime
	if err := windows.GetProcessTimes(h, &creationTime, &exitTime, &kernelTime, &userTime); err != nil {
		return nil, fmt.Errorf("GetProcessTimes: %w", err)
	}
	if creationTime != up.ProcessStartTime {
		// The PID has been reused and does not actually reference the original process.
		return nil, ErrDefunctProcess
	}

	var tok windows.Token
	if err := windows.OpenProcessToken(h, windows.TOKEN_QUERY, &tok); err != nil {
		return nil, fmt.Errorf("OpenProcessToken: %w", err)
	}
	defer tok.Close()

	tsSessionID, err := TSSessionID(tok)
	if err != nil {
		return nil, fmt.Errorf("TSSessionID: %w", err)
	}

	logonSessionID, err := LogonSessionID(tok)
	if err != nil {
		return nil, fmt.Errorf("LogonSessionID: %w", err)
	}

	img, err := ProcessImageName(h)
	if err != nil {
		return nil, fmt.Errorf("ProcessImageName: %w", err)
	}

	const _RESTART_MAX_CMD_LINE = 1024
	var cmdLine [_RESTART_MAX_CMD_LINE]uint16
	cmdLineLen := uint32(len(cmdLine))
	var rmFlags uint32
	hr := getApplicationRestartSettings(h, &cmdLine[0], &cmdLineLen, &rmFlags)
	// Not found is not an error; it just means that the app never set any restart settings.
	if e := wingoes.ErrorFromHRESULT(hr); e.Failed() && e != wingoes.ErrorFromErrno(windows.ERROR_NOT_FOUND) {
		return nil, fmt.Errorf("GetApplicationRestartSettings: %w", error(e))
	}
	if (rmFlags & _RESTART_NO_PATCH) != 0 {
		// The application explicitly stated that it cannot be restarted during
		// an upgrade.
		return nil, ErrProcessNotRestartable
	}

	var logonSID string
	// Non-fatal, so we'll proceed with best-effort.
	if tokenGroups, err := tok.GetTokenGroups(); err == nil {
		for _, group := range tokenGroups.AllGroups() {
			if (group.Attributes & windows.SE_GROUP_LOGON_ID) != 0 {
				logonSID = group.Sid.String()
				break
			}
		}
	}

	var userSID string
	// Non-fatal, so we'll proceed with best-effort.
	if tokenUser, err := tok.GetTokenUser(); err == nil {
		// Save the user's SID so that we can later check it against the currently
		// logged-in Tailscale profile.
		userSID = tokenUser.User.Sid.String()
	}

	result := &RestartableProcess{
		Process: *up,
		SessionInfo: SessionID{
			LogonSession: logonSessionID,
			TSSession:    tsSessionID,
		},
		CommandLineInfo: CommandLineInfo{
			ExePath: img,
			Args:    windows.UTF16ToString(cmdLine[:cmdLineLen]),
		},
		LogonSID: logonSID,
		UserSID:  userSID,
		handle:   h,
	}

	runtime.SetFinalizer(result, func(rp *RestartableProcess) { rp.Close() })
	h = 0
	return result, nil
}

// RestartableProcess contains the necessary information to uniquely identify
// an existing process, as well as the necessary information to be able to
// terminate it and later start a new instance in the identical logon session
// to the previous instance.
type RestartableProcess struct {
	// Process uniquely identifies the existing process.
	Process UniqueProcess
	// SessionInfo uniquely identifies the Terminal Services (RDP) and logon
	// sessions the existing process is running under.
	SessionInfo SessionID
	// CommandLineInfo contains the command line information necessary for restarting.
	CommandLineInfo CommandLineInfo
	// LogonSID contains the stringified SID of the existing process's token's logon session.
	LogonSID string
	// UserSID contains the stringified SID of the existing process's token's user.
	UserSID string
	// handle specifies the Win32 HANDLE associated with the existing process.
	// When non-zero, it includes access rights for querying, terminating, and synchronizing.
	handle windows.Handle
	// hasExitCode is true when the exitCode field is valid.
	hasExitCode bool
	// exitCode contains exit code returned by this RestartableProcess once
	// its termination has been recorded by (RestartableProcesses).Terminate.
	// It is only valid when hasExitCode == true.
	exitCode uint32
}

func (rp *RestartableProcess) Close() error {
	if rp.handle == 0 {
		return nil
	}
	windows.CloseHandle(rp.handle)
	runtime.SetFinalizer(rp, nil)
	rp.handle = 0
	return nil
}

// RestartableProcesses is a map of PID to *RestartableProcess instance.
type RestartableProcesses map[uint32]*RestartableProcess

// NewRestartableProcesses instantiates a new RestartableProcesses.
func NewRestartableProcesses() RestartableProcesses {
	return make(RestartableProcesses)
}

// Add inserts rp into rps.
func (rps RestartableProcesses) Add(rp *RestartableProcess) {
	if rp != nil {
		rps[rp.Process.PID] = rp
	}
}

// Delete removes rp from rps.
func (rps RestartableProcesses) Delete(rp *RestartableProcess) {
	if rp != nil {
		delete(rps, rp.Process.PID)
	}
}

// Close invokes (*RestartableProcess).Close on every value in rps, and then
// clears rps.
func (rps RestartableProcesses) Close() error {
	for _, v := range rps {
		v.Close()
	}
	clear(rps)
	return nil
}

// _MAXIMUM_WAIT_OBJECTS is the Win32 constant for the maximum number of
// handles that a call to WaitForMultipleObjects may receive at once.
const _MAXIMUM_WAIT_OBJECTS = 64

// Terminate forcibly terminates all processes in rps using exitCode, and then
// waits for their process handles to signal, up to timeout.
func (rps RestartableProcesses) Terminate(logf logger.Logf, exitCode uint32, timeout time.Duration) error {
	if len(rps) == 0 {
		return nil
	}

	millis, err := wingoes.DurationToTimeoutMilliseconds(timeout)
	if err != nil {
		return err
	}

	errs := make([]error, 0, len(rps))
	procs := make([]*RestartableProcess, 0, len(rps))
	handles := make([]windows.Handle, 0, len(rps))
	for _, v := range rps {
		if err := windows.TerminateProcess(v.handle, exitCode); err != nil {
			if err == windows.ERROR_ACCESS_DENIED {
				// If v terminated before we attempted to terminate, we'll receive
				// ERROR_ACCESS_DENIED, which is not really an error worth reporting in
				// our use case. Just obtain the exit code and then close the process.
				if err := windows.GetExitCodeProcess(v.handle, &v.exitCode); err != nil {
					logf("GetExitCodeProcess failed: %v", err)
				} else {
					v.hasExitCode = true
				}
				v.Close()
			} else {
				errs = append(errs, &terminationError{rp: v, err: err})
			}
			continue
		}
		procs = append(procs, v)
		handles = append(handles, v.handle)
	}

	for len(handles) > 0 {
		// WaitForMultipleObjects can only wait on _MAXIMUM_WAIT_OBJECTS handles per
		// call, so we batch them as necessary.
		count := uint32(min(len(handles), _MAXIMUM_WAIT_OBJECTS))
		waitCode, err := windows.WaitForMultipleObjects(handles[:count], true, millis)
		if err != nil {
			errs = append(errs, fmt.Errorf("waiting on terminated process handles: %w", err))
			break
		}
		if e := windows.Errno(waitCode); e == windows.WAIT_TIMEOUT {
			errs = append(errs, fmt.Errorf("waiting on terminated process handles: %w", error(e)))
			break
		}
		if waitCode >= windows.WAIT_OBJECT_0 && waitCode < (windows.WAIT_OBJECT_0+count) {
			// The first count process handles have all been signaled. Close them out.
			for _, proc := range procs[:count] {
				if err := windows.GetExitCodeProcess(proc.handle, &proc.exitCode); err != nil {
					logf("GetExitCodeProcess failed: %v", err)
				} else {
					proc.hasExitCode = true
				}
				proc.Close()
			}
			procs = procs[count:]
			handles = handles[count:]
			continue
		}
		// We really shouldn't be reaching this point
		panic(fmt.Sprintf("unexpected state from WaitForMultipleObjects: %d", waitCode))
	}

	if len(errs) != 0 {
		return multierr.New(errs...)
	}
	return nil
}

type terminationError struct {
	rp  *RestartableProcess
	err error
}

func (te *terminationError) Error() string {
	pid := te.rp.Process.PID
	return fmt.Sprintf("terminating process %d (%#X): %v", pid, pid, te.err)
}

func (te *terminationError) Unwrap() error {
	return te.err
}

// SessionID encapsulates the necessary information for uniquely identifying
// sessions. In particular, SessionID contains enough information to detect
// reuse of Terminal Service session IDs.
type SessionID struct {
	// LogonSession is the NT logon session ID.
	LogonSession windows.LUID
	// TSSession is the terminal services session ID.
	TSSession uint32
}

// OpenToken obtains the security token associated with sessID.
func (sessID *SessionID) OpenToken() (windows.Token, error) {
	var token windows.Token
	if err := windows.WTSQueryUserToken(sessID.TSSession, &token); err != nil {
		return 0, err
	}

	var err error
	defer func() {
		if err != nil {
			token.Close()
		}
	}()

	tokenLogonSession, err := LogonSessionID(token)
	if err != nil {
		return 0, err
	}

	if tokenLogonSession != sessID.LogonSession {
		err = windows.ERROR_NO_SUCH_LOGON_SESSION
		return 0, err
	}

	return token, nil
}

// ContainsToken determines whether token is contained within sessID.
func (sessID *SessionID) ContainsToken(token windows.Token) (bool, error) {
	tokenTSSessionID, err := TSSessionID(token)
	if err != nil {
		return false, err
	}

	if tokenTSSessionID != sessID.TSSession {
		return false, nil
	}

	tokenLogonSession, err := LogonSessionID(token)
	if err != nil {
		return false, err
	}

	return tokenLogonSession == sessID.LogonSession, nil
}

// This is the Window Station and Desktop within a particular session that must
// be specified for interactive processes: "Winsta0\\default\x00"
var defaultDesktop = unsafe.SliceData([]uint16{'W', 'i', 'n', 's', 't', 'a', '0', '\\', 'd', 'e', 'f', 'a', 'u', 'l', 't', 0})

// CommandLineInfo manages the necessary information for creating a Win32
// process using a specific command line.
type CommandLineInfo struct {
	// ExePath must be a fully-qualified path to a Windows executable binary.
	ExePath string
	// Args must be any arguments supplied to the process, excluding the
	// path to the binary itself. Args must be properly quoted according to
	// Windows path rules. To create a properly quoted Args from scratch, call the
	// SetArgs method instead.
	Args string `json:",omitempty"`
}

// SetArgs converts args to a string quoted as necessary to satisfy the rules
// for Win32 command lines, and sets cli.Args to that string.
func (cli *CommandLineInfo) SetArgs(args []string) {
	var buf strings.Builder
	for _, arg := range args {
		if buf.Len() > 0 {
			buf.WriteByte(' ')
		}
		buf.WriteString(windows.EscapeArg(arg))
	}

	cli.Args = buf.String()
}

// Validate ensures that cli.ExePath contains an absolute path.
func (cli *CommandLineInfo) Validate() error {
	if cli == nil {
		return windows.ERROR_INVALID_PARAMETER
	}

	if !filepath.IsAbs(cli.ExePath) {
		return fmt.Errorf("%w: CommandLineInfo requires absolute ExePath", windows.ERROR_BAD_PATHNAME)
	}

	return nil
}

// Resolve converts the information in cli to a format compatible with the Win32
// CreateProcess* family of APIs, as pointers to C-style UTF-16 strings. It also
// returns the full command line as a Go string for logging purposes.
func (cli *CommandLineInfo) Resolve() (exePath *uint16, cmdLine *uint16, cmdLineStr string, err error) {
	// Resolve cmdLine first since that also does a Validate.
	cmdLineStr, cmdLine, err = cli.resolveArgsAsUTF16Ptr()
	if err != nil {
		return nil, nil, "", err
	}

	exePath, err = windows.UTF16PtrFromString(cli.ExePath)
	if err != nil {
		return nil, nil, "", err
	}

	return exePath, cmdLine, cmdLineStr, nil
}

// resolveArgs quotes cli.ExePath as necessary, appends Args, and returns the result.
func (cli *CommandLineInfo) resolveArgs() (string, error) {
	if err := cli.Validate(); err != nil {
		return "", err
	}

	var cmdLineBuf strings.Builder
	cmdLineBuf.WriteString(windows.EscapeArg(cli.ExePath))
	if args := cli.Args; args != "" {
		cmdLineBuf.WriteByte(' ')
		cmdLineBuf.WriteString(args)
	}

	return cmdLineBuf.String(), nil
}

func (cli *CommandLineInfo) resolveArgsAsUTF16Ptr() (string, *uint16, error) {
	s, err := cli.resolveArgs()
	if err != nil {
		return "", nil, err
	}
	s16, err := windows.UTF16PtrFromString(s)
	if err != nil {
		return "", nil, err
	}
	return s, s16, nil
}

// StartProcessInSession creates a new process using cmdLineInfo that will
// reside inside the session identified by sessID, with the security token whose
// logon is associated with sessID. The child process's environment will be
// inherited from the session token's environment.
func StartProcessInSession(sessID SessionID, cmdLineInfo CommandLineInfo) error {
	return StartProcessInSessionWithHandler(sessID, cmdLineInfo, nil)
}

// PostCreateProcessHandler is a function that is invoked by
// StartProcessInSessionWithHandler when the child process has been successfully
// created. It is the responsibility of the handler to close the pi.Thread and
// pi.Process handles.
type PostCreateProcessHandler func(pi *windows.ProcessInformation)

// StartProcessInSessionWithHandler creates a new process using cmdLineInfo that
// will reside inside the session identified by sessID, with the security token
// whose logon is associated with sessID. The child process's environment will be
// inherited from the session token's environment. When the child process has
// been successfully created, handler is invoked with the windows.ProcessInformation
// that was returned by the OS.
func StartProcessInSessionWithHandler(sessID SessionID, cmdLineInfo CommandLineInfo, handler PostCreateProcessHandler) error {
	pi, err := startProcessInSessionInternal(sessID, cmdLineInfo, 0)
	if err != nil {
		return err
	}
	if handler != nil {
		handler(pi)
		return nil
	}
	windows.CloseHandle(pi.Process)
	windows.CloseHandle(pi.Thread)
	return nil
}

// RunProcessInSession creates a new process and waits up to timeout for that
// child process to complete its execution. The process is created using
// cmdLineInfo and will reside inside the session identified by sessID, with the
// security token whose logon is associated with sessID. The child process's
// environment will be inherited from the session token's environment.
func RunProcessInSession(sessID SessionID, cmdLineInfo CommandLineInfo, timeout time.Duration) (uint32, error) {
	timeoutMillis, err := wingoes.DurationToTimeoutMilliseconds(timeout)
	if err != nil {
		return 1, err
	}

	pi, err := startProcessInSessionInternal(sessID, cmdLineInfo, 0)
	if err != nil {
		return 1, err
	}
	windows.CloseHandle(pi.Thread)
	defer windows.CloseHandle(pi.Process)

	waitCode, err := windows.WaitForSingleObject(pi.Process, timeoutMillis)
	if err != nil {
		return 1, fmt.Errorf("WaitForSingleObject: %w", err)
	}
	if e := windows.Errno(waitCode); e == windows.WAIT_TIMEOUT {
		return 1, e
	}
	if waitCode != windows.WAIT_OBJECT_0 {
		// This should not be possible; log
		return 1, fmt.Errorf("unexpected state from WaitForSingleObject: %d", waitCode)
	}

	var exitCode uint32
	if err := windows.GetExitCodeProcess(pi.Process, &exitCode); err != nil {
		return 1, err
	}
	return exitCode, nil
}

func startProcessInSessionInternal(sessID SessionID, cmdLineInfo CommandLineInfo, extraFlags uint32) (*windows.ProcessInformation, error) {
	if err := cmdLineInfo.Validate(); err != nil {
		return nil, err
	}

	token, err := sessID.OpenToken()
	if err != nil {
		return nil, fmt.Errorf("(*SessionID).OpenToken: %w", err)
	}
	defer token.Close()

	exePath16, commandLine16, _, err := cmdLineInfo.Resolve()
	if err != nil {
		return nil, fmt.Errorf("(*CommandLineInfo).Resolve(): %w", err)
	}

	wd16, err := windows.UTF16PtrFromString(filepath.Dir(cmdLineInfo.ExePath))
	if err != nil {
		return nil, fmt.Errorf("UTF16PtrFromString(wd): %w", err)
	}

	env, err := token.Environ(false)
	if err != nil {
		return nil, fmt.Errorf("token environment: %w", err)
	}
	env16 := newEnvBlock(env)

	// The privileges in privNames are required for CreateProcessAsUser to be
	// able to start processes as other users in other logon sessions.
	privNames := []string{
		"SeAssignPrimaryTokenPrivilege",
		"SeIncreaseQuotaPrivilege",
	}
	dropPrivs, err := EnableCurrentThreadPrivileges(privNames)
	if err != nil {
		return nil, fmt.Errorf("EnableCurrentThreadPrivileges(%#v): %w", privNames, err)
	}
	defer dropPrivs()

	createFlags := extraFlags | windows.CREATE_UNICODE_ENVIRONMENT | windows.DETACHED_PROCESS
	si := windows.StartupInfo{
		Cb:      uint32(unsafe.Sizeof(windows.StartupInfo{})),
		Desktop: defaultDesktop,
	}
	var pi windows.ProcessInformation
	if err := windows.CreateProcessAsUser(token, exePath16, commandLine16, nil, nil,
		false, createFlags, env16, wd16, &si, &pi); err != nil {
		return nil, fmt.Errorf("CreateProcessAsUser: %w", err)
	}
	return &pi, nil
}

func newEnvBlock(env []string) *uint16 {
	// Intentionally using bytes.Buffer here because we're writing nul bytes (the standard library does this too).
	var buf bytes.Buffer
	for _, v := range env {
		buf.WriteString(v)
		buf.WriteByte(0)
	}
	if buf.Len() == 0 {
		// So that we end with a double-null in the empty env case
		buf.WriteByte(0)
	}
	buf.WriteByte(0)
	return unsafe.SliceData(utf16.Encode([]rune(string(buf.Bytes()))))
}