68 lines
1.8 KiB
Go
68 lines
1.8 KiB
Go
package netforward
|
||
|
||
import (
|
||
"b612.me/stario"
|
||
"b612.me/starlog"
|
||
"github.com/spf13/cobra"
|
||
"os"
|
||
"os/signal"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
var f = new(NetForward)
|
||
var dialTimeout, udpTimeout int64
|
||
|
||
func init() {
|
||
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")
|
||
CmdNetforward.Flags().BoolVarP(&f.EnableUDP, "enable-udp-forward", "u", true, "enable udp forward mode")
|
||
CmdNetforward.Flags().Int64VarP(&dialTimeout, "dial-timeout", "d", 10000, "dial timeout milliseconds")
|
||
CmdNetforward.Flags().Int64VarP(&udpTimeout, "udp-timeout", "D", 60000, "udp connection timeout milliseconds")
|
||
}
|
||
|
||
var CmdNetforward = &cobra.Command{
|
||
Use: "forward",
|
||
Short: "端口转发工具",
|
||
Long: "端口转发工具，支持tcp和udp转发",
|
||
Run: func(cmd *cobra.Command, args []string) {
|
||
if len(args) == 0 {
|
||
starlog.Errorln("please enter a target uri")
|
||
os.Exit(1)
|
||
}
|
||
f.RemoteURI = strings.TrimSpace(args[0])
|
||
if dialTimeout == 0 {
|
||
dialTimeout = 10000
|
||
}
|
||
if udpTimeout == 0 {
|
||
udpTimeout = 60000
|
||
}
|
||
f.DialTimeout = time.Duration(dialTimeout) * time.Millisecond
|
||
f.UDPTimeout = time.Duration(udpTimeout) * time.Millisecond
|
||
if err := f.Run(); err != nil {
|
||
starlog.Errorln("run net forward failed:", err)
|
||
os.Exit(2)
|
||
}
|
||
sign := make(chan os.Signal)
|
||
signal.Notify(sign, os.Interrupt, os.Kill)
|
||
for {
|
||
select {
|
||
case <-sign:
|
||
starlog.Noticeln("Recv Stop Signal From User")
|
||
f.stopFn()
|
||
case <-stario.WaitUntilFinished(func() error {
|
||
for {
|
||
if f.Status() == 0 {
|
||
return nil
|
||
}
|
||
time.Sleep(time.Second)
|
||
}
|
||
}):
|
||
starlog.Infoln("Service Stoped")
|
||
return
|
||
}
|
||
}
|
||
},
|
||
}
|