maskshell · August 29, 2017 06:07 · May 12, 2015
diff --git a/winsvc.go b/winsvc.go
@@ -0,0 +1,303 @@
+// Copyright 2015 <chaishushan{AT}gmail.com>. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build ingore
+
+package main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"path/filepath"
+	"time"
+
+	"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"
+)
+
+var (
+	flagServiceName = flag.String("service-name", "winsvc-demo", "Set service name")
+	flagServiceDesc = flag.String("service-desc", "Winsvc Demo Server", "Set service description")
+
+	flagServiceInstall   = flag.Bool("service-install", false, "Install service")
+	flagServiceUninstall = flag.Bool("service-remove", false, "Remove service")
+	flagServiceStart     = flag.Bool("service-start", false, "Start service")
+	flagServiceStop      = flag.Bool("service-stop", false, "Stop service")
+
+	flagHelp = flag.Bool("help", false, "Show usage and exit.")
+)
+
+func init() {
+	flag.Usage = func() {
+		fmt.Fprintf(os.Stderr, `Usage:
+  demo [options]...
+
+Options:
+`)
+		flag.PrintDefaults()
+		fmt.Fprintf(os.Stderr, "%s\n", `
+Example:
+  # run demo server
+  $ go build -o demo.exe demo.go
+  $ demo.exe
+
+  # install demo as windows service
+  $ demo.exe -service-install
+
+  # start/stop demo service
+  $ demo.exe -service-start
+  $ demo.exe -service-stop
+
+  # remove demo service
+  $ demo.exe -service-remove
+
+  # help
+  $ demo.exe -h
+
+Report bugs to <chaishushan{AT}gmail.com>.`)
+	}
+
+	// change to current dir
+	exepath, err := exePath()
+	if err != nil {
+		log.Fatal(err)
+	}
+	if err := os.Chdir(filepath.Dir(exepath)); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func main() {
+	flag.Parse()
+
+	// install service
+	if *flagServiceInstall {
+		if err := installService(*flagServiceName, *flagServiceDesc); err != nil {
+			log.Fatalf("installService(%s, %s): %v\n", *flagServiceName, *flagServiceDesc, err)
+		}
+		fmt.Printf("Done\n")
+		return
+	}
+
+	// remove service
+	if *flagServiceUninstall {
+		if err := removeService(*flagServiceName); err != nil {
+			log.Fatalln("removeService:", err)
+		}
+		fmt.Printf("Done\n")
+		return
+	}
+
+	// start service
+	if *flagServiceStart {
+		if err := startService(*flagServiceName); err != nil {
+			log.Fatalln("startService:", err)
+		}
+		fmt.Printf("Done\n")
+		return
+	}
+
+	// stop service
+	if *flagServiceStop {
+		if err := controlService(*flagServiceName, svc.Stop, svc.Stopped); err != nil {
+			log.Fatalln("stopService:", err)
+		}
+		fmt.Printf("Done\n")
+		return
+	}
+
+	// run as service
+	isIntSess, err := svc.IsAnInteractiveSession()
+	if err != nil {
+		log.Fatalf("svc.IsAnInteractiveSession: %v\n", err)
+	}
+	if !isIntSess {
+		log.Println("main:", "runService")
+		if err := svc.Run(*flagServiceName, &demoService{}); err != nil {
+			log.Fatalf("svc.Run: %v\n", err)
+		}
+		return
+	}
+
+	// run as normal
+	StartServer()
+}
+
+func StartServer() {
+	log.Println("StartServer, port = 8080")
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, "winsrv server", time.Now())
+	})
+	http.ListenAndServe(":8080", nil)
+}
+
+func StopServer() {
+	log.Println("StopServer")
+}
+
+type demoService struct{}
+
+func (m *demoService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
+	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown | svc.AcceptPauseAndContinue
+	changes <- svc.Status{State: svc.StartPending}
+	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
+	go StartServer()
+
+loop:
+	for {
+		select {
+		case c := <-r:
+			switch c.Cmd {
+			case svc.Interrogate:
+				changes <- c.CurrentStatus
+				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:
+				log.Printf("unexpected control request #%d", c)
+			}
+		}
+	}
+
+	log.Println("demoService.Execute:", "svc.StopPending")
+	changes <- svc.Status{State: svc.StopPending}
+	StopServer()
+	return
+}
+
+func installService(name, desc string) error {
+	exepath, err := exePath()
+	if err != nil {
+		return err
+	}
+	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("service %s already exists", name)
+	}
+	s, err = m.CreateService(name, exepath, mgr.Config{
+		DisplayName: desc,
+		StartType:   windows.SERVICE_AUTO_START,
+	})
+	if err != nil {
+		return err
+	}
+	defer s.Close()
+	err = eventlog.InstallAsEventCreate(name, eventlog.Error|eventlog.Warning|eventlog.Info)
+	if err != nil {
+		s.Delete()
+		return fmt.Errorf("SetupEventLogSource() failed: %s", err)
+	}
+	return nil
+}
+
+func removeService(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("service %s is not installed", name)
+	}
+	defer s.Close()
+	err = s.Delete()
+	if err != nil {
+		return err
+	}
+	err = eventlog.Remove(name)
+	if err != nil {
+		return fmt.Errorf("RemoveEventLogSource() failed: %s", err)
+	}
+	return nil
+}
+
+func startService(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("could not access service: %v", err)
+	}
+	defer s.Close()
+	err = s.Start("p1", "p2", "p3")
+	if err != nil {
+		return fmt.Errorf("could not start service: %v", err)
+	}
+	return nil
+}
+
+func controlService(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("could not access service: %v", err)
+	}
+	defer s.Close()
+	status, err := s.Control(c)
+	if err != nil {
+		return fmt.Errorf("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("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("could not retrieve service status: %v", err)
+		}
+	}
+	return nil
+}
+
+func exePath() (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("%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("%s is directory", p)
+		}
+	}
+	return "", err
+}
No results found