From 2f401589456b7109ba18ddc599d8839306549404 Mon Sep 17 00:00:00 2001 From: Starainrt Date: Fri, 22 Mar 2024 15:16:14 +0800 Subject: [PATCH] add hook for net forward --- netforward/cmd.go | 3 +- netforward/forward.go | 106 +++++++++++++++++++++--------------------- 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/netforward/cmd.go b/netforward/cmd.go index 609766a..81f4785 100644 --- a/netforward/cmd.go +++ b/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 { diff --git a/netforward/forward.go b/netforward/forward.go index b02e361..29c21d7 100644 --- a/netforward/forward.go +++ b/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) } }