From ac767d655b305d4e9612f5f6e33120b9176c4ad4 Mon Sep 17 00:00:00 2001
From: madiganz <zacharywmadigan@gmail.com>
Date: Sat, 7 Jul 2018 20:57:06 -0700
Subject: [PATCH] windows/svc/mgr: add ability to set a reboot message and
 command when a service fails

Added configuration options for a windows service recovery settings.
New configurations include modifying the reboot message, or command
to be run when a service fails, and getting the current reboot message
or command.

Fixes golang/go#23239

Change-Id: I3e501d66e97745b7536fd654aee2bba488083e6d
Reviewed-on: https://go-review.googlesource.com/122579
Run-TryBot: Alex Brainman <alex.brainman@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Alex Brainman <alex.brainman@gmail.com>
---
 windows/svc/mgr/mgr_test.go | 32 ++++++++++++++++++++++++++++++
 windows/svc/mgr/recovery.go | 39 +++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/windows/svc/mgr/mgr_test.go b/windows/svc/mgr/mgr_test.go
index 13f1f381..9171f5bc 100644
--- a/windows/svc/mgr/mgr_test.go
+++ b/windows/svc/mgr/mgr_test.go
@@ -174,6 +174,34 @@ func testSetRecoveryActions(t *testing.T, s *mgr.Service) {
 	testResetPeriod(t, s, 0)
 }
 
+func testRebootMessage(t *testing.T, s *mgr.Service, should string) {
+	err := s.SetRebootMessage(should)
+	if err != nil {
+		t.Fatalf("SetRebootMessage failed: %v", err)
+	}
+	is, err := s.RebootMessage()
+	if err != nil {
+		t.Fatalf("RebootMessage failed: %v", err)
+	}
+	if should != is {
+		t.Errorf("reboot message mismatch: message is %q, but should have %q", is, should)
+	}
+}
+
+func testRecoveryCommand(t *testing.T, s *mgr.Service, should string) {
+	err := s.SetRecoveryCommand(should)
+	if err != nil {
+		t.Fatalf("SetRecoveryCommand failed: %v", err)
+	}
+	is, err := s.RecoveryCommand()
+	if err != nil {
+		t.Fatalf("RecoveryCommand failed: %v", err)
+	}
+	if should != is {
+		t.Errorf("recovery command mismatch: command is %q, but should have %q", is, should)
+	}
+}
+
 func remove(t *testing.T, s *mgr.Service) {
 	err := s.Delete()
 	if err != nil {
@@ -245,6 +273,10 @@ func TestMyService(t *testing.T) {
 	}
 
 	testSetRecoveryActions(t, s)
+	testRebootMessage(t, s, "myservice failed")
+	testRebootMessage(t, s, "") // delete reboot message
+	testRecoveryCommand(t, s, "sc query myservice")
+	testRecoveryCommand(t, s, "") // delete recovery command
 
 	remove(t, s)
 }
diff --git a/windows/svc/mgr/recovery.go b/windows/svc/mgr/recovery.go
index 9243dca2..71ce2b81 100644
--- a/windows/svc/mgr/recovery.go
+++ b/windows/svc/mgr/recovery.go
@@ -8,6 +8,7 @@ package mgr
 
 import (
 	"errors"
+	"syscall"
 	"time"
 	"unsafe"
 
@@ -94,3 +95,41 @@ func (s *Service) ResetPeriod() (uint32, error) {
 	p := (*windows.SERVICE_FAILURE_ACTIONS)(unsafe.Pointer(&b[0]))
 	return p.ResetPeriod, nil
 }
+
+// SetRebootMessage sets service s reboot message.
+// If msg is "", the reboot message is deleted and no message is broadcast.
+func (s *Service) SetRebootMessage(msg string) error {
+	rActions := windows.SERVICE_FAILURE_ACTIONS{
+		RebootMsg: syscall.StringToUTF16Ptr(msg),
+	}
+	return windows.ChangeServiceConfig2(s.Handle, windows.SERVICE_CONFIG_FAILURE_ACTIONS, (*byte)(unsafe.Pointer(&rActions)))
+}
+
+// RebootMessage is broadcast to server users before rebooting in response to the ComputerReboot service controller action.
+func (s *Service) RebootMessage() (string, error) {
+	b, err := s.queryServiceConfig2(windows.SERVICE_CONFIG_FAILURE_ACTIONS)
+	if err != nil {
+		return "", err
+	}
+	p := (*windows.SERVICE_FAILURE_ACTIONS)(unsafe.Pointer(&b[0]))
+	return toString(p.RebootMsg), nil
+}
+
+// SetRecoveryCommand sets the command line of the process to execute in response to the RunCommand service controller action.
+// If cmd is "", the command is deleted and no program is run when the service fails.
+func (s *Service) SetRecoveryCommand(cmd string) error {
+	rActions := windows.SERVICE_FAILURE_ACTIONS{
+		Command: syscall.StringToUTF16Ptr(cmd),
+	}
+	return windows.ChangeServiceConfig2(s.Handle, windows.SERVICE_CONFIG_FAILURE_ACTIONS, (*byte)(unsafe.Pointer(&rActions)))
+}
+
+// RecoveryCommand is the command line of the process to execute in response to the RunCommand service controller action. This process runs under the same account as the service.
+func (s *Service) RecoveryCommand() (string, error) {
+	b, err := s.queryServiceConfig2(windows.SERVICE_CONFIG_FAILURE_ACTIONS)
+	if err != nil {
+		return "", err
+	}
+	p := (*windows.SERVICE_FAILURE_ACTIONS)(unsafe.Pointer(&b[0]))
+	return toString(p.Command), nil
+}