From 550ec5fe735900e3f65d350ad4f7e564f8dff795 Mon Sep 17 00:00:00 2001
From: Lucus Lee <lixin9311@gmail.com>
Date: Mon, 23 Sep 2019 02:35:38 +0900
Subject: [PATCH] Better plugin life-cycle management

---
 main.go   |  1 +
 plugin.go | 27 ++++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index e1a7a428..042a10de 100644
--- a/main.go
+++ b/main.go
@@ -173,6 +173,7 @@ func main() {
 	sigCh := make(chan os.Signal, 1)
 	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
 	<-sigCh
+	killPlugin()
 }
 
 func parseURL(s string) (addr, cipher, password string, err error) {
diff --git a/plugin.go b/plugin.go
index 0d2ac420..237296e5 100644
--- a/plugin.go
+++ b/plugin.go
@@ -5,8 +5,13 @@ import (
 	"net"
 	"os"
 	"os/exec"
+	"path/filepath"
+	"syscall"
+	"time"
 )
 
+var pluginCmd *exec.Cmd
+
 func startPlugin(plugin, pluginOpts, ssAddr string, isServer bool) (newAddr string, err error) {
 	logf("starting plugin (%s) with option (%s)....", plugin, pluginOpts)
 	freePort, err := getFreePort()
@@ -31,9 +36,28 @@ func startPlugin(plugin, pluginOpts, ssAddr string, isServer bool) (newAddr stri
 	return
 }
 
+func killPlugin() {
+	if pluginCmd != nil {
+		pluginCmd.Process.Signal(syscall.SIGTERM)
+		waitCh := make(chan struct{})
+		go func() {
+			pluginCmd.Wait()
+			close(waitCh)
+		}()
+		timeout := time.After(3 * time.Second)
+		select {
+		case <-waitCh:
+		case <-timeout:
+			pluginCmd.Process.Kill()
+		}
+	}
+}
+
 func execPlugin(plugin, pluginOpts, remoteHost, remotePort, localHost, localPort string) error {
 	if fileExists(plugin) {
-		plugin = "./" + plugin
+		if !filepath.IsAbs(plugin) {
+			plugin = "./" + plugin
+		}
 	}
 	logH := newLogHelper("[" + plugin + "]: ")
 	env := append(os.Environ(),
@@ -53,6 +77,7 @@ func execPlugin(plugin, pluginOpts, remoteHost, remotePort, localHost, localPort
 	if err := cmd.Start(); err != nil {
 		return err
 	}
+	pluginCmd = cmd
 	go func() {
 		if err := cmd.Wait(); err != nil {
 			logf("plugin exited (%v)\n", err)