gitee.com/h79/goutils@v1.22.10/common/svc/winsvc.go

//go:build windows

package svc

import (
	"flag"
	"fmt"
	"gitee.com/h79/goutils/common/logger"
	"gitee.com/h79/goutils/common/system"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/mgr"
	"os"
	"path/filepath"
	"time"
)

const (
	kServiceName    = "serviceName"
	kServiceDesc    = "serviceDesc"
	kServiceInstall = "serviceInstall"
	kServiceRemove  = "serviceRemove"
	kServiceStart   = "serviceStart"
	kServiceReStart = "serviceReStart"
	kServiceStop    = "serviceStop"
)

var (
	appPath     string
	serviceName string
	serviceDesc string

	flagServiceInstall = flag.Bool(kServiceInstall, false, "Install service")
	flagServiceRemove  = flag.Bool(kServiceRemove, false, "Remove service")
	flagServiceStart   = flag.Bool(kServiceStart, false, "Start service")
	flagServiceReStart = flag.Bool(kServiceReStart, false, "ReStart service")
	flagServiceStop    = flag.Bool(kServiceStop, false, "Stop service")
)

func init() {

	flag.StringVar(&serviceName, kServiceName, "", "Set service name")
	flag.StringVar(&serviceDesc, kServiceDesc, "", "Set service description")

	// change to current dir
	var err error
	if appPath, err = getExePath(); err != nil {
		logger.Fatal("err= %v", err)
	}
	if err = os.Chdir(filepath.Dir(appPath)); err != nil {
		logger.Fatal("err= %v", err)
	}
}

func DoService(conf Config, startService func(isWindowsService bool), stopService func()) {
	if len(serviceName) <= 0 {
		serviceName = conf.ServiceName
	}
	if len(serviceDesc) <= 0 {
		serviceDesc = conf.ServiceDesc
	}
	// install service
	if *flagServiceInstall {
		if len(serviceName) <= 0 || len(serviceDesc) <= 0 {
			logger.Error("failed to install service: serviceName OR serviceDesc is empty")
			os.Exit(-2)
		}
		args := os.Args[1:]
		args = system.StripArgs(args, "-"+kServiceName)
		args = system.StripArgs(args, "-"+kServiceDesc)
		args = system.StripArgs(args, "-"+kServiceInstall)
		if err := Install(appPath, serviceName, serviceDesc, args...); err != nil {
			logger.Error("failed to install service: %v, err: %v", serviceName, err)
		}
		logger.Info("done")
		os.Exit(0)
	}

	// uninstall service
	if *flagServiceRemove {
		if len(serviceName) <= 0 {
			logger.Error("failed to uninstall service: serviceName is empty")
			os.Exit(-2)
		}
		if err := Remove(serviceName); err != nil {
			logger.Error("failed to uninstall service: %v, err: %v", serviceName, err)
		} else {
			logger.Info("done")
		}
		os.Exit(0)
	}

	// start service
	if *flagServiceStart {
		if len(serviceName) <= 0 {
			logger.Error("failed to start service: serviceName is empty")
			os.Exit(-2)
		}
		if err := Start(serviceName); err != nil {
			logger.Error("failed to start service: %v, err: %v", serviceName, err)
		} else {
			logger.Info("done")
		}
		os.Exit(0)
	}

	// restart service
	if *flagServiceReStart {
		if len(serviceName) <= 0 {
			logger.Error("failed to restart service: serviceName is empty")
			os.Exit(-2)
		}
		if err := ReStart(serviceName); err != nil {
			logger.Error("failed to restart service: %v, err: %v", serviceName, err)
		} else {
			logger.Info("done")
		}
		os.Exit(0)
	}

	// stop service
	if *flagServiceStop {
		if len(serviceName) <= 0 {
			logger.Error("failed to stop service: serviceName is empty")
			os.Exit(-2)
		}
		if err := Stop(serviceName); err != nil {
			logger.Error("failed to stop service: %v, err: %v", serviceName, err)
		} else {
			logger.Info("done")
		}
		os.Exit(0)
	}

	if !IsWindowsService() {
		startService(false)
		return
	}

	if len(serviceName) <= 0 {
		logger.Error("failed to Run service: serviceName is empty")
		os.Exit(-2)
	}
	// run as service
	if err := RunAs(serviceName, startService, stopService, false); err != nil {
		logger.Error("Run failure, err= %v", err)
	}
}

func getExePath() (string, error) {
	prog := os.Args[0]
	p, err := filepath.Abs(prog)
	if err != nil {
		return "", err
	}
	fi, err := os.Stat(p)
	if err == nil {
		if !fi.Mode().IsDir() {
			return p, nil
		}
		err = fmt.Errorf("GetAppPath: %s is directory", p)
	}
	if filepath.Ext(p) == "" {
		p += ".exe"
		fi, err = os.Stat(p)
		if err == nil {
			if !fi.Mode().IsDir() {
				return p, nil
			}
			err = fmt.Errorf("GetAppPath: %s is directory", p)
		}
	}
	return "", err
}

