star/netforward/cmd.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().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", "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")
	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")
	CmdNetforward.Flags().BoolVarP(&f.UsingKeepAlive, "keepalive", "k", true, "enable keepalive")
	CmdNetforward.Flags().IntVarP(&f.KeepAlivePeriod, "keepalive-period", "P", 10, "keepalive retry period (seconds)")
	CmdNetforward.Flags().IntVarP(&f.KeepAliveIdel, "keepalive-idel", "I", 15, "keepalive idel time (seconds)")
	CmdNetforward.Flags().IntVarP(&f.KeepAliveCount, "keepalive-count", "C", 3, "keepalive count")
	CmdNetforward.Flags().IntVarP(&f.UserTimeout, "user-timeout", "U", 0, "user timeout (milliseconds)")
	CmdNetforward.Flags().BoolVarP(&f.IgnoreEof, "ignore-eof", "E", false, "ignore eof")
}

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)
		}
		time.Sleep(time.Millisecond * 500)
		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
			}
		}
	},
}