Add support for running as a service on Windows

Signed-off-by: Brian Dwyer <[email protected]>

README.md

@@ -16,6 +16,10 @@ It reads the configured proxy from the Windows Registry, or can be set via the `
 
 By default, GoNTLM-Proxy listens locally on port 3128, however this can be set via the `GONTLM_BIND` environment variable.
 
+## Service
+If you run this as a service, it will run as NT AUTHORITY/SYSTEM.  If you wish to run it as another user, you can edit the service after installation.
+
+Please note that when running the service as `SYSTEM`, dynamically-generated CA certificate and key will be located in the `SYSTEM` user's home folder at `C:\WINDOWS\system32\config\systemprofile\.gontlm-ca.pem`.  You can always replace these with your own if you already have cert/key.
 
 ## Install
 Release binaries are available under the GitHub Releases page.  Alternatively, you can do this the Go way.

cmd/gontlm-proxy/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"flag"
 	"fmt"
+	"runtime"
 
 	proxy "github.com/bdwyertech/gontlm-proxy/pkg/gontlm-proxy"
 )
@@ -29,5 +30,10 @@ func main() {
 		showVersion()
 		return
 	}
-	proxy.Run()
+
+	if runtime.GOOS == "windows" {
+		proxy.RunWindows()
+	} else {
+		proxy.Run()
+	}
 }

pkg/gontlm-proxy/gontlm-proxy.go

@@ -2,6 +2,7 @@ package ntlm_proxy
 
 import (
 	"crypto/tls"
+	"flag"
 	"log"
 	"net/http"
 	"net/url"
@@ -13,6 +14,14 @@ import (
 	"github.com/elazarl/goproxy"
 )
 
+var proxyServer string
+
+func init() {
+	if flag.Lookup("proxy") == nil {
+		flag.StringVar(&proxyServer, "proxy", getProxyServer(), "Forwarding proxy server")
+	}
+}
+
 func Run() {
 	proxyServer := getEnv("GONTLM_PROXY", getProxyServer())
 	bind := getEnv("GONTLM_BIND", ":3128")

pkg/gontlm-proxy/helpers_other.go

@@ -2,11 +2,7 @@
 
 package ntlm_proxy
 
-import (
-	"os"
-)
-
 func getProxyServer() (proxyServer string) {
-	proxyServer = os.Args[1]
+	proxyServer = getEnv("GONTLM_PROXY", "")
 	return
 }

pkg/gontlm-proxy/helpers_windows.go

@@ -2,22 +2,28 @@
 package ntlm_proxy
 
 import (
-	"golang.org/x/sys/windows/registry"
 	"log"
+	"os"
+
+	"golang.org/x/sys/windows/registry"
 )
 
 func getProxyServer() (proxyServer string) {
+	// Check Environment
+	if proxyFromEnv, ok := os.LookupEnv("GONTLM_PROXY"); ok {
+		proxyServer = proxyFromEnv
+		return
+	}
 	// Pull Proxy from the Registry
 	k, err := registry.OpenKey(registry.CURRENT_USER, `SOFTWARE\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.QUERY_VALUE)
 	if err != nil {
 		log.Fatal(err)
 	}
 	defer k.Close()
 
-	// proxyServer := os.Args[1]
 	proxyServer, _, err = k.GetStringValue("ProxyServer")
 	if err != nil {
-		log.Fatal(err)
+		log.Println(err)
 	}
 	return
 }

pkg/gontlm-proxy/service_other.go

@@ -0,0 +1,11 @@
+// +build !windows
+
+package ntlm_proxy
+
+import (
+	"log"
+)
+
+func RunWindows() {
+	log.Fatal("RunWindows should never be called on a platform other than Windows!")
+}

pkg/gontlm-proxy/service_windows.go

@@ -0,0 +1,82 @@
+// +build windows
+
+// Windows Service Handler
+
+package ntlm_proxy
+
+import (
+	"flag"
+	"log"
+
+	"github.com/kardianos/service"
+)
+
+type program struct {
+	exit chan struct{}
+}
+
+var svcFlag string
+
+func init() {
+	if flag.Lookup("service") == nil {
+		flag.StringVar(&svcFlag, "service", "", "Control the Windows System service.")
+	}
+}
+
+func RunWindows() {
+	svcConfig := &service.Config{
+		Name:        "gontlm-proxy",
+		DisplayName: "GoNTLM Proxy",
+		Description: "GoNTLM Forwarding Proxy",
+		Arguments:   []string{"-proxy", proxyServer},
+	}
+
+	prg := &program{}
+	s, err := service.New(prg, svcConfig)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	if len(svcFlag) != 0 {
+		err := service.Control(s, svcFlag)
+		if err != nil {
+			log.Printf("Valid actions: %q\n", service.ControlAction)
+			log.Fatal(err)
+		}
+		return
+	}
+
+	if err = s.Run(); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func (p *program) Start(s service.Service) (err error) {
+	if service.Interactive() {
+		log.Println("Running in terminal.")
+	} else {
+		log.Println("Running under service manager.")
+	}
+	p.exit = make(chan struct{})
+
+	go p.run()
+
+	return
+}
+
+func (p *program) run() (err error) {
+	// Run the Proxy
+	go Run()
+	// Wait for Exit Signal
+	for {
+		select {
+		case <-p.exit:
+			return
+		}
+	}
+}
+
+func (p *program) Stop(s service.Service) (err error) {
+	close(p.exit)
+	return
+}