func IsWindowsService() bool {
	ok, err := svc.IsWindowsService()
	if err != nil {
		fmt.Println("IsWindowsService:  err = ", err)
	}
	return ok
}

func Install(appPath, name, desc string, params ...string) error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()
	s, err := m.OpenService(name)
	if err == nil {
		s.Close()
		return fmt.Errorf("InstallService: service %s already exists", name)
	}
	s, err = m.CreateService(name, appPath,
		mgr.Config{
			DisplayName: desc,
			Description: desc,
			StartType:   windows.SERVICE_AUTO_START,
		},
		params...,
	)
	if err != nil {
		return err
	}
	defer s.Close()
	//_ = eventlog.InstallAsEventCreate(name, eventlog.Error|eventlog.Warning|eventlog.Info)
	return nil
}

func Remove(name string) error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()
	s, err := m.OpenService(name)
	if err != nil {
		return fmt.Errorf("RemoveService: service %s is not installed", name)
	}
	defer s.Close()
	err = s.Delete()
	if err != nil {
		return err
	}
	//_ = eventlog.Remove(name)
	return nil
}

func Start(name string) error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()
	s, err := m.OpenService(name)
	if err != nil {
		return fmt.Errorf("StartService: could not access service: %v", err)
	}
	defer s.Close()
	err = s.Start("p1", "p2", "p3")
	if err != nil {
		return fmt.Errorf("StartService: could not start service: %v", err)
	}
	return nil
}

func ReStart(name string) error {
	if err := Control(name, svc.Stop, svc.Stopped); err != nil {
		return err
	}
	return Start(name)
}

func Stop(name string) error {
	if err := Control(name, svc.Stop, svc.Stopped); err != nil {
		return err
	}
	return nil
}

func Query(name string) (status string, err error) {
	m, err := mgr.Connect()
	if err != nil {
		return
	}
	defer m.Disconnect()
	s, err := m.OpenService(name)
	if err != nil {
		err = fmt.Errorf("Query: could not access service: %v", err)
		return
	}
	defer s.Close()

	statusCode, err := s.Query()
	if err != nil {
		return
	}
	switch statusCode.State {
	case svc.Stopped:
		return "Stopped", nil
	case svc.StartPending:
		return "StartPending", nil
	case svc.StopPending:
		return "StopPending", nil
	case svc.Running:
		return "Running", nil
	case svc.ContinuePending:
		return "ContinuePending", nil
	case svc.PausePending:
		return "PausePending", nil
	case svc.Paused:
		return "Paused", nil
	}
	panic("unreached")
}

func RunAs(name string, start func(isWindowsService bool), stop func(), isDebug bool) (err error) {

	run := svc.Run
	if isDebug {
		run = debug.Run
	}

	fmt.Println("RunAs: starting service: ", name)
	if err = run(name, &winService{Start: start, Stop: stop}); err != nil {
		fmt.Println(name, "service failed: ", err)
		return
	}
	fmt.Println("RunAs:", name, "service stopped")
	return
}

func Control(name string, c svc.Cmd, to svc.State) error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()
	s, err := m.OpenService(name)
	if err != nil {
		return fmt.Errorf("Control: could not access service: %v", err)
	}
	defer s.Close()
	status, err := s.Control(c)
	if err != nil {
		return fmt.Errorf("Control: could not send control=%d: %v", c, err)
	}
	timeout := time.Now().Add(10 * time.Second)
	for status.State != to {
		if timeout.Before(time.Now()) {
			return fmt.Errorf("Control: timeout waiting for service to go to state=%d", to)
		}
		time.Sleep(300 * time.Millisecond)
		status, err = s.Query()
		if err != nil {
			return fmt.Errorf("Control: could not retrieve service status: %v", err)
		}
	}
	return nil
}

type winService struct {
	Start func(isWindowsService bool)
	Stop  func()
}

func (p *winService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	logger.Info("Execute: begin")
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown | svc.AcceptPauseAndContinue
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}

	var exit = make(chan struct{})
	system.ChildRunning(func() {
		p.Start(true)
		close(exit)
	})
	var stop = false
loop:
	for {
		select {
		case <-exit:
			stop = true
			break
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
				// testing deadlock from https://code.google.com/p/winsvc/issues/detail?id=4
				time.Sleep(100 * time.Millisecond)
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				break loop
			case svc.Pause:
				changes <- svc.Status{State: svc.Paused, Accepts: cmdsAccepted}
			case svc.Continue:
				changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
			default:
				logger.Error("Execute:: unexpected control request #%d", c)
			}
		}
		if stop {
			break
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	p.Stop()

	logger.Info("Execute: end")
	return
}