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().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
			}
		}
	},
}