Support event hook

Author: hzyitc

README.md

@@ -58,6 +58,14 @@ Flags:
 
   -u, --disable-upnp     Disable UPnP
 
+  -x, --event-hook       Execute command when event triggered:
+                           escape:
+                             %%: percent sign
+                             %e: Event: connecting fail success disconnected
+                             %m: Error message
+                             %p: Local hole port
+                             %a: Hole addr
+
   -h, --help             help for mnh
 ```
 

README_zh.md

@@ -59,6 +59,14 @@ Flags:
 
   -u, --disable-upnp     禁用UPnP
 
+  -x, --event-hook       在事件触发后执行命令:
+                           转义符:
+                             %%: 百分号
+                             %e: 事件: connecting fail success disconnected
+                             %m: 错误消息
+                             %p: 本地洞端口
+                             %a: 洞地址
+
   -h, --help             输出本帮助
 ```
 

go.mod

@@ -3,6 +3,7 @@ module github.com/hzyitc/mnh
 go 1.16
 
 require (
+	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/hzyitc/netutils v0.1.0
 	github.com/libp2p/go-reuseport v0.0.2
 	github.com/spf13/cobra v1.1.3

go.sum

@@ -62,6 +62,8 @@ github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXi
 github.com/google/pprof v0.0.0-20181206194817-3ea8567a2e57/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc=
 github.com/google/pprof v0.0.0-20190515194954-54271f7e092f/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=

main.go

@@ -1,11 +1,15 @@
 package main
 
 import (
+	"net"
 	"os"
+	"os/exec"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
+	"github.com/google/shlex"
 	"github.com/spf13/cobra"
 
 	"github.com/hzyitc/mnh/TCPMode"
@@ -47,6 +51,8 @@ var (
 	service string
 
 	upnpD bool
+
+	eventHook string
 )
 
 func commonCmdRegister(cmd *cobra.Command) {
@@ -58,6 +64,8 @@ func commonCmdRegister(cmd *cobra.Command) {
 	cmd.PersistentFlags().StringVarP(&service, "service", "t", "127.0.0.1:80", "Target service address. Only need in proxy mode")
 
 	cmd.PersistentFlags().BoolVarP(&upnpD, "disable-upnp", "u", false, "Disable UPnP")
+
+	cmd.PersistentFlags().StringVarP(&eventHook, "event-hook", "x", "", "Execute command when event triggered")
 }
 
 func tcpCmdRegister(cmd *cobra.Command) {
@@ -74,6 +82,34 @@ func udpCmdRegister(cmd *cobra.Command) {
 	cmd.AddCommand(udpCmd)
 }
 
+func runHook(event string, errmsg string, port string, addr string) {
+	if eventHook == "" {
+		return
+	}
+
+	cmdline := strings.NewReplacer(
+		"%%", "%",
+		"%e", event,
+		"%m", errmsg,
+		"%p", port,
+		"%a", addr,
+	).Replace(eventHook)
+	log.Debug("Running hook:", cmdline)
+
+	args, err := shlex.Split(cmdline)
+	if err != nil {
+		log.Error("Split hook error:", err.Error())
+		return
+	}
+
+	cmd := exec.Command(args[0], args[1:]...)
+	err = cmd.Start()
+	if err != nil {
+		log.Error("Run hook error:", err.Error())
+		return
+	}
+}
+
 func main() {
 	tcpCmdRegister(rootCmd)
 	udpCmdRegister(rootCmd)
@@ -122,9 +158,12 @@ func tcp() {
 
 	for {
 		func() {
+			runHook("connecting", "", "", "")
+
 			protocol, err := TCPProtocol.NewMnhv1(mode, server, id)
 			if err != nil {
 				log.Error("NewMnhv1 error:", err.Error())
+				runHook("fail", err.Error(), "", "")
 				return
 			}
 			defer protocol.Close()
@@ -136,6 +175,11 @@ func tcp() {
 
 			log.Info("\n\nNow you can use " + protocol.RemoteHoleAddr().String() + " to access your service")
 
+			_, port, _ := net.SplitHostPort(protocol.LocalHoleAddr().String())
+			addr := protocol.RemoteHoleAddr().String()
+			runHook("success", "", port, addr)
+			defer runHook("disconnected", "", port, addr)
+
 			select {
 			case <-protocol.ClosedChan():
 				return
@@ -193,9 +237,12 @@ func udp() {
 
 	for {
 		func() {
+			runHook("connecting", "", "", "")
+
 			protocol, err := UDPProtocol.NewMnhv1(mode, server, id)
 			if err != nil {
 				log.Error("NewMnhv1 error:", err.Error())
+				runHook("fail", err.Error(), "", "")
 				return
 			}
 			defer protocol.Close()
@@ -207,6 +254,11 @@ func udp() {
 
 			log.Info("\n\nNow you can use " + protocol.RemoteHoleAddr().String() + " to access your service")
 
+			_, port, _ := net.SplitHostPort(protocol.LocalHoleAddr().String())
+			addr := protocol.RemoteHoleAddr().String()
+			runHook("success", "", port, addr)
+			defer runHook("disconnected", "", port, addr)
+
 			select {
 			case <-protocol.ClosedChan():
 				return