wincmd/svc.go

package wincmd

import (
	"errors"
	"fmt"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/eventlog"
	"golang.org/x/sys/windows/svc/mgr"
	"time"
)

type SvcStatus svc.State

const (
	Stopped         = SvcStatus(svc.Stopped)
	StartPending    = SvcStatus(svc.StartPending)
	StopPending     = SvcStatus(svc.StopPending)
	Running         = SvcStatus(svc.Running)
	ContinuePending = SvcStatus(svc.ContinuePending)
	PausePending    = SvcStatus(svc.PausePending)
	Paused          = SvcStatus(svc.Paused)

	StartManual    = windows.SERVICE_DEMAND_START // the service must be started manually
	StartAutomatic = windows.SERVICE_AUTO_START   // the service will start by itself whenever the computer reboots
	StartDisabled  = windows.SERVICE_DISABLED     // the service cannot be started

	// The severity of the error, and action taken,
	// if this service fails to start.
	ErrorCritical = windows.SERVICE_ERROR_CRITICAL
	ErrorIgnore   = windows.SERVICE_ERROR_IGNORE
	ErrorNormal   = windows.SERVICE_ERROR_NORMAL
	ErrorSevere   = windows.SERVICE_ERROR_SEVERE
)

type WinSvcExecute struct {
	Run         func()
	Stop        func()
	Interrupt   func()
	Pause       func()
	Continue    func()
	OtherMethod func(svc.Cmd)
	Name        string
	Accepted    []svc.Accepted
}

type WinSvcInput struct {
	Name             string
	DisplayName      string
	ExecPath         string
	DelayedAutoStart bool
	Description      string
	StartType        uint32
	Args             []string
}

type WinSvc struct {
	*mgr.Service
}

func IsServiceExists(name string) (bool, error) {
	if !Isas() {
		return false, errors.New("permission deny")
	}
	winmgr, err := mgr.Connect()
	if err != nil {
		return false, err
	}
	defer winmgr.Disconnect()
	lists, err := winmgr.ListServices()
	if err != nil {
		return false, err
	}
	for _, v := range lists {
		if name == v {
			return true, nil
		}
	}
	return false, nil
}

func CreateService(mysvc WinSvcInput) (*WinSvc, error) {
	if !Isas() {
		return nil, errors.New("permission deny")
	}
	if exists, err := IsServiceExists(mysvc.Name); err != nil {
		return nil, err
	} else if exists {
		return nil, errors.New("service already exists")
	}
	winmgr, err := mgr.Connect()
	if err != nil {
		return nil, err
	}
	defer winmgr.Disconnect()
	mycfg := mgr.Config{
		DisplayName:      mysvc.DisplayName,
		StartType:        mysvc.StartType,
		DelayedAutoStart: mysvc.DelayedAutoStart,
		Description:      mysvc.Description,
	}
	gsvc, err := winmgr.CreateService(mysvc.Name, mysvc.ExecPath, mycfg, mysvc.Args...)
	if err != nil {
		return nil, err
	}
	err = eventlog.InstallAsEventCreate(mysvc.Name, eventlog.Error|eventlog.Warning|eventlog.Info)
	if err != nil {
		gsvc.Delete()
		return nil, fmt.Errorf("winsvc.InstallService: InstallAsEventCreate failed, err = %v", err)
	}
	var result WinSvc
	result.Service = gsvc
	return &result, nil
}

func OpenService(name string) (*WinSvc, error) {
	if !Isas() {
		return nil, errors.New("permission deny")
	}
	if exists, err := IsServiceExists(name); err != nil {
		return nil, err
	} else if !exists {
		return nil, errors.New("service not exists")
	}
	winmgr, err := mgr.Connect()
	if err != nil {
		return nil, err
	}
	defer winmgr.Disconnect()
	gsvc, err := winmgr.OpenService(name)
	if err != nil {
		return nil, err
	}
	var result WinSvc
	result.Service = gsvc
	return &result, nil
}

func DeleteService(name string) error {
	mysvc, err := OpenService(name)
	if err != nil {
		return err
	}
	err = mysvc.Service.Delete()
	if err != nil {
		mysvc.Close()
		return err
	}
	mysvc.Close()
	err = eventlog.Remove(name)
	if err != nil {
		return err
	}
	var count int
	for {
		if ok, err := IsServiceExists(name); err != nil {
			return err
		} else if !ok {
			return nil
		}
		time.Sleep(time.Millisecond * 300)
		count++
		if count > 100 {
			return errors.New("timeout")
		}
	}
}

