From 92203a574cd472fa71f17a98a25259df6db8f121 Mon Sep 17 00:00:00 2001
From: Gerhard Tan <gwohau.tan@gmail.com>
Date: Tue, 5 Sep 2023 01:55:34 +0800
Subject: [PATCH] Add runtime from external binary

---
 services/client.go   |   8 ++++
 services/external.go | 107 +++++++++++++++++++++++++++++++++++++++++++
 services/service.go  |   7 +++
 3 files changed, 122 insertions(+)
 create mode 100644 services/external.go

diff --git a/services/client.go b/services/client.go
index 92d0a7af..50ea5231 100644
--- a/services/client.go
+++ b/services/client.go
@@ -61,3 +61,11 @@ func (s *FrpClientService) Reload() error {
 	}
 	return s.svr.ReloadConf(pxyCfgs, visitorCfgs)
 }
+
+func (s *FrpClientService) Close() error {
+	return nil
+}
+
+func (s *FrpClientService) Done() <-chan struct{} {
+	return nil
+}
diff --git a/services/external.go b/services/external.go
new file mode 100644
index 00000000..e765518e
--- /dev/null
+++ b/services/external.go
@@ -0,0 +1,107 @@
+package services
+
+import (
+	"context"
+	"os"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+	"golang.org/x/sys/windows/registry"
+)
+
+var (
+	modKernel32       = syscall.NewLazyDLL("kernel32.dll")
+	procAttachConsole = modKernel32.NewProc("AttachConsole")
+)
+
+type FrpClientExternalService struct {
+	jo     windows.Handle
+	ctx    context.Context
+	cancel context.CancelFunc
+	si     windows.StartupInfo
+	pi     windows.ProcessInformation
+	cmd    *uint16
+}
+
+func NewFrpClientExternalService(path, args string, cfgFile string) (*FrpClientExternalService, error) {
+	jo, err := windows.CreateJobObject(nil, nil)
+	if err != nil {
+		return nil, err
+	}
+	var info = windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{
+		BasicLimitInformation: windows.JOBOBJECT_BASIC_LIMIT_INFORMATION{
+			LimitFlags: windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+		},
+	}
+	if _, err = windows.SetInformationJobObject(jo, windows.JobObjectExtendedLimitInformation,
+		uintptr(unsafe.Pointer(&info)), uint32(unsafe.Sizeof(info))); err != nil {
+		return nil, err
+	}
+	if err = os.Setenv("FRP_CONF", syscall.EscapeArg(cfgFile)); err != nil {
+		return nil, err
+	}
+	if args == "" {
+		args = "-c %FRP_CONF%"
+	}
+	args, err = registry.ExpandString(args)
+	if err != nil {
+		return nil, err
+	}
+	cmd, err := windows.UTF16PtrFromString(syscall.EscapeArg(path) + " " + args)
+	if err != nil {
+		return nil, err
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	return &FrpClientExternalService{
+		jo:     jo,
+		ctx:    ctx,
+		cancel: cancel,
+		si:     windows.StartupInfo{Cb: uint32(unsafe.Sizeof(windows.StartupInfo{}))},
+		cmd:    cmd,
+	}, nil
+}
+
+func (s *FrpClientExternalService) Run() {
+	defer s.cancel()
+	if err := windows.AssignProcessToJobObject(s.jo, windows.CurrentProcess()); err != nil {
+		panic(err)
+	}
+	if err := windows.CreateProcess(nil, s.cmd, nil, nil, false,
+		0, nil, nil, &s.si, &s.pi); err != nil {
+		panic(err)
+	}
+	windows.WaitForSingleObject(s.pi.Process, windows.INFINITE)
+}
+
+func (s *FrpClientExternalService) Reload() error {
+	return nil
+}
+
+func (s *FrpClientExternalService) Stop(wait bool) {
+	var object uint32
+	r, _, err := procAttachConsole.Call(uintptr(s.pi.ProcessId))
+	if r == 0 || err != nil {
+		goto kill
+	}
+	if err = windows.GenerateConsoleCtrlEvent(windows.CTRL_C_EVENT, 0); err != nil {
+		goto kill
+	}
+	if object, err = windows.WaitForSingleObject(s.pi.Process, windows.INFINITE); err != nil {
+		goto kill
+	}
+	if object == windows.WAIT_OBJECT_0 {
+		return
+	}
+kill:
+	windows.TerminateProcess(s.pi.Process, 0)
+}
+
+func (s *FrpClientExternalService) Done() <-chan struct{} {
+	return s.ctx.Done()
+}
+
+func (s *FrpClientExternalService) Close() error {
+	s.cancel()
+	return windows.CloseHandle(s.jo)
+}
diff --git a/services/service.go b/services/service.go
index 17ae7634..c165397f 100644
--- a/services/service.go
+++ b/services/service.go
@@ -3,6 +3,7 @@ package services
 import (
 	"crypto/md5"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"path/filepath"
@@ -16,12 +17,15 @@ import (
 )
 
 type Service interface {
+	io.Closer
 	// Run service in blocking mode.
 	Run()
 	// Reload config file.
 	Reload() error
 	// Stop service and cleanup resources.
 	Stop(wait bool)
+	// Done returns a channel that's closed when service exits.
+	Done() <-chan struct{}
 }
 
 func ServiceNameOfClient(name string) string {
@@ -79,6 +83,7 @@ func (service *frpService) Execute(args []string, r <-chan svc.ChangeRequest, ch
 	if err != nil {
 		return
 	}
+	defer svr.Close()
 
 	go svr.Run()
 
@@ -104,6 +109,8 @@ func (service *frpService) Execute(args []string, r <-chan svc.ChangeRequest, ch
 			svr.Stop(false)
 			deleteFrpConfig(args[0], service.configPath, cc)
 			return
+		case <-svr.Done():
+			return
 		}
 	}
 }