add hook for net forward

netforward/cmd.go

@@ -16,7 +16,7 @@ var dialTimeout, udpTimeout int64
 func init() {
 	CmdNetforward.Flags().IntVarP(&f.DelayToward, "delay-toward", "T", 0, "delay toward milliseconds,0 for both,1 for local,2 for remote")
 	CmdNetforward.Flags().BoolVarP(&f.StdinMode, "stdin", "s", false, "enable stdin mode")
-	CmdNetforward.Flags().IntVarP(&f.DelayMilSec, "delay", "D", 0, "delay milliseconds")
+	CmdNetforward.Flags().IntVarP(&f.DelayMilSec, "delay", "S", 0, "delay milliseconds")
 	CmdNetforward.Flags().StringVarP(&f.LocalAddr, "local", "l", "0.0.0.0", "bind address")
 	CmdNetforward.Flags().IntVarP(&f.LocalPort, "port", "p", 11270, "local listen port")
 	CmdNetforward.Flags().BoolVarP(&f.EnableTCP, "enable-tcp-forward", "t", true, "enable tcp forward mode")
@@ -47,6 +47,7 @@ var CmdNetforward = &cobra.Command{
 			starlog.Errorln("run net forward failed:", err)
 			os.Exit(2)
 		}
+		time.Sleep(time.Millisecond * 500)
 		sign := make(chan os.Signal)
 		signal.Notify(sign, os.Interrupt, os.Kill)
 		for {

netforward/forward.go

@@ -35,7 +35,7 @@ func (n *NetForward) Close() {
 }
 
 func (n *NetForward) Status() int32 {
-	return n.running
+	return atomic.LoadInt32(&n.running)
 }
 
 func (n *NetForward) Run() error {
@@ -48,61 +48,63 @@ func (n *NetForward) Run() error {
 		n.DialTimeout = time.Second * 5
 	}
 	if n.StdinMode {
-	breakfor:
-		for {
-			cmd := strings.TrimSpace(stario.MessageBoxRaw("", "").MustString())
-			for strings.Contains(cmd, "  ") {
-				cmd = strings.Replace(cmd, "  ", " ", -1)
-			}
-			cmds := strings.Split(cmd, " ")
-			if len(cmds) < 3 {
-				starlog.Errorln("Invalid Command", cmd)
-				continue
-			}
-			switch cmds[0] + cmds[1] {
-			case "setremote":
-				n.RemoteURI = cmds[2]
-				starlog.Noticef("Remote URI Set to %s\n", n.RemoteURI)
-			case "setdelaytoward":
-				tmp, err := strconv.Atoi(cmds[2])
-				if err != nil {
-					starlog.Errorln("Invalid Delay Toward Value", cmds[2])
-					continue
-				}
-				n.DelayToward = tmp
-				starlog.Noticef("Delay Toward Set to %d\n", n.DelayToward)
-			case "setdelay":
-				tmp, err := strconv.Atoi(cmds[2])
-				if err != nil {
-					starlog.Errorln("Invalid Delay Value", cmds[2])
-					continue
-				}
-				n.DelayMilSec = tmp
-				starlog.Noticef("Delay Set to %d\n", n.DelayMilSec)
-			case "setdialtimeout":
-				tmp, err := strconv.Atoi(cmds[2])
-				if err != nil {
-					starlog.Errorln("Invalid Dial Timeout Value", cmds[2])
-					continue
+		go func() {
+			for {
+				cmd := strings.TrimSpace(stario.MessageBox("", "").MustString())
+				for strings.Contains(cmd, "  ") {
+					cmd = strings.Replace(cmd, "  ", " ", -1)
 				}
-				n.DialTimeout = time.Millisecond * time.Duration(tmp)
-				starlog.Noticef("Dial Timeout Set to %d\n", n.DialTimeout)
-			case "setudptimeout":
-				tmp, err := strconv.Atoi(cmds[2])
-				if err != nil {
-					starlog.Errorln("Invalid UDP Timeout Value", cmds[2])
+				starlog.Debugf("Recv Command %s\n", cmd)
+				cmds := strings.Split(cmd, " ")
+				if len(cmds) < 3 {
+					starlog.Errorln("Invalid Command", cmd)
 					continue
 				}
-				n.UDPTimeout = time.Millisecond * time.Duration(tmp)
-				starlog.Noticef("UDP Timeout Set to %d\n", n.UDPTimeout)
-			case "setstdin":
-				if cmds[2] == "off" {
-					n.StdinMode = false
-					starlog.Noticef("Stdin Mode Off\n")
-					break breakfor
+				switch cmds[0] + cmds[1] {
+				case "setremote":
+					n.RemoteURI = cmds[2]
+					starlog.Noticef("Remote URI Set to %s\n", n.RemoteURI)
+				case "setdelaytoward":
+					tmp, err := strconv.Atoi(cmds[2])
+					if err != nil {
+						starlog.Errorln("Invalid Delay Toward Value", cmds[2])
+						continue
+					}
+					n.DelayToward = tmp
+					starlog.Noticef("Delay Toward Set to %d\n", n.DelayToward)
+				case "setdelay":
+					tmp, err := strconv.Atoi(cmds[2])
+					if err != nil {
+						starlog.Errorln("Invalid Delay Value", cmds[2])
+						continue
+					}
+					n.DelayMilSec = tmp
+					starlog.Noticef("Delay Set to %d\n", n.DelayMilSec)
+				case "setdialtimeout":
+					tmp, err := strconv.Atoi(cmds[2])
+					if err != nil {
+						starlog.Errorln("Invalid Dial Timeout Value", cmds[2])
+						continue
+					}
+					n.DialTimeout = time.Millisecond * time.Duration(tmp)
+					starlog.Noticef("Dial Timeout Set to %d\n", n.DialTimeout)
+				case "setudptimeout":
+					tmp, err := strconv.Atoi(cmds[2])
+					if err != nil {
+						starlog.Errorln("Invalid UDP Timeout Value", cmds[2])
+						continue
+					}
+					n.UDPTimeout = time.Millisecond * time.Duration(tmp)
+					starlog.Noticef("UDP Timeout Set to %d\n", n.UDPTimeout)
+				case "setstdin":
+					if cmds[2] == "off" {
+						n.StdinMode = false
+						starlog.Noticef("Stdin Mode Off\n")
+						return
+					}
 				}
 			}
-		}
+		}()
 	}
 	if n.EnableTCP {
 		go n.runTCP()
@@ -148,7 +150,7 @@ func (n *NetForward) runTCP() error {
 			}
 			log.Infof("TCP Connect %s <==> %s\n", conn.RemoteAddr().String(), rmt.RemoteAddr().String())
 			n.copy(rmt, conn)
-			log.Noticef("TCP Connection Closed  %s <==> %s", conn.RemoteAddr().String(), n.RemoteURI)
+			log.Noticef("TCP Connection Closed  %s <==> %s\n", conn.RemoteAddr().String(), n.RemoteURI)
 		}(conn)
 	}
 }