func StopService(name string) error {
	mysvc, err := OpenService(name)
	if err != nil {
		return err
	}
	defer mysvc.Close()
	_, err = mysvc.Service.Control(svc.Stop)
	if err != nil {
		return err
	}
	var count int
	for {
		status, err := mysvc.Service.Query()
		if err != nil {
			return err
		}
		if status.State == svc.Stopped {
			return nil
		}
		time.Sleep(time.Millisecond * 100)
		count++
		if count > 100 {
			return errors.New("timeout")
		}
	}
}

func StartService(name string) error {
	mysvc, err := OpenService(name)
	if err != nil {
		return err
	}
	defer mysvc.Close()
	err = mysvc.Service.Start()
	if err != nil {
		return err
	}
	var count int
	for {
		status, err := mysvc.Service.Query()
		if err != nil {
			return err
		}
		if status.State == svc.Running {
			return nil
		}
		time.Sleep(time.Millisecond * 100)
		count++
		if count > 100 {
			return errors.New("timeout")
		}
	}
}

func ServiceStatus(name string) (SvcStatus, error) {
	mysvc, err := OpenService(name)
	if err != nil {
		return Stopped, err
	}
	defer mysvc.Close()
	status, err := mysvc.Service.Query()
	return SvcStatus(status.State), err
}

func InService() (bool, error) {
	if !Isas() {
		return false, nil
	}
	return svc.IsWindowsService()
}

func (w *WinSvc) Stop() error {
	return StopService(w.Name)
}

func (w *WinSvc) Delete() error {
	if err := w.Close(); err != nil {
		return err
	}
	return DeleteService(w.Name)
}

func (w *WinSvc) StartService() error {
	err := w.Service.Start()
	if err != nil {
		return err
	}
	var count int
	for {
		sts, err := w.Query()
		if err != nil {
			return err
		}
		if SvcStatus(sts.State) == Running {
			return nil
		}
		time.Sleep(time.Millisecond * 100)
		count++
		if count > 100 {
			return errors.New("timeout")
		}
	}
}

func InServiceBool() bool {
	ok, _ := svc.IsWindowsService()
	return ok
}

func (w *WinSvcExecute) Execute(args []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (svcSpecificEC bool, exitCode uint32) {
	var sva svc.Accepted
	alreadyStoped := make(chan int)
	for _, v := range w.Accepted {
		sva = sva | v
	}
	s <- svc.Status{State: svc.StartPending}
	go func() {
		s <- svc.Status{State: svc.Running, Accepts: sva}
		w.Run()
		alreadyStoped <- 1
	}()

	for {
		select {
		case <-alreadyStoped:
			return
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				s <- c.CurrentStatus
				w.Interrupt()
				s <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s <- svc.Status{State: svc.StopPending}
				w.Stop()
				s <- svc.Status{State: svc.Stopped}
				return
			case svc.Pause:
				s <- svc.Status{State: svc.PausePending, Accepts: sva}
				if w.Pause != nil {
					w.Pause()
				}
				s <- svc.Status{State: svc.Paused, Accepts: sva}
			case svc.Continue:
				s <- svc.Status{State: svc.ContinuePending, Accepts: sva}
				if w.Continue != nil {
					w.Continue()
				}
				s <- svc.Status{State: svc.Running, Accepts: sva}
			default:
				if w.OtherMethod != nil {
					w.OtherMethod(c.Cmd)
				}
				s <- svc.Status{State: svc.Running, Accepts: sva}
			}
		}
	}
}

func NewWinSvcExecute(name string, run, stop func()) *WinSvcExecute {
	var res WinSvcExecute
	res.Run = run
	res.Stop = stop
	res.Interrupt = func() {
		time.Sleep(time.Millisecond)
	}
	res.Accepted = []svc.Accepted{svc.AcceptStop, svc.AcceptShutdown, svc.AcceptPauseAndContinue}
	return &res
}

func (w *WinSvcExecute) StartService() error {
	return svc.Run(w.Name, w)
}

func (w *WinSvcExecute) InService() (bool, error) {
	if !Isas() {
		return false, nil
	}
	return svc.IsWindowsService()
}

func (w *WinSvcExecute) InServiceBool() bool {
	ok, _ := svc.IsWindowsService()
	return ok
}