restartmgr_windows.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842
  1. // Copyright (c) Tailscale Inc & AUTHORS
  2. // SPDX-License-Identifier: BSD-3-Clause
  3. package winutil
  4. import (
  5. "bytes"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "os"
  10. "path/filepath"
  11. "runtime"
  12. "strings"
  13. "time"
  14. "unicode/utf16"
  15. "unsafe"
  16. "github.com/dblohm7/wingoes"
  17. "golang.org/x/sys/windows"
  18. "tailscale.com/types/logger"
  19. "tailscale.com/util/multierr"
  20. )
  21. var (
  22. // ErrDefunctProcess is returned by (*UniqueProcess).AsRestartableProcess
  23. // when the process no longer exists.
  24. ErrDefunctProcess = errors.New("process is defunct")
  25. // ErrProcessNotRestartable is returned by (*UniqueProcess).AsRestartableProcess
  26. // when the process has previously indicated that it must not be restarted
  27. // during a patch/upgrade.
  28. ErrProcessNotRestartable = errors.New("process is not restartable")
  29. )
  30. // Implementation note: the code in this file will be invoked from within
  31. // MSI custom actions, so please try to return windows.Errno error codes
  32. // whenever possible; this makes the action return more accurate errors to
  33. // the installer engine.
  34. const (
  35. _RESTART_NO_CRASH = 1
  36. _RESTART_NO_HANG = 2
  37. _RESTART_NO_PATCH = 4
  38. _RESTART_NO_REBOOT = 8
  39. )
  40. func registerForRestart(opts RegisterForRestartOpts) error {
  41. var flags uint32
  42. if !opts.RestartOnCrash {
  43. flags |= _RESTART_NO_CRASH
  44. }
  45. if !opts.RestartOnHang {
  46. flags |= _RESTART_NO_HANG
  47. }
  48. if !opts.RestartOnUpgrade {
  49. flags |= _RESTART_NO_PATCH
  50. }
  51. if !opts.RestartOnReboot {
  52. flags |= _RESTART_NO_REBOOT
  53. }
  54. var cmdLine *uint16
  55. if opts.UseCmdLineArgs {
  56. if len(opts.CmdLineArgs) == 0 {
  57. // re-use our current args, excluding the exe name itself
  58. opts.CmdLineArgs = os.Args[1:]
  59. }
  60. var b strings.Builder
  61. for _, arg := range opts.CmdLineArgs {
  62. if b.Len() > 0 {
  63. b.WriteByte(' ')
  64. }
  65. b.WriteString(windows.EscapeArg(arg))
  66. }
  67. if b.Len() > 0 {
  68. var err error
  69. cmdLine, err = windows.UTF16PtrFromString(b.String())
  70. if err != nil {
  71. return err
  72. }
  73. }
  74. }
  75. hr := registerApplicationRestart(cmdLine, flags)
  76. if e := wingoes.ErrorFromHRESULT(hr); e.Failed() {
  77. return e
  78. }
  79. return nil
  80. }
  81. type _RMHANDLE uint32
  82. // See https://web.archive.org/web/20231128212837/https://learn.microsoft.com/en-us/windows/win32/rstmgr/using-restart-manager-with-a-secondary-installer
  83. const _INVALID_RMHANDLE = ^_RMHANDLE(0)
  84. type _RM_UNIQUE_PROCESS struct {
  85. PID uint32
  86. ProcessStartTime windows.Filetime
  87. }
  88. type _RM_APP_TYPE int32
  89. const (
  90. _RmUnknownApp _RM_APP_TYPE = 0
  91. _RmMainWindow _RM_APP_TYPE = 1
  92. _RmOtherWindow _RM_APP_TYPE = 2
  93. _RmService _RM_APP_TYPE = 3
  94. _RmExplorer _RM_APP_TYPE = 4
  95. _RmConsole _RM_APP_TYPE = 5
  96. _RmCritical _RM_APP_TYPE = 1000
  97. )
  98. type _RM_APP_STATUS uint32
  99. const (
  100. //lint:ignore U1000 maps to a win32 API
  101. _RmStatusUnknown _RM_APP_STATUS = 0x0
  102. _RmStatusRunning _RM_APP_STATUS = 0x1
  103. _RmStatusStopped _RM_APP_STATUS = 0x2
  104. _RmStatusStoppedOther _RM_APP_STATUS = 0x4
  105. _RmStatusRestarted _RM_APP_STATUS = 0x8
  106. _RmStatusErrorOnStop _RM_APP_STATUS = 0x10
  107. _RmStatusErrorOnRestart _RM_APP_STATUS = 0x20
  108. _RmStatusShutdownMasked _RM_APP_STATUS = 0x40
  109. _RmStatusRestartMasked _RM_APP_STATUS = 0x80
  110. )
  111. type _RM_PROCESS_INFO struct {
  112. Process _RM_UNIQUE_PROCESS
  113. AppName [256]uint16
  114. ServiceShortName [64]uint16
  115. AppType _RM_APP_TYPE
  116. AppStatus _RM_APP_STATUS
  117. TSSessionID uint32
  118. Restartable int32 // Win32 BOOL
  119. }
  120. // RestartManagerSession represents an open Restart Manager session.
  121. type RestartManagerSession interface {
  122. io.Closer
  123. // AddPaths adds the fully-qualified paths in fqPaths to the set of binaries
  124. // that will be monitored by this restart manager session. NOTE: This
  125. // method is expensive to call, so it is better to make a single call with
  126. // a larger slice than to make multiple calls with smaller slices.
  127. AddPaths(fqPaths []string) error
  128. // AffectedProcesses returns the UniqueProcess information for all running
  129. // processes that utilize the binaries previously specified by calls to
  130. // AddPaths.
  131. AffectedProcesses() ([]UniqueProcess, error)
  132. // Key returns the session key associated with this instance.
  133. Key() string
  134. }
  135. // rmSession encapsulates the necessary information to represent an open
  136. // restart manager session.
  137. //
  138. // Implementation note: rmSession methods that return errors should use
  139. // windows.Errno codes whenever possible, as we call them from the custom
  140. // action DLL. MSI custom actions are expected to return windows.Errno values;
  141. // to ensure our compliance with this expectation, we should also use those
  142. // values. Failure to do so will result in a generic windows.Errno being
  143. // returned to the Windows Installer, which obviously is less than ideal.
  144. type rmSession struct {
  145. session _RMHANDLE
  146. key string
  147. logf logger.Logf
  148. }
  149. const _CCH_RM_SESSION_KEY = 32 // (excludes NUL terminator)
  150. // NewRestartManagerSession creates a new RestartManagerSession that utilizes
  151. // logf for logging.
  152. func NewRestartManagerSession(logf logger.Logf) (RestartManagerSession, error) {
  153. var sessionKeyBuf [_CCH_RM_SESSION_KEY + 1]uint16
  154. result := rmSession{
  155. logf: logf,
  156. }
  157. if err := rmStartSession(&result.session, 0, &sessionKeyBuf[0]); err != nil {
  158. return nil, err
  159. }
  160. result.key = windows.UTF16ToString(sessionKeyBuf[:_CCH_RM_SESSION_KEY])
  161. return &result, nil
  162. }
  163. // AttachRestartManagerSession opens a connection to an existing session
  164. // specified by sessionKey, using logf for logging.
  165. func AttachRestartManagerSession(logf logger.Logf, sessionKey string) (RestartManagerSession, error) {
  166. sessionKey16, err := windows.UTF16PtrFromString(sessionKey)
  167. if err != nil {
  168. return nil, err
  169. }
  170. result := rmSession{
  171. key: sessionKey,
  172. logf: logf,
  173. }
  174. if err := rmJoinSession(&result.session, sessionKey16); err != nil {
  175. return nil, err
  176. }
  177. return &result, nil
  178. }
  179. func (rms *rmSession) Close() error {
  180. if rms == nil || rms.session == _INVALID_RMHANDLE {
  181. return nil
  182. }
  183. if err := rmEndSession(rms.session); err != nil {
  184. return err
  185. }
  186. rms.session = _INVALID_RMHANDLE
  187. return nil
  188. }
  189. func (rms *rmSession) Key() string {
  190. return rms.key
  191. }
  192. func (rms *rmSession) AffectedProcesses() ([]UniqueProcess, error) {
  193. infos, err := rms.processList()
  194. if err != nil {
  195. return nil, err
  196. }
  197. result := make([]UniqueProcess, 0, len(infos))
  198. for _, info := range infos {
  199. result = append(result, UniqueProcess{
  200. _RM_UNIQUE_PROCESS: info.Process,
  201. CanReceiveGUIMsgs: info.AppType == _RmMainWindow || info.AppType == _RmOtherWindow,
  202. })
  203. }
  204. return result, nil
  205. }
  206. func (rms *rmSession) processList() ([]_RM_PROCESS_INFO, error) {
  207. const maxAttempts = 5
  208. var avail, rebootReasons uint32
  209. needed := uint32(1)
  210. var buf []_RM_PROCESS_INFO
  211. err := error(windows.ERROR_MORE_DATA)
  212. numAttempts := 0
  213. for err == windows.ERROR_MORE_DATA && numAttempts < maxAttempts {
  214. numAttempts++
  215. buf = make([]_RM_PROCESS_INFO, needed)
  216. avail = needed
  217. err = rmGetList(rms.session, &needed, &avail, unsafe.SliceData(buf), &rebootReasons)
  218. }
  219. if err != nil {
  220. if err == windows.ERROR_SESSION_CREDENTIAL_CONFLICT {
  221. // Add some more context about the meaning of this error.
  222. err = fmt.Errorf("%w (the Restart Manager does not permit calling RmGetList from a process that did not originally create the session)", err)
  223. }
  224. return nil, err
  225. }
  226. return buf[:avail], nil
  227. }
  228. func (rms *rmSession) AddPaths(fqPaths []string) error {
  229. if len(fqPaths) == 0 {
  230. return nil
  231. }
  232. fqPaths16 := make([]*uint16, 0, len(fqPaths))
  233. for _, fqPath := range fqPaths {
  234. if !filepath.IsAbs(fqPath) {
  235. return fmt.Errorf("%w: paths must be fully-qualified", windows.ERROR_BAD_PATHNAME)
  236. }
  237. fqPath16, err := windows.UTF16PtrFromString(fqPath)
  238. if err != nil {
  239. return err
  240. }
  241. fqPaths16 = append(fqPaths16, fqPath16)
  242. }
  243. return rmRegisterResources(rms.session, uint32(len(fqPaths16)), unsafe.SliceData(fqPaths16), 0, nil, 0, nil)
  244. }
  245. // UniqueProcess contains the necessary information to uniquely identify a
  246. // process in the face of potential PID reuse.
  247. type UniqueProcess struct {
  248. _RM_UNIQUE_PROCESS
  249. // CanReceiveGUIMsgs is true when the process has open top-level windows.
  250. CanReceiveGUIMsgs bool
  251. }
  252. // AsRestartableProcess obtains a RestartableProcess populated using the
  253. // information obtained from up.
  254. func (up *UniqueProcess) AsRestartableProcess() (*RestartableProcess, error) {
  255. // We need PROCESS_QUERY_INFORMATION instead of PROCESS_QUERY_LIMITED_INFORMATION
  256. // in order for ProcessImageName to be able to work from within a privileged
  257. // Windows Installer process.
  258. // We need PROCESS_VM_READ for GetApplicationRestartSettings.
  259. // We need PROCESS_TERMINATE and SYNCHRONIZE to terminate the process and
  260. // to be able to wait for the terminated process's handle to signal.
  261. access := uint32(windows.PROCESS_QUERY_INFORMATION | windows.PROCESS_TERMINATE | windows.PROCESS_VM_READ | windows.SYNCHRONIZE)
  262. h, err := windows.OpenProcess(access, false, up.PID)
  263. if err != nil {
  264. return nil, fmt.Errorf("OpenProcess(%d[%#X]): %w", up.PID, up.PID, err)
  265. }
  266. defer func() {
  267. if h == 0 {
  268. return
  269. }
  270. windows.CloseHandle(h)
  271. }()
  272. var creationTime, exitTime, kernelTime, userTime windows.Filetime
  273. if err := windows.GetProcessTimes(h, &creationTime, &exitTime, &kernelTime, &userTime); err != nil {
  274. return nil, fmt.Errorf("GetProcessTimes: %w", err)
  275. }
  276. if creationTime != up.ProcessStartTime {
  277. // The PID has been reused and does not actually reference the original process.
  278. return nil, ErrDefunctProcess
  279. }
  280. var tok windows.Token
  281. if err := windows.OpenProcessToken(h, windows.TOKEN_QUERY, &tok); err != nil {
  282. return nil, fmt.Errorf("OpenProcessToken: %w", err)
  283. }
  284. defer tok.Close()
  285. tsSessionID, err := TSSessionID(tok)
  286. if err != nil {
  287. return nil, fmt.Errorf("TSSessionID: %w", err)
  288. }
  289. logonSessionID, err := LogonSessionID(tok)
  290. if err != nil {
  291. return nil, fmt.Errorf("LogonSessionID: %w", err)
  292. }
  293. img, err := ProcessImageName(h)
  294. if err != nil {
  295. return nil, fmt.Errorf("ProcessImageName: %w", err)
  296. }
  297. const _RESTART_MAX_CMD_LINE = 1024
  298. var cmdLine [_RESTART_MAX_CMD_LINE]uint16
  299. cmdLineLen := uint32(len(cmdLine))
  300. var rmFlags uint32
  301. hr := getApplicationRestartSettings(h, &cmdLine[0], &cmdLineLen, &rmFlags)
  302. // Not found is not an error; it just means that the app never set any restart settings.
  303. if e := wingoes.ErrorFromHRESULT(hr); e.Failed() && e != wingoes.ErrorFromErrno(windows.ERROR_NOT_FOUND) {
  304. return nil, fmt.Errorf("GetApplicationRestartSettings: %w", error(e))
  305. }
  306. if (rmFlags & _RESTART_NO_PATCH) != 0 {
  307. // The application explicitly stated that it cannot be restarted during
  308. // an upgrade.
  309. return nil, ErrProcessNotRestartable
  310. }
  311. var logonSID string
  312. // Non-fatal, so we'll proceed with best-effort.
  313. if tokenGroups, err := tok.GetTokenGroups(); err == nil {
  314. for _, group := range tokenGroups.AllGroups() {
  315. if (group.Attributes & windows.SE_GROUP_LOGON_ID) != 0 {
  316. logonSID = group.Sid.String()
  317. break
  318. }
  319. }
  320. }
  321. var userSID string
  322. // Non-fatal, so we'll proceed with best-effort.
  323. if tokenUser, err := tok.GetTokenUser(); err == nil {
  324. // Save the user's SID so that we can later check it against the currently
  325. // logged-in Tailscale profile.
  326. userSID = tokenUser.User.Sid.String()
  327. }
  328. result := &RestartableProcess{
  329. Process: *up,
  330. SessionInfo: SessionID{
  331. LogonSession: logonSessionID,
  332. TSSession: tsSessionID,
  333. },
  334. CommandLineInfo: CommandLineInfo{
  335. ExePath: img,
  336. Args: windows.UTF16ToString(cmdLine[:cmdLineLen]),
  337. },
  338. LogonSID: logonSID,
  339. UserSID: userSID,
  340. handle: h,
  341. }
  342. runtime.SetFinalizer(result, func(rp *RestartableProcess) { rp.Close() })
  343. h = 0
  344. return result, nil
  345. }
  346. // RestartableProcess contains the necessary information to uniquely identify
  347. // an existing process, as well as the necessary information to be able to
  348. // terminate it and later start a new instance in the identical logon session
  349. // to the previous instance.
  350. type RestartableProcess struct {
  351. // Process uniquely identifies the existing process.
  352. Process UniqueProcess
  353. // SessionInfo uniquely identifies the Terminal Services (RDP) and logon
  354. // sessions the existing process is running under.
  355. SessionInfo SessionID
  356. // CommandLineInfo contains the command line information necessary for restarting.
  357. CommandLineInfo CommandLineInfo
  358. // LogonSID contains the stringified SID of the existing process's token's logon session.
  359. LogonSID string
  360. // UserSID contains the stringified SID of the existing process's token's user.
  361. UserSID string
  362. // handle specifies the Win32 HANDLE associated with the existing process.
  363. // When non-zero, it includes access rights for querying, terminating, and synchronizing.
  364. handle windows.Handle
  365. // hasExitCode is true when the exitCode field is valid.
  366. hasExitCode bool
  367. // exitCode contains exit code returned by this RestartableProcess once
  368. // its termination has been recorded by (RestartableProcesses).Terminate.
  369. // It is only valid when hasExitCode == true.
  370. exitCode uint32
  371. }
  372. func (rp *RestartableProcess) Close() error {
  373. if rp.handle == 0 {
  374. return nil
  375. }
  376. windows.CloseHandle(rp.handle)
  377. runtime.SetFinalizer(rp, nil)
  378. rp.handle = 0
  379. return nil
  380. }
  381. // RestartableProcesses is a map of PID to *RestartableProcess instance.
  382. type RestartableProcesses map[uint32]*RestartableProcess
  383. // NewRestartableProcesses instantiates a new RestartableProcesses.
  384. func NewRestartableProcesses() RestartableProcesses {
  385. return make(RestartableProcesses)
  386. }
  387. // Add inserts rp into rps.
  388. func (rps RestartableProcesses) Add(rp *RestartableProcess) {
  389. if rp != nil {
  390. rps[rp.Process.PID] = rp
  391. }
  392. }
  393. // Delete removes rp from rps.
  394. func (rps RestartableProcesses) Delete(rp *RestartableProcess) {
  395. if rp != nil {
  396. delete(rps, rp.Process.PID)
  397. }
  398. }
  399. // Close invokes (*RestartableProcess).Close on every value in rps, and then
  400. // clears rps.
  401. func (rps RestartableProcesses) Close() error {
  402. for _, v := range rps {
  403. v.Close()
  404. }
  405. clear(rps)
  406. return nil
  407. }
  408. // _MAXIMUM_WAIT_OBJECTS is the Win32 constant for the maximum number of
  409. // handles that a call to WaitForMultipleObjects may receive at once.
  410. const _MAXIMUM_WAIT_OBJECTS = 64
  411. // Terminate forcibly terminates all processes in rps using exitCode, and then
  412. // waits for their process handles to signal, up to timeout.
  413. func (rps RestartableProcesses) Terminate(logf logger.Logf, exitCode uint32, timeout time.Duration) error {
  414. if len(rps) == 0 {
  415. return nil
  416. }
  417. millis, err := wingoes.DurationToTimeoutMilliseconds(timeout)
  418. if err != nil {
  419. return err
  420. }
  421. errs := make([]error, 0, len(rps))
  422. procs := make([]*RestartableProcess, 0, len(rps))
  423. handles := make([]windows.Handle, 0, len(rps))
  424. for _, v := range rps {
  425. if err := windows.TerminateProcess(v.handle, exitCode); err != nil {
  426. if err == windows.ERROR_ACCESS_DENIED {
  427. // If v terminated before we attempted to terminate, we'll receive
  428. // ERROR_ACCESS_DENIED, which is not really an error worth reporting in
  429. // our use case. Just obtain the exit code and then close the process.
  430. if err := windows.GetExitCodeProcess(v.handle, &v.exitCode); err != nil {
  431. logf("GetExitCodeProcess failed: %v", err)
  432. } else {
  433. v.hasExitCode = true
  434. }
  435. v.Close()
  436. } else {
  437. errs = append(errs, &terminationError{rp: v, err: err})
  438. }
  439. continue
  440. }
  441. procs = append(procs, v)
  442. handles = append(handles, v.handle)
  443. }
  444. for len(handles) > 0 {
  445. // WaitForMultipleObjects can only wait on _MAXIMUM_WAIT_OBJECTS handles per
  446. // call, so we batch them as necessary.
  447. count := uint32(min(len(handles), _MAXIMUM_WAIT_OBJECTS))
  448. waitCode, err := windows.WaitForMultipleObjects(handles[:count], true, millis)
  449. if err != nil {
  450. errs = append(errs, fmt.Errorf("waiting on terminated process handles: %w", err))
  451. break
  452. }
  453. if e := windows.Errno(waitCode); e == windows.WAIT_TIMEOUT {
  454. errs = append(errs, fmt.Errorf("waiting on terminated process handles: %w", error(e)))
  455. break
  456. }
  457. if waitCode >= windows.WAIT_OBJECT_0 && waitCode < (windows.WAIT_OBJECT_0+count) {
  458. // The first count process handles have all been signaled. Close them out.
  459. for _, proc := range procs[:count] {
  460. if err := windows.GetExitCodeProcess(proc.handle, &proc.exitCode); err != nil {
  461. logf("GetExitCodeProcess failed: %v", err)
  462. } else {
  463. proc.hasExitCode = true
  464. }
  465. proc.Close()
  466. }
  467. procs = procs[count:]
  468. handles = handles[count:]
  469. continue
  470. }
  471. // We really shouldn't be reaching this point
  472. panic(fmt.Sprintf("unexpected state from WaitForMultipleObjects: %d", waitCode))
  473. }
  474. if len(errs) != 0 {
  475. return multierr.New(errs...)
  476. }
  477. return nil
  478. }
  479. type terminationError struct {
  480. rp *RestartableProcess
  481. err error
  482. }
  483. func (te *terminationError) Error() string {
  484. pid := te.rp.Process.PID
  485. return fmt.Sprintf("terminating process %d (%#X): %v", pid, pid, te.err)
  486. }
  487. func (te *terminationError) Unwrap() error {
  488. return te.err
  489. }
  490. // SessionID encapsulates the necessary information for uniquely identifying
  491. // sessions. In particular, SessionID contains enough information to detect
  492. // reuse of Terminal Service session IDs.
  493. type SessionID struct {
  494. // LogonSession is the NT logon session ID.
  495. LogonSession windows.LUID
  496. // TSSession is the terminal services session ID.
  497. TSSession uint32
  498. }
  499. // OpenToken obtains the security token associated with sessID.
  500. func (sessID *SessionID) OpenToken() (windows.Token, error) {
  501. var token windows.Token
  502. if err := windows.WTSQueryUserToken(sessID.TSSession, &token); err != nil {
  503. return 0, err
  504. }
  505. var err error
  506. defer func() {
  507. if err != nil {
  508. token.Close()
  509. }
  510. }()
  511. tokenLogonSession, err := LogonSessionID(token)
  512. if err != nil {
  513. return 0, err
  514. }
  515. if tokenLogonSession != sessID.LogonSession {
  516. err = windows.ERROR_NO_SUCH_LOGON_SESSION
  517. return 0, err
  518. }
  519. return token, nil
  520. }
  521. // ContainsToken determines whether token is contained within sessID.
  522. func (sessID *SessionID) ContainsToken(token windows.Token) (bool, error) {
  523. tokenTSSessionID, err := TSSessionID(token)
  524. if err != nil {
  525. return false, err
  526. }
  527. if tokenTSSessionID != sessID.TSSession {
  528. return false, nil
  529. }
  530. tokenLogonSession, err := LogonSessionID(token)
  531. if err != nil {
  532. return false, err
  533. }
  534. return tokenLogonSession == sessID.LogonSession, nil
  535. }
  536. // This is the Window Station and Desktop within a particular session that must
  537. // be specified for interactive processes: "Winsta0\\default\x00"
  538. var defaultDesktop = unsafe.SliceData([]uint16{'W', 'i', 'n', 's', 't', 'a', '0', '\\', 'd', 'e', 'f', 'a', 'u', 'l', 't', 0})
  539. // CommandLineInfo manages the necessary information for creating a Win32
  540. // process using a specific command line.
  541. type CommandLineInfo struct {
  542. // ExePath must be a fully-qualified path to a Windows executable binary.
  543. ExePath string
  544. // Args must be any arguments supplied to the process, excluding the
  545. // path to the binary itself. Args must be properly quoted according to
  546. // Windows path rules. To create a properly quoted Args from scratch, call the
  547. // SetArgs method instead.
  548. Args string `json:",omitempty"`
  549. }
  550. // SetArgs converts args to a string quoted as necessary to satisfy the rules
  551. // for Win32 command lines, and sets cli.Args to that string.
  552. func (cli *CommandLineInfo) SetArgs(args []string) {
  553. var buf strings.Builder
  554. for _, arg := range args {
  555. if buf.Len() > 0 {
  556. buf.WriteByte(' ')
  557. }
  558. buf.WriteString(windows.EscapeArg(arg))
  559. }
  560. cli.Args = buf.String()
  561. }
  562. // Validate ensures that cli.ExePath contains an absolute path.
  563. func (cli *CommandLineInfo) Validate() error {
  564. if cli == nil {
  565. return windows.ERROR_INVALID_PARAMETER
  566. }
  567. if !filepath.IsAbs(cli.ExePath) {
  568. return fmt.Errorf("%w: CommandLineInfo requires absolute ExePath", windows.ERROR_BAD_PATHNAME)
  569. }
  570. return nil
  571. }
  572. // Resolve converts the information in cli to a format compatible with the Win32
  573. // CreateProcess* family of APIs, as pointers to C-style UTF-16 strings. It also
  574. // returns the full command line as a Go string for logging purposes.
  575. func (cli *CommandLineInfo) Resolve() (exePath *uint16, cmdLine *uint16, cmdLineStr string, err error) {
  576. // Resolve cmdLine first since that also does a Validate.
  577. cmdLineStr, cmdLine, err = cli.resolveArgsAsUTF16Ptr()
  578. if err != nil {
  579. return nil, nil, "", err
  580. }
  581. exePath, err = windows.UTF16PtrFromString(cli.ExePath)
  582. if err != nil {
  583. return nil, nil, "", err
  584. }
  585. return exePath, cmdLine, cmdLineStr, nil
  586. }
  587. // resolveArgs quotes cli.ExePath as necessary, appends Args, and returns the result.
  588. func (cli *CommandLineInfo) resolveArgs() (string, error) {
  589. if err := cli.Validate(); err != nil {
  590. return "", err
  591. }
  592. var cmdLineBuf strings.Builder
  593. cmdLineBuf.WriteString(windows.EscapeArg(cli.ExePath))
  594. if args := cli.Args; args != "" {
  595. cmdLineBuf.WriteByte(' ')
  596. cmdLineBuf.WriteString(args)
  597. }
  598. return cmdLineBuf.String(), nil
  599. }
  600. func (cli *CommandLineInfo) resolveArgsAsUTF16Ptr() (string, *uint16, error) {
  601. s, err := cli.resolveArgs()
  602. if err != nil {
  603. return "", nil, err
  604. }
  605. s16, err := windows.UTF16PtrFromString(s)
  606. if err != nil {
  607. return "", nil, err
  608. }
  609. return s, s16, nil
  610. }
  611. // StartProcessInSession creates a new process using cmdLineInfo that will
  612. // reside inside the session identified by sessID, with the security token whose
  613. // logon is associated with sessID. The child process's environment will be
  614. // inherited from the session token's environment.
  615. func StartProcessInSession(sessID SessionID, cmdLineInfo CommandLineInfo) error {
  616. return StartProcessInSessionWithHandler(sessID, cmdLineInfo, nil)
  617. }
  618. // PostCreateProcessHandler is a function that is invoked by
  619. // StartProcessInSessionWithHandler when the child process has been successfully
  620. // created. It is the responsibility of the handler to close the pi.Thread and
  621. // pi.Process handles.
  622. type PostCreateProcessHandler func(pi *windows.ProcessInformation)
  623. // StartProcessInSessionWithHandler creates a new process using cmdLineInfo that
  624. // will reside inside the session identified by sessID, with the security token
  625. // whose logon is associated with sessID. The child process's environment will be
  626. // inherited from the session token's environment. When the child process has
  627. // been successfully created, handler is invoked with the windows.ProcessInformation
  628. // that was returned by the OS.
  629. func StartProcessInSessionWithHandler(sessID SessionID, cmdLineInfo CommandLineInfo, handler PostCreateProcessHandler) error {
  630. pi, err := startProcessInSessionInternal(sessID, cmdLineInfo, 0)
  631. if err != nil {
  632. return err
  633. }
  634. if handler != nil {
  635. handler(pi)
  636. return nil
  637. }
  638. windows.CloseHandle(pi.Process)
  639. windows.CloseHandle(pi.Thread)
  640. return nil
  641. }
  642. // RunProcessInSession creates a new process and waits up to timeout for that
  643. // child process to complete its execution. The process is created using
  644. // cmdLineInfo and will reside inside the session identified by sessID, with the
  645. // security token whose logon is associated with sessID. The child process's
  646. // environment will be inherited from the session token's environment.
  647. func RunProcessInSession(sessID SessionID, cmdLineInfo CommandLineInfo, timeout time.Duration) (uint32, error) {
  648. timeoutMillis, err := wingoes.DurationToTimeoutMilliseconds(timeout)
  649. if err != nil {
  650. return 1, err
  651. }
  652. pi, err := startProcessInSessionInternal(sessID, cmdLineInfo, 0)
  653. if err != nil {
  654. return 1, err
  655. }
  656. windows.CloseHandle(pi.Thread)
  657. defer windows.CloseHandle(pi.Process)
  658. waitCode, err := windows.WaitForSingleObject(pi.Process, timeoutMillis)
  659. if err != nil {
  660. return 1, fmt.Errorf("WaitForSingleObject: %w", err)
  661. }
  662. if e := windows.Errno(waitCode); e == windows.WAIT_TIMEOUT {
  663. return 1, e
  664. }
  665. if waitCode != windows.WAIT_OBJECT_0 {
  666. // This should not be possible; log
  667. return 1, fmt.Errorf("unexpected state from WaitForSingleObject: %d", waitCode)
  668. }
  669. var exitCode uint32
  670. if err := windows.GetExitCodeProcess(pi.Process, &exitCode); err != nil {
  671. return 1, err
  672. }
  673. return exitCode, nil
  674. }
  675. func startProcessInSessionInternal(sessID SessionID, cmdLineInfo CommandLineInfo, extraFlags uint32) (*windows.ProcessInformation, error) {
  676. if err := cmdLineInfo.Validate(); err != nil {
  677. return nil, err
  678. }
  679. token, err := sessID.OpenToken()
  680. if err != nil {
  681. return nil, fmt.Errorf("(*SessionID).OpenToken: %w", err)
  682. }
  683. defer token.Close()
  684. exePath16, commandLine16, _, err := cmdLineInfo.Resolve()
  685. if err != nil {
  686. return nil, fmt.Errorf("(*CommandLineInfo).Resolve(): %w", err)
  687. }
  688. wd16, err := windows.UTF16PtrFromString(filepath.Dir(cmdLineInfo.ExePath))
  689. if err != nil {
  690. return nil, fmt.Errorf("UTF16PtrFromString(wd): %w", err)
  691. }
  692. env, err := token.Environ(false)
  693. if err != nil {
  694. return nil, fmt.Errorf("token environment: %w", err)
  695. }
  696. env16 := newEnvBlock(env)
  697. // The privileges in privNames are required for CreateProcessAsUser to be
  698. // able to start processes as other users in other logon sessions.
  699. privNames := []string{
  700. "SeAssignPrimaryTokenPrivilege",
  701. "SeIncreaseQuotaPrivilege",
  702. }
  703. dropPrivs, err := EnableCurrentThreadPrivileges(privNames)
  704. if err != nil {
  705. return nil, fmt.Errorf("EnableCurrentThreadPrivileges(%#v): %w", privNames, err)
  706. }
  707. defer dropPrivs()
  708. createFlags := extraFlags | windows.CREATE_UNICODE_ENVIRONMENT | windows.DETACHED_PROCESS
  709. si := windows.StartupInfo{
  710. Cb: uint32(unsafe.Sizeof(windows.StartupInfo{})),
  711. Desktop: defaultDesktop,
  712. }
  713. var pi windows.ProcessInformation
  714. if err := windows.CreateProcessAsUser(token, exePath16, commandLine16, nil, nil,
  715. false, createFlags, env16, wd16, &si, &pi); err != nil {
  716. return nil, fmt.Errorf("CreateProcessAsUser: %w", err)
  717. }
  718. return &pi, nil
  719. }
  720. func newEnvBlock(env []string) *uint16 {
  721. // Intentionally using bytes.Buffer here because we're writing nul bytes (the standard library does this too).
  722. var buf bytes.Buffer
  723. for _, v := range env {
  724. buf.WriteString(v)
  725. buf.WriteByte(0)
  726. }
  727. if buf.Len() == 0 {
  728. // So that we end with a double-null in the empty env case
  729. buf.WriteByte(0)
  730. }
  731. buf.WriteByte(0)
  732. return unsafe.SliceData(utf16.Encode([]rune(string(buf.Bytes()))))
  733. }