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,12 +48,13 @@ func (n *NetForward) Run() error {
 		n.DialTimeout = time.Second * 5
 	}
 	if n.StdinMode {
-	breakfor:
+		go func() {
 			for {
-			cmd := strings.TrimSpace(stario.MessageBoxRaw("", "").MustString())
+				cmd := strings.TrimSpace(stario.MessageBox("", "").MustString())
 				for strings.Contains(cmd, "  ") {
 					cmd = strings.Replace(cmd, "  ", " ", -1)
 				}
+				starlog.Debugf("Recv Command %s\n", cmd)
 				cmds := strings.Split(cmd, " ")
 				if len(cmds) < 3 {
 					starlog.Errorln("Invalid Command", cmd)
@@ -99,10 +100,11 @@ func (n *NetForward) Run() error {
 					if cmds[2] == "off" {
 						n.StdinMode = false
 						starlog.Noticef("Stdin Mode Off\n")
-					break breakfor
+						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)
 	}
 }