Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 38 additions & 3 deletions os/windows.go
Original file line number Diff line number Diff line change
Expand Up @@ -177,11 +177,46 @@ func (c Windows) CommandExist(h Host, cmd string) bool {
return h.Execf("where /q %s", cmd) == nil
}

// Reboot executes the reboot command
// Reboot triggers an immediate forced restart by scheduling a SYSTEM-context
// one-shot task that runs 'shutdown /r /f /t 5', then immediately triggering
// and deleting it within the 5-second countdown window.
//
// Running via a scheduled task bypasses the filtered Administrator token used
// by WinRM sessions (e.g. AWS EC2) which lacks SeShutdownPrivilege. Issuing
// 'shutdown /r' directly in the WinRM session is silently ignored in that
// context.
//
// /sc onstart is used instead of /sc once to avoid schtasks writing a
// stderr warning about the start time being in the past, which rig treats
// as an error. The task is deleted immediately after triggering (while the
// 5-second timer counts down) so it does not re-fire on subsequent startups.
Comment on lines +180 to +192
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment calls the scheduled task "one-shot", but the implementation uses schtasks /sc onstart, which will run on every startup unless the task is successfully deleted. Either adjust the wording to reflect that it's an ONSTART task deleted as part of the reboot flow, or switch to a truly one-shot schedule (e.g., /sc once) if feasible.

Copilot uses AI. Check for mistakes.
func (c Windows) Reboot(h Host) error {
if err := h.Exec("shutdown /r /t 5"); err != nil {
return fmt.Errorf("failed to reboot: %w", err)
const taskName = "RigReboot"
// Create a SYSTEM-context ONSTART task that runs 'shutdown /r /f /t 5'.
// The 5-second delay gives us time to delete the task before the OS
// actually executes the reboot, preventing it from firing again on the
// next startup.
create := fmt.Sprintf(`schtasks /create /tn "%s" /tr "shutdown /r /f /t 5" /sc onstart /f /ru SYSTEM`, taskName)
if err := h.Exec(create); err != nil {
return fmt.Errorf("failed to create reboot task: %w", err)
Comment on lines +194 to +201
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using a fixed scheduled task name ("RigReboot") with /f can overwrite an existing task with the same name and may conflict with concurrent reboot attempts. Consider generating a unique task name per invocation (e.g., include a timestamp/UUID) and/or checking for an existing task before overwriting/deleting.

Copilot uses AI. Check for mistakes.
}
run := fmt.Sprintf(`schtasks /run /tn "%s"`, taskName)
if err := h.Exec(run); err != nil {
// Tolerate connection-level errors; the OS may kill WinRM as it starts
// rebooting before the run command returns.
errMsg := err.Error()
if !strings.Contains(errMsg, "connection") && !strings.Contains(errMsg, "closed") && !strings.Contains(errMsg, "EOF") {
return fmt.Errorf("failed to run reboot task: %w", err)
}
}
// Delete the task immediately while the 5-second shutdown timer is still
// counting down. This prevents it from re-firing on subsequent startups.
del := fmt.Sprintf(`schtasks /delete /tn "%s" /f`, taskName)
// Best-effort: ignore delete errors — if the task fires before we can
// delete it, the caller is expected to delete it after reconnecting.
_ = h.Exec(del)
Comment on lines +212 to +217
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Delete failures are fully ignored, but with /sc onstart a leftover task will reboot the machine on every subsequent startup (persistent reboot loop). This is a high-risk failure mode. Prefer making the task self-cleaning (e.g., include a delete/disable step in the task action) and/or use a schedule that cannot retrigger on startup (e.g., /sc once with stderr tolerated via exec.AllowWinStderr()), and only ignore delete errors that are clearly caused by the ongoing reboot/connection drop.

Copilot uses AI. Check for mistakes.
// Allow Windows time to complete shutdown before waitForHost begins polling.
time.Sleep(15 * time.Second)
return nil
Comment on lines +218 to 220
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

time.Sleep(15 * time.Second) is a hard-coded delay that is decoupled from the configured shutdown timer (/t 5). Consider deriving this from the shutdown delay (plus a small buffer) or documenting why 15s is required, to avoid unnecessary slowdowns and make future changes safer.

Copilot uses AI. Check for mistakes.
}

Expand Down
39 changes: 6 additions & 33 deletions winrm.go
Original file line number Diff line number Diff line change
Expand Up @@ -197,27 +197,18 @@ type Command struct {
}

// Wait blocks until the command finishes
func (c *Command) Wait() (err error) { //nolint:nonamedreturns // needed for panic recovery
defer func() {
if r := recover(); err == nil && r != nil {
if strings.Contains(fmt.Sprint(r), "close of closed channel") {
log.Debugf("recovered from a panic in Command.Wait: %v", r)
} else {
panic(r)
}
}
}()

// Wait blocks until the command finishes
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The doc comment for Command.Wait is duplicated ("Wait blocks until the command finishes" appears twice). Please remove the extra line to keep godoc clean.

Suggested change
// Wait blocks until the command finishes

Copilot uses AI. Check for mistakes.
func (c *Command) Wait() error {
defer c.sh.Close()
defer c.cmd.Close()

c.wg.Wait()
c.cmd.Wait()
log.Debugf("command finished")
if c.cmd.ExitCode() != 0 {
err = fmt.Errorf("%w: exit code %d", ErrCommandFailed, c.cmd.ExitCode())
return fmt.Errorf("%w: exit code %d", ErrCommandFailed, c.cmd.ExitCode())
}
return err
return nil
}

// Close terminates the command
Expand Down Expand Up @@ -319,58 +310,40 @@ func (c *WinRM) Exec(cmd string, opts ...exec.Option) error { //nolint:cyclop
}()
}

var errors []string

wg.Add(1)
go func() {
// ignore channel close panics
defer func() {
if r := recover(); r != nil {
log.Debugf("recovered from a panic while reading stderr: %v", r)
}
}()
defer wg.Done()
if execOpts.Writer == nil {
outputScanner := bufio.NewScanner(command.Stdout)

for outputScanner.Scan() {
execOpts.AddOutput(c.String(), outputScanner.Text()+"\n", "")
}

if err := outputScanner.Err(); err != nil {
execOpts.LogErrorf("%s: %s", c, err.Error())
}
command.Stdout.Close()
} else {
if _, err := io.Copy(execOpts.Writer, command.Stdout); err != nil {
execOpts.LogErrorf("%s: failed to stream stdout: %v", c, err)
}
}
}()

var errors []string

wg.Add(1)
go func() {
// ignore channel close panics
defer func() {
if r := recover(); r != nil {
log.Debugf("recovered from a panic while reading stderr: %v", r)
}
}()
defer wg.Done()
outputScanner := bufio.NewScanner(command.Stderr)

for outputScanner.Scan() {
msg := outputScanner.Text()
if msg != "" {
errors = append(errors, msg)
execOpts.LogErrorf("%s: %s", c, msg)
}
}

if err := outputScanner.Err(); err != nil {
execOpts.LogErrorf("%s: %s", c, err.Error())
}
command.Stderr.Close()
}()

wg.Wait()
Expand Down
Loading